teraWeb/backend/internal/storage/capacity/governor.go

package capacity

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sync"
)

const (
	ClassUpdateStaging = "update_staging"
	EnforcementV1      = "application-governor-v1"
)

var ErrInsufficientSpace = errors.New("insufficient storage capacity")

type Policy struct {
	FormatVersion int    `json:"format_version"`
	Layout        string `json:"layout"`
	Mount         struct {
		Source          string   `json:"source"`
		PartLabel       string   `json:"partlabel"`
		Filesystem      string   `json:"filesystem"`
		RequiredOptions []string `json:"required_options"`
	} `json:"mount"`
	CapsBytes struct {
		UpdateStaging    int64 `json:"update_staging"`
		UpdatePackageMax int64 `json:"update_package_max"`
		PortEvents       int64 `json:"port_events"`
		AuditLogs        int64 `json:"audit_logs"`
		SystemEvents     int64 `json:"system_events"`
		StateAndIndexes  int64 `json:"state_and_indexes"`
		ExportScratch    int64 `json:"export_scratch"`
		GlobalFreeFloor  int64 `json:"global_free_floor"`
	} `json:"caps_bytes"`
	Enforcement           string `json:"enforcement"`
	PersistentUARTPayload bool   `json:"persistent_uart_payload"`
}

type Snapshot struct {
	Pressure         string `json:"pressure"`
	FreeBytes        int64  `json:"free_bytes"`
	GlobalFreeFloor  int64  `json:"global_free_floor_bytes"`
	UpdateUsedBytes  int64  `json:"update_used_bytes"`
	UpdateReserved   int64  `json:"update_reserved_bytes"`
	UpdateCapBytes   int64  `json:"update_cap_bytes"`
	UpdatePackageMax int64  `json:"update_package_max_bytes"`
}

type freeSpaceFunc func(string) (int64, error)
type usageFunc func(string) (int64, error)

type Governor struct {
	mu       sync.Mutex
	root     string
	policy   Policy
	reserved int64
	free     freeSpaceFunc
	usage    usageFunc
}

type Reservation struct {
	governor *Governor
	bytes    int64
	once     sync.Once
}

func LoadPolicy(path string) (Policy, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return Policy{}, fmt.Errorf("read storage policy: %w", err)
	}
	if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
		return Policy{}, fmt.Errorf("storage policy must be a regular file")
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return Policy{}, fmt.Errorf("read storage policy: %w", err)
	}
	var policy Policy
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&policy); err != nil {
		return Policy{}, fmt.Errorf("decode storage policy: %w", err)
	}
	if policy.FormatVersion != 1 || policy.Layout != "tera-userdata-v1" || policy.Mount.Source != "/dev/mmcblk0p8" || policy.Mount.PartLabel != "userdata" || policy.Mount.Filesystem != "ext4" || policy.PersistentUARTPayload {
		return Policy{}, fmt.Errorf("unsupported storage policy identity")
	}
	if policy.Enforcement != EnforcementV1 {
		return Policy{}, fmt.Errorf("storage capacity enforcement is not enabled")
	}
	if policy.CapsBytes.UpdateStaging <= 0 || policy.CapsBytes.UpdatePackageMax <= 0 || policy.CapsBytes.PortEvents <= 0 ||
		policy.CapsBytes.AuditLogs <= 0 || policy.CapsBytes.SystemEvents <= 0 || policy.CapsBytes.StateAndIndexes <= 0 ||
		policy.CapsBytes.ExportScratch <= 0 || policy.CapsBytes.GlobalFreeFloor <= 0 || policy.CapsBytes.UpdatePackageMax > policy.CapsBytes.UpdateStaging {
		return Policy{}, fmt.Errorf("invalid storage policy caps")
	}
	return policy, nil
}

func New(root string, policy Policy) (*Governor, error) {
	root = filepath.Clean(root)
	info, err := os.Lstat(root)
	if err != nil {
		return nil, fmt.Errorf("inspect update staging root: %w", err)
	}
	if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
		return nil, fmt.Errorf("update staging root must be a directory")
	}
	if policy.Enforcement != EnforcementV1 {
		return nil, fmt.Errorf("storage capacity enforcement is not enabled")
	}
	return &Governor{root: root, policy: policy, free: filesystemFreeBytes, usage: directoryUsage}, nil
}

func (g *Governor) PackageMaxBytes() int64 { return g.policy.CapsBytes.UpdatePackageMax }

func (g *Governor) Snapshot() (Snapshot, error) {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.snapshotLocked()
}

func (g *Governor) Reserve(bytes int64) (*Reservation, Snapshot, error) {
	if bytes <= 0 {
		return nil, Snapshot{}, fmt.Errorf("package size must be greater than zero")
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	snapshot, err := g.snapshotLocked()
	if err != nil {
		return nil, Snapshot{}, err
	}
	if bytes > snapshot.UpdatePackageMax || snapshot.UpdateUsedBytes+snapshot.UpdateReserved+bytes > snapshot.UpdateCapBytes || snapshot.FreeBytes-snapshot.UpdateReserved-bytes < snapshot.GlobalFreeFloor {
		return nil, snapshot, fmt.Errorf("%w: requested=%d free=%d floor=%d used=%d reserved=%d cap=%d", ErrInsufficientSpace, bytes, snapshot.FreeBytes, snapshot.GlobalFreeFloor, snapshot.UpdateUsedBytes, snapshot.UpdateReserved, snapshot.UpdateCapBytes)
	}
	g.reserved += bytes
	snapshot.UpdateReserved = g.reserved
	snapshot.Pressure = pressureFor(snapshot.FreeBytes-g.reserved, snapshot.GlobalFreeFloor)
	return &Reservation{governor: g, bytes: bytes}, snapshot, nil
}

func (r *Reservation) Release() {
	if r == nil || r.governor == nil {
		return
	}
	r.once.Do(func() {
		r.governor.mu.Lock()
		defer r.governor.mu.Unlock()
		r.governor.reserved -= r.bytes
		if r.governor.reserved < 0 {
			r.governor.reserved = 0
		}
	})
}

func (g *Governor) snapshotLocked() (Snapshot, error) {
	free, err := g.free(g.root)
	if err != nil {
		return Snapshot{}, fmt.Errorf("measure filesystem free bytes: %w", err)
	}
	used, err := g.usage(g.root)
	if err != nil {
		return Snapshot{}, fmt.Errorf("measure update staging usage: %w", err)
	}
	return Snapshot{
		Pressure:         pressureFor(free-g.reserved, g.policy.CapsBytes.GlobalFreeFloor),
		FreeBytes:        free,
		GlobalFreeFloor:  g.policy.CapsBytes.GlobalFreeFloor,
		UpdateUsedBytes:  used,
		UpdateReserved:   g.reserved,
		UpdateCapBytes:   g.policy.CapsBytes.UpdateStaging,
		UpdatePackageMax: g.policy.CapsBytes.UpdatePackageMax,
	}, nil
}

func pressureFor(available, floor int64) string {
	margin := available - floor
	switch {
	case margin < 0:
		return "emergency"
	case margin < floor/4:
		return "critical"
	case margin < floor:
		return "warning"
	default:
		return "normal"
	}
}

func directoryUsage(root string) (int64, error) {
	var total int64
	err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		info, err := entry.Info()
		if err != nil {
			return err
		}
		if info.Mode()&os.ModeSymlink != 0 {
			return fmt.Errorf("symlink is not allowed in update staging: %s", path)
		}
		if info.Mode().IsRegular() {
			total += info.Size()
		}
		return nil
	})
	return total, err
}