rvbox/internal/client/spool/quota.go

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
}