feat: enforce client spool quota tiers
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user