steamcache2/limitedFs/limitedFs.go

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()
}