feat: enforce client spool quota tiers

This commit is contained in:
2026-09-06 06:32:22 +00:00
parent c348d6bb16
commit 6539b9e965
6 changed files with 363 additions and 11 deletions
+86 -9
View File
@@ -28,6 +28,8 @@ type EventInput struct {
Compression uint32
RawBytes uint64
Payload []byte
Output bool
UseCloseout bool
CreatedAt time.Time
}
@@ -72,10 +74,25 @@ func (store *Store) AcceptCommand(ctx context.Context, command Command, accepted
if err != sql.ErrNoRows {
return Acceptance{}, err
}
_, err = tx.ExecContext(ctx, `INSERT INTO commands(issue_uuid, immutable_sha256, command_revision, phase, terminal, accepted_at) VALUES (?, ?, ?, ?, ?, ?)`, command.IssueUUID[:], command.ImmutableSHA256[:], command.Revision, command.Phase, boolInt(command.Terminal), acceptedAt.UnixNano())
charge, err := EstimateCharge(ChargeInput{SQLiteRows: 1, IndexEntries: 2})
if err != nil {
return Acceptance{}, err
}
clientTotal, err := clientTotalCharge(ctx, tx)
if err != nil {
return Acceptance{}, err
}
decision, err := CheckReservation(store.quotaLimits, ReservationState{ClientTotalCharged: clientTotal, CloseoutRemaining: store.quotaLimits.CloseoutReserveBytes}, ReservationRequest{ChargedBytes: charge})
if err != nil {
return Acceptance{}, err
}
_, err = tx.ExecContext(ctx, `INSERT INTO commands(issue_uuid, immutable_sha256, command_revision, phase, terminal, total_charged_bytes, closeout_remaining_bytes, accepted_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, command.IssueUUID[:], command.ImmutableSHA256[:], command.Revision, command.Phase, boolInt(command.Terminal), decision.CommandTotalCharged, decision.CloseoutRemaining, acceptedAt.UnixNano())
if err != nil {
return Acceptance{}, err
}
if err := updateClientTotalCharge(ctx, tx, decision.ClientTotalCharged); err != nil {
return Acceptance{}, err
}
if err := tx.Commit(); err != nil {
return Acceptance{}, err
}
@@ -93,19 +110,38 @@ func (store *Store) AppendEvent(ctx context.Context, issueUUID domain.UUID, inpu
return Event{}, err
}
defer tx.Rollback()
var ordinal uint64
err = tx.QueryRowContext(ctx, `UPDATE commands SET next_local_ordinal = next_local_ordinal + 1 WHERE issue_uuid = ? RETURNING next_local_ordinal - 1`, issueUUID[:]).Scan(&ordinal)
var nextOrdinal, outputCharged, totalCharged, closeout uint64
err = tx.QueryRowContext(ctx, `SELECT next_local_ordinal, output_charged_bytes, total_charged_bytes, closeout_remaining_bytes FROM commands WHERE issue_uuid = ?`, issueUUID[:]).Scan(&nextOrdinal, &outputCharged, &totalCharged, &closeout)
if err == sql.ErrNoRows {
return Event{}, ErrUnknownCommand
}
if err != nil {
return Event{}, err
}
digest := immutableDigest(input.Payload)
_, err = tx.ExecContext(ctx, `INSERT INTO events(issue_uuid, local_ordinal, event_kind, compression, raw_bytes, payload, payload_sha256, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, issueUUID[:], ordinal, input.Kind, input.Compression, input.RawBytes, input.Payload, digest[:], input.CreatedAt.UnixNano())
charge, err := EstimateCharge(ChargeInput{EncodedBytes: uint64(len(input.Payload)), SQLiteRows: 1, IndexEntries: 2})
if err != nil {
return Event{}, err
}
clientTotal, err := clientTotalCharge(ctx, tx)
if err != nil {
return Event{}, err
}
decision, err := CheckReservation(store.quotaLimits, ReservationState{CommandOutputCharged: outputCharged, CommandTotalCharged: totalCharged, ClientTotalCharged: clientTotal, CloseoutRemaining: closeout}, ReservationRequest{ChargedBytes: charge, Output: input.Output, UseCloseout: input.UseCloseout})
if err != nil {
return Event{}, err
}
ordinal := nextOrdinal
digest := immutableDigest(input.Payload)
_, err = tx.ExecContext(ctx, `INSERT INTO events(issue_uuid, local_ordinal, event_kind, compression, raw_bytes, charged_bytes, output, payload, payload_sha256, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, issueUUID[:], ordinal, input.Kind, input.Compression, input.RawBytes, charge, boolInt(input.Output), input.Payload, digest[:], input.CreatedAt.UnixNano())
if err != nil {
return Event{}, err
}
if _, err = tx.ExecContext(ctx, `UPDATE commands SET next_local_ordinal = ?, output_charged_bytes = ?, total_charged_bytes = ?, closeout_remaining_bytes = ? WHERE issue_uuid = ?`, ordinal+1, decision.CommandOutputCharged, decision.CommandTotalCharged, decision.CloseoutRemaining, issueUUID[:]); err != nil {
return Event{}, err
}
if err := updateClientTotalCharge(ctx, tx, decision.ClientTotalCharged); err != nil {
return Event{}, err
}
if err := tx.Commit(); err != nil {
return Event{}, err
}
@@ -185,8 +221,8 @@ func (store *Store) Ack(ctx context.Context, issueUUID domain.UUID, eventSeq uin
return err
}
defer tx.Rollback()
var lastAck, next uint64
err = tx.QueryRowContext(ctx, `SELECT last_server_ack, next_event_seq FROM commands WHERE issue_uuid = ?`, issueUUID[:]).Scan(&lastAck, &next)
var lastAck, next, outputCharged, totalCharged uint64
err = tx.QueryRowContext(ctx, `SELECT last_server_ack, next_event_seq, output_charged_bytes, total_charged_bytes FROM commands WHERE issue_uuid = ?`, issueUUID[:]).Scan(&lastAck, &next, &outputCharged, &totalCharged)
if err == sql.ErrNoRows {
return ErrUnknownCommand
}
@@ -199,10 +235,27 @@ func (store *Store) Ack(ctx context.Context, issueUUID domain.UUID, eventSeq uin
if eventSeq >= next {
return ErrInvalidEventAck
}
var released, releasedOutput uint64
if err := tx.QueryRowContext(ctx, `SELECT COALESCE(sum(charged_bytes), 0), COALESCE(sum(CASE WHEN output = 1 THEN charged_bytes ELSE 0 END), 0) FROM events WHERE issue_uuid = ? AND event_seq <= ?`, issueUUID[:], eventSeq).Scan(&released, &releasedOutput); err != nil {
return err
}
if released > totalCharged || releasedOutput > outputCharged {
return fmt.Errorf("client spool charge counter mismatch")
}
clientTotal, err := clientTotalCharge(ctx, tx)
if err != nil {
return err
}
if released > clientTotal {
return fmt.Errorf("client spool aggregate charge counter mismatch")
}
if _, err := tx.ExecContext(ctx, `DELETE FROM events WHERE issue_uuid = ? AND event_seq <= ?`, issueUUID[:], eventSeq); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, `UPDATE commands SET last_server_ack = ? WHERE issue_uuid = ?`, eventSeq, issueUUID[:]); err != nil {
if _, err := tx.ExecContext(ctx, `UPDATE commands SET last_server_ack = ?, output_charged_bytes = ?, total_charged_bytes = ? WHERE issue_uuid = ?`, eventSeq, outputCharged-releasedOutput, totalCharged-released, issueUUID[:]); err != nil {
return err
}
if err := updateClientTotalCharge(ctx, tx, clientTotal-released); err != nil {
return err
}
return tx.Commit()
@@ -239,7 +292,8 @@ func (store *Store) CleanupTerminal(ctx context.Context, issueUUID domain.UUID,
defer tx.Rollback()
var hash []byte
var terminal int
err = tx.QueryRowContext(ctx, `SELECT immutable_sha256, terminal FROM commands WHERE issue_uuid = ?`, issueUUID[:]).Scan(&hash, &terminal)
var totalCharged uint64
err = tx.QueryRowContext(ctx, `SELECT immutable_sha256, terminal, total_charged_bytes FROM commands WHERE issue_uuid = ?`, issueUUID[:]).Scan(&hash, &terminal, &totalCharged)
if err == sql.ErrNoRows {
return ErrUnknownCommand
}
@@ -256,12 +310,22 @@ func (store *Store) CleanupTerminal(ctx context.Context, issueUUID domain.UUID,
if pending != 0 {
return ErrEventsPending
}
clientTotal, err := clientTotalCharge(ctx, tx)
if err != nil {
return err
}
if totalCharged > clientTotal {
return fmt.Errorf("client spool aggregate charge counter mismatch")
}
if _, err := tx.ExecContext(ctx, `INSERT INTO command_tombstones(issue_uuid, immutable_sha256, acknowledged_at) VALUES (?, ?, ?)`, issueUUID[:], hash, acknowledgedAt.UnixNano()); err != nil {
return err
}
if _, err := tx.ExecContext(ctx, `DELETE FROM commands WHERE issue_uuid = ?`, issueUUID[:]); err != nil {
return err
}
if err := updateClientTotalCharge(ctx, tx, clientTotal-totalCharged); err != nil {
return err
}
trim, err := tombstonesToTrim(ctx, tx, store.tombstoneLimit)
if err != nil {
return err
@@ -336,6 +400,19 @@ func boolInt(value bool) int {
return 0
}
func clientTotalCharge(ctx context.Context, tx *sql.Tx) (uint64, error) {
var total uint64
if err := tx.QueryRowContext(ctx, `SELECT client_total_charged_bytes FROM spool_counters WHERE singleton = 1`).Scan(&total); err != nil {
return 0, err
}
return total, nil
}
func updateClientTotalCharge(ctx context.Context, tx *sql.Tx, total uint64) error {
_, err := tx.ExecContext(ctx, `UPDATE spool_counters SET client_total_charged_bytes = ? WHERE singleton = 1`, total)
return err
}
func validUUID(value domain.UUID) bool {
_, err := domain.ParseUUIDv7(value.String())
return err == nil
+11
View File
@@ -60,6 +60,9 @@ CREATE TABLE commands (
command_revision INTEGER NOT NULL CHECK(command_revision > 0),
phase INTEGER NOT NULL CHECK(phase BETWEEN 1 AND 11),
terminal INTEGER NOT NULL DEFAULT 0 CHECK(terminal IN (0, 1)),
output_charged_bytes INTEGER NOT NULL DEFAULT 0 CHECK(output_charged_bytes >= 0),
total_charged_bytes INTEGER NOT NULL CHECK(total_charged_bytes >= 0),
closeout_remaining_bytes INTEGER NOT NULL CHECK(closeout_remaining_bytes >= 0),
next_local_ordinal INTEGER NOT NULL DEFAULT 1 CHECK(next_local_ordinal > 0),
next_event_seq INTEGER NOT NULL DEFAULT 1 CHECK(next_event_seq > 0),
last_server_ack INTEGER NOT NULL DEFAULT 0 CHECK(last_server_ack >= 0),
@@ -72,6 +75,8 @@ CREATE TABLE events (
event_kind INTEGER NOT NULL CHECK(event_kind > 0),
compression INTEGER NOT NULL CHECK(compression IN (1, 2)),
raw_bytes INTEGER NOT NULL CHECK(raw_bytes >= 0),
charged_bytes INTEGER NOT NULL CHECK(charged_bytes > 0),
output INTEGER NOT NULL CHECK(output IN (0, 1)),
payload BLOB NOT NULL,
payload_sha256 BLOB NOT NULL CHECK(length(payload_sha256) = 32),
created_at INTEGER NOT NULL,
@@ -79,6 +84,12 @@ CREATE TABLE events (
UNIQUE(issue_uuid, event_seq)
) STRICT, WITHOUT ROWID;
CREATE INDEX events_send_window ON events(issue_uuid, event_seq, local_ordinal);
CREATE TABLE spool_counters (
singleton INTEGER PRIMARY KEY CHECK(singleton = 1),
client_total_charged_bytes INTEGER NOT NULL CHECK(client_total_charged_bytes >= 0),
charge_version INTEGER NOT NULL CHECK(charge_version = 1)
) STRICT;
INSERT INTO spool_counters(singleton, client_total_charged_bytes, charge_version) VALUES (1, 0, 1);
CREATE TABLE command_tombstones (
issue_uuid BLOB PRIMARY KEY CHECK(length(issue_uuid) = 16),
immutable_sha256 BLOB NOT NULL CHECK(length(immutable_sha256) = 32),
+172
View File
@@ -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
}
+69
View File
@@ -0,0 +1,69 @@
package spool
import (
"context"
"errors"
"path/filepath"
"testing"
"time"
)
func TestCheckReservationQuotaTiers_BH_RET_01(t *testing.T) {
t.Parallel()
limits := QuotaLimits{HardAllocationBytes: 1_000, CommandOutputBytes: 1_200, CommandTotalBytes: 2_000, ClientTotalBytes: 3_000, CloseoutReserveBytes: 400}
cases := []struct {
name string
state ReservationState
request ReservationRequest
tier CapacityTier
}{
{name: "hard allocation", request: ReservationRequest{ChargedBytes: 1_001}, tier: CapacityTierHardMaximum},
{name: "output", state: ReservationState{CommandOutputCharged: 1_100}, request: ReservationRequest{ChargedBytes: 101, Output: true}, tier: CapacityTierCommandOut},
{name: "normal preserves closeout", state: ReservationState{CommandTotalCharged: 1_500, CloseoutRemaining: 400}, request: ReservationRequest{ChargedBytes: 101}, tier: CapacityTierCommand},
{name: "client", state: ReservationState{ClientTotalCharged: 2_900, CloseoutRemaining: 400}, request: ReservationRequest{ChargedBytes: 101}, tier: CapacityTierClient},
}
for _, test := range cases {
t.Run(test.name, func(t *testing.T) {
_, err := CheckReservation(limits, test.state, test.request)
var capacity *CapacityError
if !errors.As(err, &capacity) || capacity.Tier != test.tier {
t.Fatalf("CheckReservation error = %v, want tier %q", err, test.tier)
}
})
}
decision, err := CheckReservation(limits, ReservationState{CommandTotalCharged: 1_500, CloseoutRemaining: 400}, ReservationRequest{ChargedBytes: 400, UseCloseout: true})
if err != nil || decision.CloseoutRemaining != 0 || decision.CommandTotalCharged != 1_900 {
t.Fatalf("closeout reservation = %#v, %v", decision, err)
}
}
func TestSpoolAcknowledgementReleasesQuota_HP_CLIENT_08(t *testing.T) {
t.Parallel()
ctx := context.Background()
limits := QuotaLimits{HardAllocationBytes: 1_000, CommandOutputBytes: 800, CommandTotalBytes: 2_000, ClientTotalBytes: 3_000, CloseoutReserveBytes: 400}
store, err := Open(ctx, Options{DataDir: filepath.Join(t.TempDir(), "spool"), BusyTimeout: time.Second, QuotaLimits: limits})
if err != nil {
t.Fatal(err)
}
t.Cleanup(func() { _ = store.Close() })
issue := testUUID(t, "019c46f1-1d02-7000-8000-000000000021")
if _, err := store.AcceptCommand(ctx, testCommand(issue, []byte("quota")), time.Now().UTC()); err != nil {
t.Fatal(err)
}
input := EventInput{Kind: 1, Compression: 1, RawBytes: 200, Payload: make([]byte, 200), Output: true, CreatedAt: time.Now().UTC()}
if _, err := store.AppendEvent(ctx, issue, input); err != nil {
t.Fatal(err)
}
if _, err := store.AppendEvent(ctx, issue, input); !errors.Is(err, ErrCapacityExhausted) {
t.Fatalf("second output append error = %v, want ErrCapacityExhausted", err)
}
if _, err := store.AssignSendWindow(ctx, issue, 1, 1<<20); err != nil {
t.Fatal(err)
}
if err := store.Ack(ctx, issue, 1); err != nil {
t.Fatal(err)
}
if _, err := store.AppendEvent(ctx, issue, input); err != nil {
t.Fatalf("append after acknowledgement did not release quota: %v", err)
}
}
+9 -1
View File
@@ -37,6 +37,7 @@ type Options struct {
DataDir string
BusyTimeout time.Duration
TombstoneLimit uint64
QuotaLimits QuotaLimits
}
type Store struct {
@@ -45,6 +46,7 @@ type Store struct {
dataDir string
identity domain.UUID
tombstoneLimit uint64
quotaLimits QuotaLimits
mu sync.Mutex
}
@@ -61,6 +63,12 @@ func Open(ctx context.Context, options Options) (*Store, error) {
if options.TombstoneLimit == 0 {
options.TombstoneLimit = DefaultTombstoneLimit
}
if options.QuotaLimits == (QuotaLimits{}) {
options.QuotaLimits = DefaultQuotaLimits()
}
if err := options.QuotaLimits.Validate(); err != nil {
return nil, err
}
if err := ensurePrivateDirectory(options.DataDir); err != nil {
return nil, err
}
@@ -92,7 +100,7 @@ func Open(ctx context.Context, options Options) (*Store, error) {
}
db.SetMaxOpenConns(1)
db.SetMaxIdleConns(1)
store := &Store{db: db, unlock: unlock, dataDir: options.DataDir, identity: identity, tombstoneLimit: options.TombstoneLimit}
store := &Store{db: db, unlock: unlock, dataDir: options.DataDir, identity: identity, tombstoneLimit: options.TombstoneLimit, quotaLimits: options.QuotaLimits}
if err := db.PingContext(ctx); err != nil {
_ = store.Close()
return nil, fmt.Errorf("open client spool SQLite: %w", err)