// SPDX-License-Identifier: MIT // SPDX-FileCopyrightText: Copyright (c) 2024 Matthew Penner package ufs import ( "sync/atomic" ) type Quota struct { // fs is the underlying filesystem that runs the actual I/O operations. *UnixFS // limit is the size limit of the filesystem. // // limit is atomic to allow the limit to be safely changed after the // filesystem was created. // // A limit of `-1` disables any write operation from being performed. // A limit of `0` disables any limit checking. limit atomic.Int64 // usage is the current usage of the filesystem. // // If usage is set to `-1`, it hasn't been calculated yet. usage atomic.Int64 } func NewQuota(fs *UnixFS, limit int64) *Quota { qfs := Quota{UnixFS: fs} qfs.limit.Store(limit) return &qfs } // Close closes the filesystem. func (fs *Quota) Close() (err error) { err = fs.UnixFS.Close() return } // Limit returns the limit of the filesystem. func (fs *Quota) Limit() int64 { return fs.limit.Load() } // SetLimit returns the limit of the filesystem. func (fs *Quota) SetLimit(newLimit int64) int64 { return fs.limit.Swap(newLimit) } // Usage returns the current usage of the filesystem. func (fs *Quota) Usage() int64 { return fs.usage.Load() } // SetUsage updates the total usage of the filesystem. func (fs *Quota) SetUsage(newUsage int64) int64 { return fs.usage.Swap(newUsage) } // Add adds `i` to the tracked usage total. func (fs *Quota) Add(i int64) int64 { usage := fs.Usage() // If adding `i` to the usage will put us below 0, cap it. (`i` can be negative) if usage+i < 0 { fs.usage.Store(0) return 0 } return fs.usage.Add(i) } // CanFit checks if the given size can fit in the filesystem without exceeding // the limit of the filesystem. func (fs *Quota) CanFit(size int64) bool { // Get the size limit of the filesystem. limit := fs.Limit() switch limit { case -1: // A limit of -1 means no write operations are allowed. return false case 0: // A limit of 0 means unlimited. return true } // Any other limit is a value we need to check. usage := fs.Usage() if usage == -1 { // We don't know what the current usage is yet. return true } // If the current usage + the requested size are under the limit of the // filesystem, allow it. if usage+size <= limit { return true } // Welp, the size would exceed the limit of the filesystem, deny it. return false } func (fs *Quota) Remove(name string) error { // For information on why this interface is used here, check its // documentation. s, err := fs.RemoveStat(name) if err != nil { return err } // Don't reduce the quota's usage as `name` is not a regular file. if !s.Mode().IsRegular() { return nil } // Remove the size of the deleted file from the quota usage. fs.Add(-s.Size()) return nil } // RemoveAll removes path and any children it contains. // // It removes everything it can but returns the first error // it encounters. If the path does not exist, RemoveAll // returns nil (no error). // // If there is an error, it will be of type *PathError. func (fs *Quota) RemoveAll(name string) error { name, err := fs.unsafePath(name) if err != nil { return err } // While removeAll internally checks this, I want to make sure we check it // and return the proper error so our tests can ensure that this will never // be a possibility. if name == "." { return &PathError{ Op: "removeall", Path: name, Err: ErrBadPathResolution, } } return fs.removeAll(name) } func (fs *Quota) removeAll(path string) error { return removeAll(fs, path) } func (fs *Quota) unlinkat(dirfd int, name string, flags int) error { if flags == 0 { s, err := fs.Lstatat(dirfd, name) if err == nil && s.Mode().IsRegular() { fs.Add(-s.Size()) } } return fs.UnixFS.unlinkat(dirfd, name, flags) }