syft/internal/archive/limits.go

package archive

import (
	"errors"
	"sync"
)

// ErrDiskLimitReached reports that placing an archive's content would exceed the disk limit.
var ErrDiskLimitReached = errors.New("archive content would exceed the disk limit")

// Limits bound how much archive content one scan holds at once. Each limit falls as archives are
// released, so it bounds peak usage rather than the total ever written.
//
// For both limits: positive is the bound, zero forbids the resource, negative is unbounded.
type Limits struct {
	// MaxMemoryBytes bounds archive content held in memory. Content that does not fit is written to
	// disk instead.
	MaxMemoryBytes int64

	// MaxDiskBytes bounds archive content written to disk. Content that does not fit is dropped and
	// the archive is skipped or truncated.
	MaxDiskBytes int64
}

// Limiter tracks how much archive content a scan holds right now against its Limits. A nil *Limiter
// enforces nothing.
type Limiter struct {
	limits Limits

	mu     sync.Mutex
	memory int64
	disk   int64

	peakMemory int64
	peakDisk   int64
}

func NewLimiter(limits Limits) *Limiter {
	return &Limiter{limits: limits}
}

// Peak reports the most memory and disk held at any one moment.
func (l *Limiter) Peak() (memory, disk int64) {
	if l == nil {
		return 0, 0
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.peakMemory, l.peakDisk
}

// InUse reports the memory and disk held right now.
func (l *Limiter) InUse() (memory, disk int64) {
	if l == nil {
		return 0, 0
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.memory, l.disk
}

// charge starts accounting for one archive. Releasing the charge gives back everything it took.
func (l *Limiter) charge() *charge {
	if l == nil {
		return nil
	}
	return &charge{limiter: l}
}

// charge is one archive's draw on a Limiter. A nil *charge admits everything.
//
// Bytes are charged as they land, never from the sizes an archive declares, since those can lie.
type charge struct {
	limiter  *Limiter
	inMemory int64
	onDisk   int64
}

func admits(limit, held, n int64) bool {
	switch {
	case limit < 0:
		return true
	case limit == 0:
		return false
	default:
		return held+n <= limit
	}
}

// memory charges n bytes held in memory, or reports false (charging nothing) when the memory limit
// does not admit them.
func (c *charge) memory(n int64) bool {
	if c == nil || n <= 0 {
		return true
	}
	l := c.limiter
	l.mu.Lock()
	defer l.mu.Unlock()
	if !admits(l.limits.MaxMemoryBytes, l.memory, n) {
		return false
	}
	l.memory += n
	c.inMemory += n
	l.peakMemory = max(l.peakMemory, l.memory)
	return true
}

// disk charges n bytes written to disk, or reports false (charging nothing) when the disk limit does
// not admit them.
func (c *charge) disk(n int64) bool {
	if c == nil || n <= 0 {
		return true
	}
	l := c.limiter
	l.mu.Lock()
	defer l.mu.Unlock()
	if !admits(l.limits.MaxDiskBytes, l.disk, n) {
		return false
	}
	l.disk += n
	c.onDisk += n
	l.peakDisk = max(l.peakDisk, l.disk)
	return true
}

// index charges n bytes of index bookkeeping, which lives in memory whatever the limits say. A zero
// memory limit is a valid configuration that must still index archives, so only then does the cost
// fall back to the disk budget; a bounded memory limit bounds the index too.
func (c *charge) index(n int64) bool {
	if c.memory(n) {
		return true
	}
	return c != nil && c.limiter.limits.MaxMemoryBytes == 0 && c.disk(n)
}

// refundMemory gives back n bytes no longer held in memory.
func (c *charge) refundMemory(n int64) {
	if c == nil || n <= 0 {
		return
	}
	l := c.limiter
	l.mu.Lock()
	defer l.mu.Unlock()
	n = min(n, c.inMemory)
	l.memory -= n
	c.inMemory -= n
}

// refundDisk gives back n bytes no longer on disk.
func (c *charge) refundDisk(n int64) {
	if c == nil || n <= 0 {
		return
	}
	l := c.limiter
	l.mu.Lock()
	defer l.mu.Unlock()
	n = min(n, c.onDisk)
	l.disk -= n
	c.onDisk -= n
}

// release gives back everything this charge took. Safe to call more than once.
func (c *charge) release() {
	if c == nil {
		return
	}
	l := c.limiter
	l.mu.Lock()
	defer l.mu.Unlock()
	l.memory -= c.inMemory
	l.disk -= c.onDisk
	c.inMemory, c.onDisk = 0, 0
}