add limitedstore

This commit is contained in:
Acconut 2015-04-11 13:40:39 +02:00
parent efe70bb0f3
commit e0fa7556b4
3 changed files with 196 additions and 1 deletions

View File

@ -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
}

View File

@ -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")
}
}

View File

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