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 }