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