add limitedstore

limitedstore/limitedstore.go

@@ -0,0 +1,98 @@
+// Package limitedstore implements a simple wrapper around existing
+// datastores (tusd.DataStore) while limiting the used storage size.
+// It will start terminating existing uploads if not enough space is left in
+// order to create a new upload.
+// This package's functionality is very limited and naive. It will terminate
+// uploads whether they are finished yet or not and it won't terminate them
+// intelligently (e.g. bigger uploads first). Only one datastore is allowed to
+// access the underlying storage else the limited store will not function
+// properly. Two tusd.FileStore instances using the same directory, for example.
+// In addition the limited store will keep a list of the uploads' ids in memory
+// which may create a growing memory leak.
+package limitedstore
+
+import (
+	"github.com/tus/tusd"
+	"sync"
+)
+
+type LimitedStore struct {
+	StoreSize int64
+	tusd.DataStore
+
+	uploads  map[string]int64
+	usedSize int64
+
+	mutex *sync.Mutex
+}
+
+// Create a new limited store with the given size as the maximum storage size
+func New(storeSize int64, dataStore tusd.DataStore) *LimitedStore {
+	return &LimitedStore{
+		StoreSize: storeSize,
+		DataStore: dataStore,
+		uploads:   make(map[string]int64),
+		mutex:     new(sync.Mutex),
+	}
+}
+
+func (store *LimitedStore) NewUpload(info tusd.FileInfo) (string, error) {
+	store.mutex.Lock()
+	defer store.mutex.Unlock()
+
+	if err := store.ensureSpace(info.Size); err != nil {
+		return "", err
+	}
+
+	id, err := store.DataStore.NewUpload(info)
+	if err != nil {
+		return "", err
+	}
+
+	store.usedSize += info.Size
+	store.uploads[id] = info.Size
+
+	return id, nil
+}
+
+func (store *LimitedStore) Terminate(id string) error {
+	store.mutex.Lock()
+	defer store.mutex.Unlock()
+
+	return store.terminate(id)
+}
+
+func (store *LimitedStore) terminate(id string) error {
+	err := store.DataStore.Terminate(id)
+	if err != nil {
+		return err
+	}
+
+	size := store.uploads[id]
+	delete(store.uploads, id)
+	store.usedSize -= size
+
+	return nil
+}
+
+// Ensure enough space is available to store an upload of the specified size.
+// It will terminate uploads until enough space is freed.
+func (store *LimitedStore) ensureSpace(size int64) error {
+	if (store.usedSize + size) <= store.StoreSize {
+		// Enough space is available to store the new upload
+		return nil
+	}
+
+	for id, _ := range store.uploads {
+		if err := store.terminate(id); err != nil {
+			return err
+		}
+
+		if (store.usedSize + size) <= store.StoreSize {
+			// Enough space has been freed to store the new upload
+			return nil
+		}
+	}
+
+	return nil
+}

limitedstore/limitedstore_test.go

@@ -0,0 +1,83 @@
+package limitedstore
+
+import (
+	"github.com/tus/tusd"
+	"io"
+	"testing"
+)
+
+type dataStore struct {
+	t                  *testing.T
+	firstUploadCreated bool
+	uploadTerminated   bool
+}
+
+func (store *dataStore) NewUpload(info tusd.FileInfo) (string, error) {
+	if !store.firstUploadCreated {
+		if info.Size != 80 {
+			store.t.Errorf("expect size to be 80, got %v", info.Size)
+		}
+		store.firstUploadCreated = true
+
+		return "1", nil
+	}
+
+	if info.Size != 50 {
+		store.t.Errorf("expect size to be 50, got %v", info.Size)
+	}
+	return "2", nil
+}
+
+func (store *dataStore) WriteChunk(id string, offset int64, src io.Reader) (int64, error) {
+	return 0, nil
+}
+
+func (store *dataStore) GetInfo(id string) (tusd.FileInfo, error) {
+	return tusd.FileInfo{}, nil
+}
+
+func (store *dataStore) GetReader(id string) (io.Reader, error) {
+	return nil, tusd.ErrNotImplemented
+}
+
+func (store *dataStore) Terminate(id string) error {
+	if id != "1" {
+		store.t.Errorf("expect first upload to be terminated, got %v", id)
+	}
+	store.uploadTerminated = true
+
+	return nil
+}
+
+func TestLimitedStore(t *testing.T) {
+	dataStore := &dataStore{
+		t: t,
+	}
+	store := New(100, dataStore)
+
+	// Create new upload (80 bytes)
+	id, err := store.NewUpload(tusd.FileInfo{
+		Size: 80,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if id != "1" {
+		t.Errorf("expected first upload to be created, got %v", id)
+	}
+
+	// Create new upload (50 bytes)
+	id, err = store.NewUpload(tusd.FileInfo{
+		Size: 50,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if id != "2" {
+		t.Errorf("expected second upload to be created, got %v", id)
+	}
+
+	if !dataStore.uploadTerminated {
+		t.Error("expected first upload to be terminated")
+	}
+}

tusd/main.go

@@ -4,6 +4,7 @@ import (
 	"flag"
 	"github.com/tus/tusd"
 	"github.com/tus/tusd/filestore"
+	"github.com/tus/tusd/limitedstore"
 	"log"
 	"net/http"
 	"os"
@@ -13,6 +14,7 @@ var httpHost string
 var httpPort string
 var maxSize int64
 var dir string
+var storeSize int64
 
 var stdout = log.New(os.Stdout, "[tusd] ", 0)
 var stderr = log.New(os.Stderr, "[tusd] ", 0)
@@ -22,6 +24,7 @@ func init() {
 	flag.StringVar(&httpPort, "port", "1080", "Port to bind HTTP server to")
 	flag.Int64Var(&maxSize, "max-size", 0, "Maximum size of uploads in bytes")
 	flag.StringVar(&dir, "dir", "./data", "Directory to store uploads in")
+	flag.Int64Var(&storeSize, "store-size", 0, "Size of disk space allowed to storage")
 
 	flag.Parse()
 }
@@ -33,10 +36,21 @@ func main() {
 		stderr.Fatalf("Unable to ensure directory exists: %s", err)
 	}
 
-	store := filestore.FileStore{
+	var store tusd.DataStore
+	store = filestore.FileStore{
 		Path: dir,
 	}
 
+	if storeSize > 0 {
+		store = limitedstore.New(storeSize, store)
+		stdout.Printf("Using %.2fMB as storage size.\n", float64(storeSize)/1024/1024)
+
+		// We need to ensure that a single upload can fit into the storage size
+		if maxSize > storeSize || maxSize == 0 {
+			maxSize = storeSize
+		}
+	}
+
 	stdout.Printf("Using %.2fMB as maximum size.\n", float64(maxSize)/1024/1024)
 
 	handler, err := tusd.NewHandler(tusd.Config{