From e0fa7556b43a320ab27b5391f4816a5520c258f0 Mon Sep 17 00:00:00 2001 From: Acconut Date: Sat, 11 Apr 2015 13:40:39 +0200 Subject: [PATCH] add limitedstore --- limitedstore/limitedstore.go | 98 +++++++++++++++++++++++++++++++ limitedstore/limitedstore_test.go | 83 ++++++++++++++++++++++++++ tusd/main.go | 16 ++++- 3 files changed, 196 insertions(+), 1 deletion(-) create mode 100644 limitedstore/limitedstore.go create mode 100644 limitedstore/limitedstore_test.go diff --git a/limitedstore/limitedstore.go b/limitedstore/limitedstore.go new file mode 100644 index 0000000..8a3b5ad --- /dev/null +++ b/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 +} diff --git a/limitedstore/limitedstore_test.go b/limitedstore/limitedstore_test.go new file mode 100644 index 0000000..d4f75ea --- /dev/null +++ b/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") + } +} diff --git a/tusd/main.go b/tusd/main.go index 98c77d0..1135789 100644 --- a/tusd/main.go +++ b/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{