204 lines
8.3 KiB
Go
204 lines
8.3 KiB
Go
package spool
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/klauspost/compress/zstd"
|
|
rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
|
|
"github.com/rvbox/rvbox/internal/domain"
|
|
"google.golang.org/protobuf/proto"
|
|
)
|
|
|
|
const MaxOutputChunkBytes = 64 << 10
|
|
|
|
var ErrPinnedOutput = errors.New("client output is pinned awaiting server acknowledgement")
|
|
|
|
type OutputInput struct {
|
|
Stream rvboxv1.StreamKind
|
|
Raw []byte
|
|
ObservedAt time.Time
|
|
}
|
|
|
|
// AppendOutput stores one raw pipe chunk as a Zstandard OutputChunk. Under
|
|
// offline pressure it replaces only the oldest unsequenced output with an
|
|
// ordered loss marker; assigned wire events are immutable and never evicted.
|
|
func (store *Store) AppendOutput(ctx context.Context, issueUUID domain.UUID, input OutputInput) (Event, error) {
|
|
if (input.Stream != rvboxv1.StreamKind_STREAM_STDOUT && input.Stream != rvboxv1.StreamKind_STREAM_STDERR) || uint64(len(input.Raw)) > MaxOutputChunkBytes || input.ObservedAt.IsZero() {
|
|
return Event{}, errors.New("invalid raw output chunk")
|
|
}
|
|
payload, err := encodeOutputChunk(input.Stream, input.Raw)
|
|
if err != nil {
|
|
return Event{}, err
|
|
}
|
|
eventInput := EventInput{Kind: EventKindOutput, Compression: 2, RawBytes: uint64(len(input.Raw)), Payload: payload, Output: true, CreatedAt: input.ObservedAt}
|
|
for {
|
|
event, err := store.AppendEvent(ctx, issueUUID, eventInput)
|
|
if !errors.Is(err, ErrCapacityExhausted) {
|
|
return event, err
|
|
}
|
|
rotated, rotateErr := store.RotateOldestUnassignedOutput(ctx, issueUUID, input.ObservedAt)
|
|
if rotateErr != nil {
|
|
return Event{}, rotateErr
|
|
}
|
|
if !rotated {
|
|
return Event{}, err
|
|
}
|
|
}
|
|
}
|
|
|
|
// RotateOldestUnassignedOutput replaces exactly one earliest contiguous run.
|
|
// Retaining the first removed local ordinal for the marker preserves durable
|
|
// command-local order while leaving every assigned event pinned.
|
|
func (store *Store) RotateOldestUnassignedOutput(ctx context.Context, issueUUID domain.UUID, observedAt time.Time) (bool, error) {
|
|
if !validUUID(issueUUID) || observedAt.IsZero() {
|
|
return false, errors.New("invalid output rotation request")
|
|
}
|
|
tx, err := store.db.BeginTx(ctx, nil)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
defer tx.Rollback()
|
|
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 false, ErrUnknownCommand
|
|
}
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
rows, err := tx.QueryContext(ctx, `SELECT local_ordinal, raw_bytes, charged_bytes, payload FROM events WHERE issue_uuid = ? AND event_seq IS NULL AND output = 1 ORDER BY local_ordinal`, issueUUID[:])
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
var run []outputRow
|
|
for rows.Next() {
|
|
var current outputRow
|
|
if err := rows.Scan(¤t.ordinal, ¤t.rawBytes, ¤t.charge, ¤t.payload); err != nil {
|
|
_ = rows.Close()
|
|
return false, err
|
|
}
|
|
if len(run) > 0 && current.ordinal != run[len(run)-1].ordinal+1 {
|
|
break
|
|
}
|
|
run = append(run, current)
|
|
}
|
|
if err := rows.Close(); err != nil {
|
|
return false, err
|
|
}
|
|
if len(run) == 0 {
|
|
return false, ErrPinnedOutput
|
|
}
|
|
markerPayload, removedRaw, removedCompressed, removedCharge, err := outputLossMarker(run)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
markerCharge, err := EstimateCharge(ChargeInput{EncodedBytes: uint64(len(markerPayload)), SQLiteRows: 1, IndexEntries: 2})
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if removedCharge > totalCharged || removedCharge > outputCharged {
|
|
return false, fmt.Errorf("client spool output charge counter mismatch")
|
|
}
|
|
clientTotal, err := clientTotalCharge(ctx, tx)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
if removedCharge > clientTotal {
|
|
return false, fmt.Errorf("client spool aggregate charge counter mismatch")
|
|
}
|
|
remainingTotal := totalCharged - removedCharge
|
|
remainingOutput := outputCharged - removedCharge
|
|
remainingClient := clientTotal - removedCharge
|
|
decision, err := CheckReservation(store.quotaLimits, ReservationState{CommandOutputCharged: remainingOutput, CommandTotalCharged: remainingTotal, ClientTotalCharged: remainingClient, CloseoutRemaining: closeout}, ReservationRequest{ChargedBytes: markerCharge})
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
for _, current := range run {
|
|
if _, err := tx.ExecContext(ctx, `DELETE FROM events WHERE issue_uuid = ? AND local_ordinal = ? AND event_seq IS NULL`, issueUUID[:], current.ordinal); err != nil {
|
|
return false, err
|
|
}
|
|
}
|
|
digest := immutableDigest(markerPayload)
|
|
if _, 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 (?, ?, ?, 2, ?, ?, 0, ?, ?, ?)`, issueUUID[:], run[0].ordinal, EventKindOutputTruncation, removedRaw, markerCharge, markerPayload, digest[:], observedAt.UnixNano()); err != nil {
|
|
return false, err
|
|
}
|
|
if _, err := tx.ExecContext(ctx, `UPDATE commands SET output_charged_bytes = ?, total_charged_bytes = ?, closeout_remaining_bytes = ?, next_local_ordinal = ? WHERE issue_uuid = ?`, decision.CommandOutputCharged, decision.CommandTotalCharged, decision.CloseoutRemaining, nextOrdinal, issueUUID[:]); err != nil {
|
|
return false, err
|
|
}
|
|
if err := updateClientTotalCharge(ctx, tx, decision.ClientTotalCharged); err != nil {
|
|
return false, err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return false, err
|
|
}
|
|
_ = removedCompressed // retained in the durable marker payload.
|
|
return true, nil
|
|
}
|
|
|
|
type outputRow struct {
|
|
ordinal uint64
|
|
rawBytes uint64
|
|
charge uint64
|
|
payload []byte
|
|
}
|
|
|
|
func encodeOutputChunk(stream rvboxv1.StreamKind, raw []byte) ([]byte, error) {
|
|
encoder, err := zstd.NewWriter(nil, zstd.WithEncoderConcurrency(1))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
compressed := encoder.EncodeAll(raw, nil)
|
|
encoder.Close()
|
|
chunk := &rvboxv1.OutputChunk{Stream: stream, Compression: rvboxv1.Compression_COMPRESSION_ZSTD, Data: compressed, UncompressedSize: uint64(len(raw)), CompressedSize: uint64(len(compressed))}
|
|
return proto.Marshal(chunk)
|
|
}
|
|
|
|
func outputLossMarker(rows []outputRow) ([]byte, uint64, uint64, uint64, error) {
|
|
var raw, compressed, charged uint64
|
|
for _, row := range rows {
|
|
var chunk rvboxv1.OutputChunk
|
|
if err := proto.Unmarshal(row.payload, &chunk); err != nil || chunk.Compression != rvboxv1.Compression_COMPRESSION_ZSTD || chunk.UncompressedSize != row.rawBytes || chunk.CompressedSize != uint64(len(chunk.Data)) {
|
|
return nil, 0, 0, 0, ErrStoredPayloadChecksum
|
|
}
|
|
if raw > ^uint64(0)-row.rawBytes || compressed > ^uint64(0)-chunk.CompressedSize || charged > ^uint64(0)-row.charge {
|
|
return nil, 0, 0, 0, ErrCapacityExhausted
|
|
}
|
|
raw += row.rawBytes
|
|
compressed += chunk.CompressedSize
|
|
charged += row.charge
|
|
}
|
|
marker := &rvboxv1.OutputTruncation{RemovedCompressedBytes: &compressed, RemovedUncompressedBytes: raw, Reason: "CLIENT_SPOOL", Source: rvboxv1.OutputTruncationSource_OUTPUT_TRUNCATION_SOURCE_CLIENT_SPOOL}
|
|
payload, err := proto.Marshal(marker)
|
|
return payload, raw, compressed, charged, err
|
|
}
|
|
|
|
func decodeOutputLossMarker(payload []byte) (*rvboxv1.OutputTruncation, error) {
|
|
var marker rvboxv1.OutputTruncation
|
|
if err := proto.Unmarshal(payload, &marker); err != nil || marker.Source != rvboxv1.OutputTruncationSource_OUTPUT_TRUNCATION_SOURCE_CLIENT_SPOOL || marker.FirstRemovedEventSeq != nil || marker.LastRemovedEventSeq != nil || marker.RemovedCompressedBytes == nil {
|
|
return nil, ErrStoredPayloadChecksum
|
|
}
|
|
return &marker, nil
|
|
}
|
|
|
|
func outputChunkRaw(payload []byte) ([]byte, error) {
|
|
var chunk rvboxv1.OutputChunk
|
|
if err := proto.Unmarshal(payload, &chunk); err != nil || (chunk.Stream != rvboxv1.StreamKind_STREAM_STDOUT && chunk.Stream != rvboxv1.StreamKind_STREAM_STDERR) || chunk.Compression != rvboxv1.Compression_COMPRESSION_ZSTD || chunk.CompressedSize != uint64(len(chunk.Data)) {
|
|
return nil, ErrStoredPayloadChecksum
|
|
}
|
|
decoder, err := zstd.NewReader(bytes.NewReader(chunk.Data), zstd.WithDecoderConcurrency(1), zstd.WithDecoderMaxMemory(MaxOutputChunkBytes+1))
|
|
if err != nil {
|
|
return nil, ErrStoredPayloadChecksum
|
|
}
|
|
defer decoder.Close()
|
|
raw, err := decoder.DecodeAll(chunk.Data, nil)
|
|
if err != nil || uint64(len(raw)) != chunk.UncompressedSize || uint64(len(raw)) > MaxOutputChunkBytes {
|
|
return nil, ErrStoredPayloadChecksum
|
|
}
|
|
return raw, nil
|
|
}
|