feat: enforce client spool quota tiers

2026-09-06 06:32:22 +00:00
parent c348d6bb16
commit 6539b9e965
6 changed files with 363 additions and 11 deletions
@@ -0,0 +1,172 @@
+package spool
+
+import (
+	"errors"
+	"fmt"
+	"math"
+)
+
+const (
+	ChargeFormulaVersion  uint32 = 1
+	ChargeSQLiteRowBytes  uint64 = 192
+	ChargeIndexEntryBytes uint64 = 64
+)
+
+var ErrCapacityExhausted = errors.New("client spool capacity exhausted")
+
+type CapacityTier string
+
+const (
+	CapacityTierHardMaximum CapacityTier = "hard_maximum"
+	CapacityTierCommandOut  CapacityTier = "command_output"
+	CapacityTierCommand     CapacityTier = "command_total"
+	CapacityTierClient      CapacityTier = "client_total"
+)
+
+type CapacityError struct {
+	Tier      CapacityTier
+	Requested uint64
+	Available uint64
+}
+
+func (failure *CapacityError) Error() string {
+	return fmt.Sprintf("%s: tier=%s requested=%d available=%d", ErrCapacityExhausted, failure.Tier, failure.Requested, failure.Available)
+}
+
+func (failure *CapacityError) Unwrap() error { return ErrCapacityExhausted }
+
+type QuotaLimits struct {
+	HardAllocationBytes  uint64
+	CommandOutputBytes   uint64
+	CommandTotalBytes    uint64
+	ClientTotalBytes     uint64
+	CloseoutReserveBytes uint64
+}
+
+func DefaultQuotaLimits() QuotaLimits {
+	return QuotaLimits{
+		HardAllocationBytes: 1 << 20, CommandOutputBytes: 10 << 20, CommandTotalBytes: 32 << 20,
+		ClientTotalBytes: 256 << 20, CloseoutReserveBytes: 64 << 10,
+	}
+}
+
+func (limits QuotaLimits) Validate() error {
+	if limits.HardAllocationBytes == 0 || limits.CommandOutputBytes == 0 || limits.CommandTotalBytes <= limits.CommandOutputBytes || limits.ClientTotalBytes <= limits.CommandTotalBytes || limits.CloseoutReserveBytes == 0 || limits.CloseoutReserveBytes >= limits.CommandTotalBytes {
+		return errors.New("invalid client spool quota limits")
+	}
+	return nil
+}
+
+type ChargeInput struct {
+	EncodedBytes uint64
+	SQLiteRows   uint64
+	IndexEntries uint64
+}
+
+func EstimateCharge(input ChargeInput) (uint64, error) {
+	rows, overflow := multiplyChecked(input.SQLiteRows, ChargeSQLiteRowBytes)
+	if overflow {
+		return 0, capacityFailure(CapacityTierHardMaximum, math.MaxUint64, 0)
+	}
+	indexes, overflow := multiplyChecked(input.IndexEntries, ChargeIndexEntryBytes)
+	if overflow {
+		return 0, capacityFailure(CapacityTierHardMaximum, math.MaxUint64, 0)
+	}
+	result, overflow := addChecked(input.EncodedBytes, rows)
+	if !overflow {
+		result, overflow = addChecked(result, indexes)
+	}
+	if overflow {
+		return 0, capacityFailure(CapacityTierHardMaximum, math.MaxUint64, 0)
+	}
+	return result, nil
+}
+
+type ReservationState struct {
+	CommandOutputCharged uint64
+	CommandTotalCharged  uint64
+	ClientTotalCharged   uint64
+	CloseoutRemaining    uint64
+}
+
+type ReservationRequest struct {
+	ChargedBytes uint64
+	Output       bool
+	UseCloseout  bool
+}
+
+type ReservationDecision struct {
+	CommandOutputCharged uint64
+	CommandTotalCharged  uint64
+	ClientTotalCharged   uint64
+	CloseoutRemaining    uint64
+}
+
+func CheckReservation(limits QuotaLimits, state ReservationState, request ReservationRequest) (ReservationDecision, error) {
+	if err := limits.Validate(); err != nil {
+		return ReservationDecision{}, err
+	}
+	if request.ChargedBytes == 0 || request.ChargedBytes > limits.HardAllocationBytes {
+		return ReservationDecision{}, capacityFailure(CapacityTierHardMaximum, request.ChargedBytes, limits.HardAllocationBytes)
+	}
+	output := state.CommandOutputCharged
+	if request.Output {
+		var overflow bool
+		output, overflow = addChecked(output, request.ChargedBytes)
+		if overflow || output > limits.CommandOutputBytes {
+			return ReservationDecision{}, capacityFailure(CapacityTierCommandOut, request.ChargedBytes, available(limits.CommandOutputBytes, state.CommandOutputCharged))
+		}
+	}
+	command, overflow := addChecked(state.CommandTotalCharged, request.ChargedBytes)
+	if overflow || command > limits.CommandTotalBytes {
+		return ReservationDecision{}, capacityFailure(CapacityTierCommand, request.ChargedBytes, available(limits.CommandTotalBytes, state.CommandTotalCharged))
+	}
+	closeout := state.CloseoutRemaining
+	if request.UseCloseout {
+		if request.ChargedBytes >= closeout {
+			closeout = 0
+		} else {
+			closeout -= request.ChargedBytes
+		}
+	} else if command > limits.CommandTotalBytes-closeout {
+		return ReservationDecision{}, capacityFailure(CapacityTierCommand, request.ChargedBytes, availableForNormalCommand(limits.CommandTotalBytes, state.CommandTotalCharged, closeout))
+	}
+	client, overflow := addChecked(state.ClientTotalCharged, request.ChargedBytes)
+	if overflow || client > limits.ClientTotalBytes {
+		return ReservationDecision{}, capacityFailure(CapacityTierClient, request.ChargedBytes, available(limits.ClientTotalBytes, state.ClientTotalCharged))
+	}
+	return ReservationDecision{CommandOutputCharged: output, CommandTotalCharged: command, ClientTotalCharged: client, CloseoutRemaining: closeout}, nil
+}
+
+func capacityFailure(tier CapacityTier, requested, availableBytes uint64) error {
+	return &CapacityError{Tier: tier, Requested: requested, Available: availableBytes}
+}
+
+func addChecked(left, right uint64) (uint64, bool) {
+	if left > math.MaxUint64-right {
+		return 0, true
+	}
+	return left + right, false
+}
+
+func multiplyChecked(left, right uint64) (uint64, bool) {
+	if left != 0 && right > math.MaxUint64/left {
+		return 0, true
+	}
+	return left * right, false
+}
+
+func available(limit, used uint64) uint64 {
+	if used >= limit {
+		return 0
+	}
+	return limit - used
+}
+
+func availableForNormalCommand(limit, used, closeout uint64) uint64 {
+	remaining := available(limit, used)
+	if closeout >= remaining {
+		return 0
+	}
+	return remaining - closeout
+}