rvbox/internal/server/store/quota.go

package store

import (
	"errors"
	"fmt"
	"math"
)

const (
	ChargeFormulaVersion  uint32 = 1
	ChargeSQLiteRowBytes  uint64 = 192
	ChargeIndexEntryBytes uint64 = 64
)

var ErrCapacityExhausted = errors.New("storage capacity exhausted")

type FreeSpaceProbe interface {
	AvailableBytes(path string) (uint64, error)
}

type FreeSpaceProbeFunc func(path string) (uint64, error)

func (function FreeSpaceProbeFunc) AvailableBytes(path string) (uint64, error) { return function(path) }

type CapacityTier string

const (
	CapacityTierHardMaximum CapacityTier = "hard_maximum"
	CapacityTierCommandOut  CapacityTier = "command_output"
	CapacityTierCommand     CapacityTier = "command_total"
	CapacityTierClient      CapacityTier = "client_total"
	CapacityTierServer      CapacityTier = "server_total"
	CapacityTierFilesystem  CapacityTier = "filesystem_floor"
)

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 ChargeInput struct {
	EncodedBytes uint64
	SQLiteRows   uint64
	IndexEntries uint64
}

func EstimateCharge(input ChargeInput) (uint64, error) {
	rowCharge, overflow := multiplyChecked(input.SQLiteRows, ChargeSQLiteRowBytes)
	if overflow {
		return 0, &CapacityError{Tier: CapacityTierHardMaximum, Requested: math.MaxUint64}
	}
	indexCharge, overflow := multiplyChecked(input.IndexEntries, ChargeIndexEntryBytes)
	if overflow {
		return 0, &CapacityError{Tier: CapacityTierHardMaximum, Requested: math.MaxUint64}
	}
	result, overflow := addChecked(input.EncodedBytes, rowCharge)
	if !overflow {
		result, overflow = addChecked(result, indexCharge)
	}
	if overflow {
		return 0, &CapacityError{Tier: CapacityTierHardMaximum, Requested: math.MaxUint64}
	}
	return result, nil
}

type QuotaLimits struct {
	HardAllocationBytes  uint64
	CommandOutputBytes   uint64
	CommandTotalBytes    uint64
	ClientTotalBytes     uint64
	ServerTotalBytes     uint64
	CloseoutReserveBytes uint64
	FilesystemFloorBytes uint64
}

func DefaultQuotaLimits() QuotaLimits {
	return QuotaLimits{
		HardAllocationBytes: DefaultSegmentLimit, CommandOutputBytes: 10 << 20, CommandTotalBytes: 32 << 20,
		ClientTotalBytes: 256 << 20, ServerTotalBytes: 4 << 30, CloseoutReserveBytes: 64 << 10,
		FilesystemFloorBytes: 256 << 20,
	}
}

func (limits QuotaLimits) Validate() error {
	if limits.HardAllocationBytes == 0 || limits.CommandOutputBytes == 0 || limits.CommandTotalBytes <= limits.CommandOutputBytes || limits.ClientTotalBytes <= limits.CommandTotalBytes || limits.ServerTotalBytes <= limits.ClientTotalBytes || limits.CloseoutReserveBytes == 0 || limits.CloseoutReserveBytes >= limits.CommandTotalBytes || limits.FilesystemFloorBytes == 0 {
		return errors.New("invalid quota limits")
	}
	return nil
}

type ReservationState struct {
	CommandOutputCharged uint64
	CommandTotalCharged  uint64
	ClientTotalCharged   uint64
	ServerTotalCharged   uint64
	CloseoutRemaining    uint64
	FilesystemFreeBytes  uint64
}

type ReservationRequest struct {
	ChargedBytes  uint64
	PhysicalBytes uint64
	Output        bool
	UseCloseout   bool
}

type ReservationDecision struct {
	CommandOutputCharged uint64
	CommandTotalCharged  uint64
	ClientTotalCharged   uint64
	ServerTotalCharged   uint64
	CloseoutRemaining    uint64
}

func CheckReservation(limits QuotaLimits, state ReservationState, request ReservationRequest) (ReservationDecision, error) {
	var result ReservationDecision
	if err := limits.Validate(); err != nil {
		return result, err
	}
	if request.ChargedBytes == 0 || request.ChargedBytes > limits.HardAllocationBytes {
		return result, capacityFailure(CapacityTierHardMaximum, request.ChargedBytes, limits.HardAllocationBytes)
	}
	output, overflow := addChecked(state.CommandOutputCharged, request.ChargedBytes)
	if request.Output && (overflow || output > limits.CommandOutputBytes) {
		return result, capacityFailure(CapacityTierCommandOut, request.ChargedBytes, available(limits.CommandOutputBytes, state.CommandOutputCharged))
	}
	command, overflow := addChecked(state.CommandTotalCharged, request.ChargedBytes)
	if overflow {
		return result, capacityFailure(CapacityTierCommand, request.ChargedBytes, available(limits.CommandTotalBytes, state.CommandTotalCharged))
	}
	closeoutRemaining := state.CloseoutRemaining
	if request.UseCloseout {
		if request.ChargedBytes >= closeoutRemaining {
			closeoutRemaining = 0
		} else {
			closeoutRemaining -= request.ChargedBytes
		}
	} else if _, overflow := addChecked(command, closeoutRemaining); overflow || command+closeoutRemaining > limits.CommandTotalBytes {
		return result, capacityFailure(CapacityTierCommand, request.ChargedBytes, availableForNormalCommand(limits.CommandTotalBytes, state.CommandTotalCharged, closeoutRemaining))
	}
	if command > limits.CommandTotalBytes {
		return result, capacityFailure(CapacityTierCommand, request.ChargedBytes, available(limits.CommandTotalBytes, state.CommandTotalCharged))
	}
	client, overflow := addChecked(state.ClientTotalCharged, request.ChargedBytes)
	if overflow || client > limits.ClientTotalBytes {
		return result, capacityFailure(CapacityTierClient, request.ChargedBytes, available(limits.ClientTotalBytes, state.ClientTotalCharged))
	}
	server, overflow := addChecked(state.ServerTotalCharged, request.ChargedBytes)
	if overflow || server > limits.ServerTotalBytes {
		return result, capacityFailure(CapacityTierServer, request.ChargedBytes, available(limits.ServerTotalBytes, state.ServerTotalCharged))
	}
	requiredFree, overflow := addChecked(limits.FilesystemFloorBytes, request.PhysicalBytes)
	if overflow || state.FilesystemFreeBytes < requiredFree {
		return result, capacityFailure(CapacityTierFilesystem, request.PhysicalBytes, available(state.FilesystemFreeBytes, limits.FilesystemFloorBytes))
	}
	if !request.Output {
		output = state.CommandOutputCharged
	}
	return ReservationDecision{
		CommandOutputCharged: output, CommandTotalCharged: command, ClientTotalCharged: client,
		ServerTotalCharged: server, CloseoutRemaining: closeoutRemaining,
	}, nil
}

func capacityFailure(tier CapacityTier, requested, availableBytes uint64) error {
	return &CapacityError{Tier: tier, Requested: requested, Available: availableBytes}
}

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
}

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
}