lottsa stuff

limitedFs/limitedFs.go

@@ -0,0 +1,192 @@
+package limitedFs
+
+import (
+	"errors"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/labstack/gommon/log"
+	"github.com/spf13/afero"
+)
+
+var (
+	// Ensure LimitedFs implements afero.Fs.
+	_ afero.Fs = &LimitedFs{}
+
+	// ErrNotEnoughSpace is returned when there's not enough space to write.
+	ErrNotEnoughSpace = errors.New("not enough space")
+)
+
+// LimitedFs wraps an afero.Fs to limit the storage space.
+type LimitedFs struct {
+	fs afero.Fs
+
+	mu        sync.Mutex
+	usedSpace int64
+	maxSpace  int64
+}
+
+// NewLimitedFs creates a new LimitedFs with the specified max storage space.
+func NewLimitedFs(source afero.Fs, maxSpace int64) *LimitedFs {
+	lfs := &LimitedFs{
+		fs:        source,
+		usedSpace: 0,
+		maxSpace:  maxSpace,
+	}
+
+	// If the source filesystem is a memory filesystem, there is no need to scan it.
+	if _, ok := source.(*afero.MemMapFs); !ok {
+		// Scan the source filesystem to calculate the used space.
+		log.Infof("Scanning source filesystem to calculate used space...")
+		err := afero.Walk(source, "/", func(path string, info os.FileInfo, err error) error {
+			if err != nil {
+				return err
+			}
+			lfs.usedSpace += info.Size()
+			return nil
+		})
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	return lfs
+}
+
+// Create implements the Create method of afero.Fs, ensuring we don't exceed storage limits.
+func (lfs *LimitedFs) Create(name string) (afero.File, error) {
+	file, err := lfs.fs.Create(name)
+	if err != nil {
+		return nil, err
+	}
+	return &LimitedFile{
+		File: file,
+		lfs:  lfs,
+	}, nil
+}
+
+func (lfs *LimitedFs) Mkdir(name string, perm os.FileMode) error {
+	return lfs.fs.Mkdir(name, perm)
+}
+
+func (lfs *LimitedFs) MkdirAll(path string, perm os.FileMode) error {
+	return lfs.fs.MkdirAll(path, perm)
+}
+
+func (lfs *LimitedFs) Open(name string) (afero.File, error) {
+	file, err := lfs.fs.Open(name)
+	if err != nil {
+		return nil, err
+	}
+	return &LimitedFile{
+		File: file,
+		lfs:  lfs,
+	}, nil
+}
+
+func (lfs *LimitedFs) OpenFile(name string, flag int, perm os.FileMode) (afero.File, error) {
+	file, err := lfs.fs.OpenFile(name, flag, perm)
+	if err != nil {
+		return nil, err
+	}
+	return &LimitedFile{
+		File: file,
+		lfs:  lfs,
+	}, nil
+}
+
+func (lfs *LimitedFs) Remove(name string) error {
+	info, err := lfs.fs.Stat(name)
+	if err != nil {
+		return err
+	}
+	lfs.mu.Lock()
+	lfs.usedSpace -= info.Size()
+	lfs.mu.Unlock()
+	return lfs.fs.Remove(name)
+}
+
+func (lfs *LimitedFs) RemoveAll(path string) error {
+	// Calculate the space used by the directory and its children.
+	var size int64
+	err := afero.Walk(lfs.fs, path, func(subpath string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		size += info.Size()
+		return nil
+	})
+
+	if err != nil {
+		return err
+	}
+
+	lfs.mu.Lock()
+	lfs.usedSpace -= size
+	lfs.mu.Unlock()
+
+	return lfs.RemoveAll(path)
+}
+
+func (lfs *LimitedFs) Rename(oldname, newname string) error {
+	return lfs.fs.Rename(oldname, newname)
+}
+
+func (lfs *LimitedFs) Stat(name string) (os.FileInfo, error) {
+	return lfs.fs.Stat(name)
+}
+
+func (lfs *LimitedFs) Name() string {
+	return "LimitedFS"
+}
+
+func (lfs *LimitedFs) Chmod(name string, mode os.FileMode) error {
+	return lfs.fs.Chmod(name, mode)
+}
+
+func (lfs *LimitedFs) Chown(name string, uid, gid int) error {
+	return lfs.fs.Chown(name, uid, gid)
+}
+
+func (lfs *LimitedFs) Chtimes(name string, atime time.Time, mtime time.Time) error {
+	return lfs.fs.Chtimes(name, atime, mtime)
+}
+
+// ensure LimitedFile implements afero.File.
+var _ afero.File = &LimitedFile{}
+
+// LimitedFile wraps an afero.File to manage space usage on write operations.
+type LimitedFile struct {
+	afero.File
+	lfs *LimitedFs
+}
+
+// Write checks if there's enough space before writing.
+func (lf *LimitedFile) Write(p []byte) (n int, err error) {
+	lf.lfs.mu.Lock()
+	if int64(len(p)) > lf.lfs.maxSpace-lf.lfs.usedSpace {
+		lf.lfs.mu.Unlock()
+		return 0, ErrNotEnoughSpace
+	}
+	lf.lfs.mu.Unlock()
+	n, err = lf.File.Write(p)
+	if err == nil {
+		lf.lfs.mu.Lock()
+		lf.lfs.usedSpace += int64(n)
+		lf.lfs.mu.Unlock()
+	}
+	return
+}
+
+// Close updates used space and calls the underlying file's Close method.
+func (lf *LimitedFile) Close() error {
+	info, err := lf.File.Stat()
+	if err != nil {
+		return err
+	}
+	lf.lfs.mu.Lock()
+	lf.lfs.usedSpace -= info.Size()
+	lf.lfs.mu.Unlock()
+	return lf.File.Close()
+}