feat: enforce tiered storage reservations

internal/server/store/event.go

@@ -42,6 +42,8 @@ type EventAppend struct {
 	RawLength        uint64
 	Payload          []byte
 	ImmutableSHA256  [32]byte
+	Output           bool
+	UseCloseout      bool
 }
 
 type EventAppendResult struct {
@@ -76,8 +78,14 @@ func (store *Store) AppendCommandEvent(ctx context.Context, event EventAppend) (
 	if err != nil {
 		return result, err
 	}
-	var lastSequence uint64
-	err = database.QueryRowContext(ctx, `SELECT last_event_seq FROM commands WHERE issue_uuid = ?`, event.IssueUUID[:]).Scan(&lastSequence)
+	var lastSequence, commandOutputCharged, commandCharged, clientCharged, serverCharged, closeoutRemaining uint64
+	var clientID string
+	err = database.QueryRowContext(ctx, `SELECT commands.last_event_seq, commands.output_charged_bytes,
+commands.charged_bytes, commands.closeout_remaining_bytes, commands.client_id, clients.charged_bytes,
+storage_counters.command_charged_bytes
+FROM commands JOIN clients ON clients.client_id = commands.client_id
+JOIN storage_counters ON storage_counters.singleton = 1 WHERE commands.issue_uuid = ?`, event.IssueUUID[:]).Scan(
+		&lastSequence, &commandOutputCharged, &commandCharged, &closeoutRemaining, &clientID, &clientCharged, &serverCharged)
 	if errors.Is(err, sql.ErrNoRows) {
 		return result, ErrCommandNotFound
 	}
@@ -97,6 +105,30 @@ func (store *Store) AppendCommandEvent(ctx context.Context, event EventAppend) (
 	if event.EventSeq != lastSequence+1 {
 		return result, ErrEventSequenceGap
 	}
+	rowsCharged, indexesCharged := uint64(1), uint64(1)
+	if store.willCreateSegment(event.IssueUUID, uint64(len(encoded))) {
+		rowsCharged++
+		indexesCharged += 2
+	}
+	chargedBytes, err := EstimateCharge(ChargeInput{EncodedBytes: uint64(len(encoded)), SQLiteRows: rowsCharged, IndexEntries: indexesCharged})
+	if err != nil {
+		return result, err
+	}
+	freeBytes, err := store.freeSpaceProbe.AvailableBytes(store.dataDir)
+	if err != nil {
+		return result, fmt.Errorf("check filesystem capacity: %w", err)
+	}
+	reservation, err := CheckReservation(store.quotaLimits, ReservationState{
+		CommandOutputCharged: commandOutputCharged, CommandTotalCharged: commandCharged,
+		ClientTotalCharged: clientCharged, ServerTotalCharged: serverCharged,
+		CloseoutRemaining: closeoutRemaining, FilesystemFreeBytes: freeBytes,
+	}, ReservationRequest{
+		ChargedBytes: chargedBytes, PhysicalBytes: uint64(len(encoded)), Output: event.Output || event.Stream > 0,
+		UseCloseout: event.UseCloseout,
+	})
+	if err != nil {
+		return result, err
+	}
 
 	active, err := store.segmentForAppend(ctx, database, event.IssueUUID, uint64(len(encoded)))
 	if err != nil {
@@ -153,8 +185,11 @@ payload, segment_ordinal, segment_record_offset, segment_record_length, immutabl
 	}
 	if err == nil {
 		var update sql.Result
-		update, err = tx.ExecContext(ctx, `UPDATE commands SET last_event_seq = ?, retained_compressed_bytes = retained_compressed_bytes + ?
-WHERE issue_uuid = ? AND last_event_seq = ?`, event.EventSeq, len(event.Payload), event.IssueUUID[:], lastSequence)
+		update, err = tx.ExecContext(ctx, `UPDATE commands SET last_event_seq = ?,
+retained_compressed_bytes = retained_compressed_bytes + ?, output_charged_bytes = ?, charged_bytes = ?,
+closeout_remaining_bytes = ? WHERE issue_uuid = ? AND last_event_seq = ?`, event.EventSeq, len(event.Payload),
+			reservation.CommandOutputCharged, reservation.CommandTotalCharged, reservation.CloseoutRemaining,
+			event.IssueUUID[:], lastSequence)
 		if err == nil {
 			var affected int64
 			affected, err = update.RowsAffected()
@@ -163,6 +198,28 @@ WHERE issue_uuid = ? AND last_event_seq = ?`, event.EventSeq, len(event.Payload)
 			}
 		}
 	}
+	if err == nil {
+		var update sql.Result
+		update, err = tx.ExecContext(ctx, `UPDATE clients SET charged_bytes = ? WHERE client_id = ? AND charged_bytes = ?`, reservation.ClientTotalCharged, clientID, clientCharged)
+		if err == nil {
+			var affected int64
+			affected, err = update.RowsAffected()
+			if err == nil && affected != 1 {
+				err = errors.New("client quota counter changed unexpectedly")
+			}
+		}
+	}
+	if err == nil {
+		var update sql.Result
+		update, err = tx.ExecContext(ctx, `UPDATE storage_counters SET command_charged_bytes = ? WHERE singleton = 1 AND command_charged_bytes = ?`, reservation.ServerTotalCharged, serverCharged)
+		if err == nil {
+			var affected int64
+			affected, err = update.RowsAffected()
+			if err == nil && affected != 1 {
+				err = errors.New("server quota counter changed unexpectedly")
+			}
+		}
+	}
 	if err != nil {
 		_ = tx.Rollback()
 		store.invalidateActiveSegment(event.IssueUUID)
@@ -180,6 +237,11 @@ WHERE issue_uuid = ? AND last_event_seq = ?`, event.EventSeq, len(event.Payload)
 	return result, nil
 }
 
+func (store *Store) willCreateSegment(owner [16]byte, recordLength uint64) bool {
+	current := store.activeSegments[owner]
+	return current == nil || (current.segment.offset != 0 && current.segment.offset+recordLength > store.segmentTarget)
+}
+
 func (store *Store) segmentForAppend(ctx context.Context, database *sql.DB, owner [16]byte, recordLength uint64) (*activeSegment, error) {
 	if current := store.activeSegments[owner]; current != nil {
 		if current.segment.offset == 0 || current.segment.offset+recordLength <= store.segmentTarget {