diff --git a/go/e2e/leg_s3_durability_test.go b/go/e2e/leg_s3_durability_test.go index af83aa0ae..82d70dbdd 100644 --- a/go/e2e/leg_s3_durability_test.go +++ b/go/e2e/leg_s3_durability_test.go @@ -205,7 +205,7 @@ func runS3Outage(t *testing.T, ctx context.Context) { if err := runS3TurnSettled(ctx, f, sessionID, channel, "durability-second"); err != nil { t.Fatalf("turn during outage: %v", err) } - // The turn settles before CommitConversationFrame lands; wait for the row segmentsCoveringMarker reads. + // The commit no longer waits on the frozen PUT, but it trails the settle; wait for the row segmentsCoveringMarker reads. if _, err := f.awaitTranscriptPersisted(ctx, st, sessionID, "marker durability-second"); err != nil { t.Fatalf("await outage transcript: %v", err) } @@ -219,6 +219,7 @@ func runS3Outage(t *testing.T, ctx context.Context) { if err := runS3TurnSettled(ctx, f, sessionID, channel, "durability-second"); err != nil { t.Fatalf("turn after outage: %v", err) } + // A pass hung on the frozen PUT ends within its 60s deadline, inside settleTimeout; a later pass archives the reply. waitArchiveContains(t, ctx, f, sessionID, "durability second reply") } diff --git a/go/internal/runnerhub/commit_frame_flush_pgtest_test.go b/go/internal/runnerhub/commit_frame_flush_pgtest_test.go new file mode 100644 index 000000000..c98005c7c --- /dev/null +++ b/go/internal/runnerhub/commit_frame_flush_pgtest_test.go @@ -0,0 +1,66 @@ +//go:build pgtest && unix + +package runnerhub + +import ( + "context" + "errors" + "testing" + "time" + + compassv1 "github.com/RigelBuild/compass/go/gen/compass/v1" +) + +// hungObjectStore holds every PUT until the pass ctx ends, like an object store +// that accepts the connection and never answers. +type hungObjectStore struct { + started chan struct{} +} + +func (h *hungObjectStore) PutSegment(ctx context.Context, _ string, _ []byte) error { + select { + case h.started <- struct{}{}: + default: + } + <-ctx.Done() + return ctx.Err() +} + +func (h *hungObjectStore) GetSegment(context.Context, string) ([]byte, error) { + return nil, errors.New("hung object store: no reads") +} + +// A checkpoint commit returns committed while its flush PUT is hung. +func TestCommitConversationFrameReturnsWhileFlushPutHangs(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + st, agent, ctxB := openTenantBSession(t, ctx) + const sessionID = "sess-b-hung" + if err := st.RecordAgentSession(ctxB, sessionID, agent.ID); err != nil { + t.Fatalf("RecordAgentSession: %v", err) + } + hub := newHubOnly() + hub.SetSessionBindingStore(st) + hub.SetTranscriptStore(st) + if _, err := hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED); err != nil { + t.Fatalf("enroll Runner: %v", err) + } + if _, _, err := st.RecordSessionBinding(ctxB, sessionID, agent.ID, "runner-1"); err != nil { + t.Fatalf("RecordSessionBinding after enroll: %v", err) + } + if _, err := hub.CommitConversationFrame(ctx, "runner-1", transcriptReq(sessionID, "k-1", 1, false, `{"e":1}`)); err != nil { + t.Fatalf("commit delta: %v", err) + } + objects := &hungObjectStore{started: make(chan struct{}, 1)} + st.SetObjectStore(objects) + + resp, err := hub.CommitConversationFrame(ctx, "runner-1", transcriptReq(sessionID, "k-2", 2, true, `{"cp":2}`)) + if err != nil || !resp.GetCommitted() { + t.Fatalf("checkpoint commit = %v, %v; want committed while the PUT hangs", resp, err) + } + select { + case <-objects.started: + case <-ctx.Done(): + t.Fatalf("flush PUT never started: %v", ctx.Err()) + } +} diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 943016e8f..f61025b22 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -261,8 +261,10 @@ type RoutingFabric interface { } // TranscriptStore is the durable transcript surface the hub's commit arm writes -// a relayed transcript_entry frame to (RIG-1667 T4, the durable counterpart to -// the loss-tolerant Deliver path). *store.Store implements it; the hub depends +// a relayed transcript_entry frame to, the durable counterpart to the +// loss-tolerant Deliver path (docs/designs/agent/compass-agent-session-persistence/design.md; +// flush timing in docs/designs/agent/compass-transcript-flush-off-commit.md). +// *store.Store implements it; the hub depends // only on this narrow surface (pattern: DeliveryStore). Wired via // SetTranscriptStore after construction so no NewHub caller signature changes, // and nil-safe: a hub with no transcript store fails a transcript commit closed diff --git a/go/internal/runnerhub/reconstruct.go b/go/internal/runnerhub/reconstruct.go index cb3e2742c..a750e6818 100644 --- a/go/internal/runnerhub/reconstruct.go +++ b/go/internal/runnerhub/reconstruct.go @@ -61,9 +61,11 @@ func (h *Hub) ReconstructSessionBody(ctx context.Context, sessionID string) ([]b } // One atomic snapshot: PG hot-tail and safety_valve manifest are read together - // so a concurrent flush can't corrupt the body. The NotFound-on-empty-tail - // return relies on the store invariant that the valve never evicts the last - // row — a future eviction change MUST preserve it or resume wrongly NotFounds. + // so a background flush pass committing mid-read can't corrupt the body, and + // the manifest read drops valve segments at or below the latest checkpoint. + // The NotFound-on-empty-tail return relies on the store invariant that the + // valve never evicts the last row — a future eviction change MUST preserve it + // or resume wrongly NotFounds. tail, segments, err := reader.SessionResumeSnapshot(ctx, sessionID) if err != nil { return nil, reconstructReadError(err) diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index 022bc095b..64a5f1faf 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -694,17 +694,20 @@ func (h *Hub) RelayCommsCall( // - A hub with no transcript store wired fails CodeUnavailable (the durable // transcript leg is not mounted — a Deliver-only hub). Checked BEFORE // session resolution, so even a bound session gets Unavailable. -// - committed=true on a fresh commit AND on an idempotent replay of an -// already-committed key (the store dedups a duplicate idempotency_key as a -// silent success). A non-commit is NEVER committed=false with a nil error — +// - committed=true once the row commits, on a fresh commit AND on an +// idempotent replay of an already-committed key (the store dedups a +// duplicate idempotency_key as a silent success). The archive flush runs +// after, on the store's flusher, so object-store faults never reach this +// path. A non-commit is NEVER committed=false with a nil error — // it is ALWAYS a Connect status error, because the Runner drives // at-least-once purely off the Connect code (err==nil => committed). // - Store errors are mapped to Connect codes (transcriptCommitError, mirroring // the comms edgeError): ErrInvalidArgument -> InvalidArgument (a malformed or // unknown-session frame, terminal), ErrConflict -> AlreadyExists (a genuine -// entry_seq collision, terminal), any other -> Internal (a transient store -// fault the relay should retry). Never a bare error connect.CodeOf would read -// as CodeUnknown — that would wrongly present as a retryable teardown. +// entry_seq collision, terminal), any other -> Internal (a transient Postgres +// fault on the commit path the relay should retry). Never a bare error +// connect.CodeOf would read as CodeUnknown — that would wrongly present as +// a retryable teardown. // - message_id and seq are not meaningful for a transcript entry (the Runner // reads neither): shipped as "" and 0. func (h *Hub) CommitConversationFrame( diff --git a/go/internal/store/agent_transcripts.go b/go/internal/store/agent_transcripts.go index 97b038cd5..d6609e118 100644 --- a/go/internal/store/agent_transcripts.go +++ b/go/internal/store/agent_transcripts.go @@ -41,10 +41,8 @@ const ( // SegmentKindSuperseded holds pre-checkpoint history flushed at compaction; // NEVER read on resume (the read view starts at the latest checkpoint). SegmentKindSuperseded SegmentKind = "superseded" - // SegmentKindSafetyValve holds post-checkpoint entries evicted by the high - // size cap; spliced back on resume (T5). A later checkpoint re-marks any now - // pre-checkpoint safety_valve row to superseded, so a safety_valve row is - // post-latest-checkpoint by construction. + // SegmentKindSafetyValve holds post-checkpoint entries evicted by the size cap. + // It is resume-eligible only above the latest checkpoint, as the segment readers' query enforces. SegmentKindSafetyValve SegmentKind = "safety_valve" // SegmentKindSessionEnd archives the retained post-checkpoint tail at // teardown for analytics completeness; NEVER read on resume (the PG tail @@ -91,8 +89,8 @@ type ObjectStore interface { // SetObjectStore injects the archive-tier object-store client. It is set once at // construction (slice B wires the real client; store tests inject an in-memory -// fake) before the store serves any flush. A nil object store fails a flush -// loudly rather than silently dropping the archive. +// fake) before the store serves any append. With none, appends request no flush +// and the rows stay in PG; FlushTranscript reports it. func (s *Store) SetObjectStore(os ObjectStore) { s.objectStore = os } @@ -132,10 +130,9 @@ func (s *Store) BindLifetime(ctx context.Context, sessionID string, account Acco // lifetimeSeq is the agent-stamped per-lifetime sequence; the store rebases it // onto the session's base at lifetime bind (entry_seq = base + lifetimeSeq). // Unknown session_id is ErrInvalidArgument (the FK). A distinct entry re-using an -// existing (session_id, entry_seq) is ErrConflict (the PK). When a new -// checkpoint row commits, the PRIMARY flush fires (all pre-checkpoint hot-tail -// rows -> one superseded segment, pruned); otherwise the SAFETY-VALVE cap is -// checked. A deduped retry short-circuits before either flush. +// existing (session_id, entry_seq) is ErrConflict (the PK). It returns once the +// row commits: a fresh row requests a background flush pass, so an object-store +// fault never fails or stalls the commit. A deduped retry requests nothing. func (s *Store) AppendTranscriptEntry(ctx context.Context, sessionID string, lifetimeSeq uint64, checkpoint bool, entryJSON, idempotencyKey string) error { switch { case sessionID == "": @@ -177,18 +174,11 @@ func (s *Store) AppendTranscriptEntry(ctx context.Context, sessionID string, lif } return fmt.Errorf("store: append transcript entry: %w", err) } - if rowsAffected == 0 { - // Duplicate idempotency_key: retry dedup. Silent success, NO flush — a - // retried checkpoint must not re-invoke the primary flush (design.md T4). - // If the original's flush failed, the next checkpoint's flush (or - // session-end) re-covers the range, so the PG read view stays correct. - return nil - } - - if checkpoint { - return s.primaryFlush(ctx, sessionID, entrySeq) + // A deduped retry owes no flush: the original's commit already requested one. + if rowsAffected > 0 && s.objectStore != nil { + s.requestFlush(ctx, sessionID) } - return s.maybeSafetyValve(ctx, sessionID) + return nil } // SessionTranscript returns the hot-tail post-supersession view: the latest @@ -237,8 +227,8 @@ func transcriptRowsFromDB(rows []db.SessionTranscriptRow) []TranscriptEntryRow { // deterministic [min..max] key and the manifest INSERT is // ON CONFLICT (session_id, object_key) DO NOTHING, so a re-run cannot wedge; a // crash before the PUT leaves PG intact. This is the SESSION-END / external flush -// entry point (the call-site is slice C); the PRIMARY and SAFETY-VALVE flushes -// fire internally from AppendTranscriptEntry. +// entry point; the PRIMARY and SAFETY-VALVE flushes run on the Store's +// background flusher after AppendTranscriptEntry commits. func (s *Store) FlushSuperseded(ctx context.Context, sessionID string, uptoEntrySeq uint64, kind SegmentKind) error { if sessionID == "" { return fmt.Errorf("%w: session id is required", ErrInvalidArgument) @@ -253,10 +243,10 @@ func (s *Store) FlushSuperseded(ctx context.Context, sessionID string, uptoEntry // SafetyValveSegments returns the manifest rows for a session whose kind is // safety_valve (post-checkpoint entries evicted to the object store by the size // cap), ordered by min_entry_seq — the only segments T5's reconstructor fetches -// on resume. By construction these are all post-latest-checkpoint: the PRIMARY -// flush re-marks any safety_valve row a later checkpoint superseded to -// superseded, so a stale pre-checkpoint segment is never returned. Empty for a -// normal session (superseded and session_end segments are excluded). +// on resume. The query keeps only segments wholly above the latest checkpoint: +// a valve pass can commit after a newer checkpoint lands, before the next pass +// re-marks its row. Empty for a normal session (superseded and session_end +// segments are excluded). func (s *Store) SafetyValveSegments(ctx context.Context, sessionID string) ([]ArchiveSegmentRow, error) { if sessionID == "" { return nil, fmt.Errorf("%w: session id is required", ErrInvalidArgument) @@ -299,7 +289,8 @@ func segmentRowsFromDB(rows []db.SafetyValveSegmentsRow) []ArchiveSegmentRow { // corrupt the reconstructed body (double-count an entry, or drop one). Under the // snapshot the at-rest disjointness invariant — safety_valve seqs are disjoint // from PG-tail seqs — is observed as a single consistent DB state, closing the -// flush-between-reads race. +// flush-between-reads race. Flushes run in the background, so a deferred or +// failed one only leaves more rows in the PG tail. // // It preserves SessionTranscript's ErrNotFound-on-empty-tail behavior and // short-circuits before reading segments, mirroring the reconstructor's normal @@ -383,83 +374,32 @@ func (s *Store) sessionBase(ctx context.Context, sessionID string) (uint64, erro return toUint64(base), nil } -// primaryFlush is the compaction/checkpoint-arrival flush: when a checkpoint row -// at checkpointSeq commits, it flushes ALL pre-checkpoint hot-tail rows -// (entry_seq < checkpointSeq) as one superseded segment and prunes them, and in -// the SAME transaction re-marks any existing safety_valve manifest row whose -// max_entry_seq < checkpointSeq to superseded — so no crash window leaves a -// superseded safety_valve segment resume-eligible. Self-healing: it targets the -// rows STILL IN the hot-tail, so a crash before this commits is idempotently -// reclaimed by the next compaction. -func (s *Store) primaryFlush(ctx context.Context, sessionID string, checkpointSeq uint64) error { - return s.flushUpto(ctx, sessionID, 0, checkpointSeq-1, SegmentKindSuperseded, true, checkpointSeq) -} - -// maybeSafetyValve is the high-size-cap check on the same AppendTranscriptEntry -// path. Two-phase: a CHEAP single-scalar aggregate sums the post-checkpoint -// hot-tail bytes in ONE query and returns early when under the cap (the common -// path — the valve sits above the compaction window, so it never engages in -// normal operation, and no per-row materialization happens). ONLY when the -// scalar exceeds the cap does it run the per-row query to pick the eviction cut -// point, evict the oldest post-checkpoint chunk as a safety_valve segment, and -// prune it, always retaining the newest entry. +// maybeSafetyValve is the high-size-cap check, run by each flush pass after the +// primary step. A cheap scalar sums the post-checkpoint bytes and returns early +// under the cap (the normal case). Over it, the per-row read picks the cut with +// valveCut, which stops before any checkpoint committed since the scalar and +// always retains the newest entry. func (s *Store) maybeSafetyValve(ctx context.Context, sessionID string) error { - cpSeq, err := s.latestCheckpointSeq(ctx, sessionID) - if err != nil { - return err - } - // Cheap gate: sum the post-checkpoint byte total in ONE scalar (no per-row - // materialization on the common path). The valve sits above the compaction - // window, so in normal operation this scalar is under the cap and returns here. - tailBytes, err := s.q.HotTailBytes(ctx, db.HotTailBytesParams{ - SessionID: sessionID, - EntrySeq: toInt64(cpSeq), - }) + tailBytes, err := s.q.HotTailBytes(ctx, sessionID) if err != nil { return fmt.Errorf("store: measure hot tail: %w", err) } - if tailBytes <= int64(s.safetyValveCapBytes) { + if s.flushes.afterValveRead != nil { + s.flushes.afterValveRead() + } + capBytes := int64(s.safetyValveCapBytes) + if tailBytes <= capBytes { return nil } - // Over the cap: NOW materialize the rows to pick the eviction cut point. - sizes, err := s.q.HotTailSizes(ctx, db.HotTailSizesParams{ - SessionID: sessionID, - EntrySeq: toInt64(cpSeq), - }) + sizes, err := s.q.HotTailSizes(ctx, sessionID) if err != nil { return fmt.Errorf("store: measure hot tail: %w", err) } - type sized struct { - seq uint64 - bytes int64 - } - entries := make([]sized, 0, len(sizes)) - var total int64 - for _, row := range sizes { - entries = append(entries, sized{seq: toUint64(row.EntrySeq), bytes: row.Bytes}) - total += row.Bytes - } - cap64 := int64(s.safetyValveCapBytes) - if total <= cap64 { - return nil - } - - // Evict the oldest post-checkpoint rows until the remainder fits under the - // cap, but never the newest entry (the just-appended one) — the eviction is - // a contiguous oldest chunk with a deterministic upper bound. - remaining := total - var upto uint64 - for i := range len(entries) - 1 { - if remaining <= cap64 { - break - } - upto = entries[i].seq - remaining -= entries[i].bytes - } + from, upto := valveCut(sizes, capBytes) if upto == 0 { return nil } - return s.flushUpto(ctx, sessionID, cpSeq+1, upto, SegmentKindSafetyValve, true, 0) + return s.flushUpto(ctx, sessionID, from, upto, SegmentKindSafetyValve, true, 0) } // latestCheckpointSeq returns the entry_seq of the session's newest checkpoint @@ -510,7 +450,7 @@ func (s *Store) flushUpto(ctx context.Context, sessionID string, fromEntrySeq, u minSeq, maxSeq := seqs[0], seqs[len(seqs)-1] if s.objectStore == nil { - return errors.New("store: object store not configured for transcript flush") + return errNoObjectStore } key := segmentKey(sessionID, minSeq, maxSeq) // PUT-before-prune: the segment must be durable in the object store before @@ -538,13 +478,19 @@ func (s *Store) flushUpto(ctx context.Context, sessionID string, fromEntrySeq, u return fmt.Errorf("store: insert archive segment: %w", err) } if prune { - if err := qtx.PruneTranscriptEntries(ctx, db.PruneTranscriptEntriesParams{ + pruned, err := qtx.PruneTranscriptEntries(ctx, db.PruneTranscriptEntriesParams{ SessionID: sessionID, EntrySeq: toInt64(minSeq), EntrySeq_2: toInt64(maxSeq), - }); err != nil { + }) + if err != nil { return fmt.Errorf("store: prune flushed entries: %w", err) } + // The database owns disjointness: a pass whose range another pass already + // pruned, or a writer grew, must not commit its manifest row. + if pruned != int64(len(seqs)) { + return fmt.Errorf("%w: archived %d rows of [%d..%d], pruned %d", errFlushRaced, len(seqs), minSeq, maxSeq, pruned) + } } if remarkBelow > 0 { if err := qtx.RemarkSafetyValveSuperseded(ctx, db.RemarkSafetyValveSupersededParams{ diff --git a/go/internal/store/agent_transcripts_noobjectstore_test.go b/go/internal/store/agent_transcripts_noobjectstore_test.go index ee2a22f5c..013fe2590 100644 --- a/go/internal/store/agent_transcripts_noobjectstore_test.go +++ b/go/internal/store/agent_transcripts_noobjectstore_test.go @@ -8,9 +8,9 @@ import ( "testing" ) -// With no object store configured, a flush-triggering append still commits its -// row, reports the flush error once, and the same-key retry the relay sends on -// that error dedups to success. Nothing is lost and nothing is pruned. +// With no object store configured, a flush-triggering append commits its row +// and returns nil; FlushTranscript reports the missing tier, and a same-key +// retry dedups to success. Nothing is lost and nothing is pruned. func TestAppendWithoutObjectStoreKeepsEntriesAndRetrySucceeds(t *testing.T) { cases := []struct { name string @@ -57,14 +57,16 @@ func TestAppendWithoutObjectStoreKeepsEntriesAndRetrySucceeds(t *testing.T) { tc.setup(t, s, sess) key := fmt.Sprintf("idem-%d", tc.seq) - err := s.AppendTranscriptEntry(t.Context(), sess, tc.seq, tc.checkpoint, tc.body, key) - if err == nil || !strings.Contains(err.Error(), "object store not configured") { - t.Fatalf("flush-triggering append err = %v, want object store not configured", err) + if err := s.AppendTranscriptEntry(t.Context(), sess, tc.seq, tc.checkpoint, tc.body, key); err != nil { + t.Fatalf("flush-triggering append = %v, want nil (the row commits)", err) + } + if err := s.FlushTranscript(t.Context(), sess); err == nil || !strings.Contains(err.Error(), "object store not configured") { + t.Fatalf("FlushTranscript = %v, want object store not configured", err) } if got := hotTailSeqs(t, s, sess); !seqsEqual(got, tc.wantTail...) { t.Fatalf("hot tail after failed flush = %v, want %v (row committed, nothing pruned)", got, tc.wantTail) } - // The relay retries a CodeInternal under the same key; dedup makes it succeed. + // A same-key retry dedups to success. if err := s.AppendTranscriptEntry(t.Context(), sess, tc.seq, tc.checkpoint, tc.body, key); err != nil { t.Fatalf("same-key retry = %v, want nil (dedup, no flush)", err) } diff --git a/go/internal/store/agent_transcripts_test.go b/go/internal/store/agent_transcripts_test.go index 3d672878b..f438a5e9b 100644 --- a/go/internal/store/agent_transcripts_test.go +++ b/go/internal/store/agent_transcripts_test.go @@ -66,6 +66,12 @@ func (f *fakeObjectStore) GetSegment(_ context.Context, key string) ([]byte, err return body, nil } +func (f *fakeObjectStore) setPutErr(err error) { + f.mu.Lock() + defer f.mu.Unlock() + f.putErr = err +} + func (f *fakeObjectStore) has(key string) bool { f.mu.Lock() defer f.mu.Unlock() @@ -482,10 +488,17 @@ func TestSessionTranscriptUnknownNotFound(t *testing.T) { // ── Flush machinery ───────────────────────────────────────────────────────── -// TestPrimaryFlushOnCheckpointArrival pins the compaction flush embedded in -// AppendTranscriptEntry: when a checkpoint row commits, all pre-checkpoint -// hot-tail rows flush to ONE superseded segment + a manifest row, and the -// hot-tail is pruned to [checkpoint..now]. +// mustFlush waits for a flush pass over the session and fails the test on error. +func mustFlush(t *testing.T, s *Store, sessionID string) { + t.Helper() + if err := s.FlushTranscript(t.Context(), sessionID); err != nil { + t.Fatalf("FlushTranscript(%q): %v", sessionID, err) + } +} + +// TestPrimaryFlushOnCheckpointArrival pins the compaction flush a checkpoint +// commit requests: all pre-checkpoint hot-tail rows flush to ONE superseded +// segment + a manifest row, and the hot-tail is pruned to [checkpoint..now]. func TestPrimaryFlushOnCheckpointArrival(t *testing.T) { s, fake := storeWithFake(t) sess := seedSession(t, s, "primary", "sess-primary") @@ -500,6 +513,7 @@ func TestPrimaryFlushOnCheckpointArrival(t *testing.T) { t.Fatalf("append %d: %v", a.seq, err) } } + mustFlush(t, s, sess) // Hot-tail pruned to [3] (the checkpoint), the pre-checkpoint 1,2 flushed. if got := hotTailSeqs(t, s, sess); !seqsEqual(got, 3) { t.Fatalf("hot tail after primary flush = %v, want [3]", got) @@ -516,6 +530,7 @@ func TestPrimaryFlushOnCheckpointArrival(t *testing.T) { if err := s.AppendTranscriptEntry(t.Context(), sess, 4, false, `{"e":4}`, "idem-4"); err != nil { t.Fatalf("append 4: %v", err) } + mustFlush(t, s, sess) if got := hotTailSeqs(t, s, sess); !seqsEqual(got, 3, 4) { t.Fatalf("hot tail = %v, want [3 4]", got) } @@ -536,10 +551,12 @@ func TestFlushPutBeforePrune(t *testing.T) { } } // Arm the PUT to fail, then land the checkpoint that triggers the flush. - fake.putErr = errors.New("object store down") - err := s.AppendTranscriptEntry(t.Context(), sess, 3, true, `{"e":3}`, "idem-3") - if err == nil { - t.Fatal("checkpoint append with a failing PutSegment must return an error") + fake.setPutErr(errors.New("object store down")) + if err := s.AppendTranscriptEntry(t.Context(), sess, 3, true, `{"e":3}`, "idem-3"); err != nil { + t.Fatalf("checkpoint append = %v, want nil (the flush runs after the commit)", err) + } + if err := s.FlushTranscript(t.Context(), sess); err == nil { + t.Fatal("flush with a failing PutSegment must return an error") } // The checkpoint row committed (the append), but the flush left PG intact: // no prune, no manifest row. @@ -649,6 +666,7 @@ func TestReconstructionIsPGOnlyAfterFlush(t *testing.T) { t.Fatalf("append %d: %v", a.seq, err) } } + mustFlush(t, s, sess) // Any GetSegment from here fails the test. fake.failOnGet = t got, err := s.SessionTranscript(t.Context(), sess) @@ -684,6 +702,7 @@ func TestSafetyValveFlushOnSizeCap(t *testing.T) { t.Fatalf("append %d: %v", seq, err) } } + mustFlush(t, s, sess) seg, err := s.SafetyValveSegments(t.Context(), sess) if err != nil { t.Fatalf("SafetyValveSegments: %v", err) @@ -703,7 +722,7 @@ func TestSafetyValveFlushOnSizeCap(t *testing.T) { // TestSafetyValveResupersededByLaterCheckpoint pins the re-mark: once a later // checkpoint supersedes a safety_valve segment (its max < the new checkpoint -// seq), the PRIMARY flush re-marks it to superseded, so SafetyValveSegments no +// seq), the next pass's primary flush re-marks it to superseded, so SafetyValveSegments no // longer returns it — a stale pre-checkpoint segment is never spliced on resume. func TestSafetyValveResupersededByLaterCheckpoint(t *testing.T) { s, _ := storeWithFake(t) @@ -723,6 +742,7 @@ func TestSafetyValveResupersededByLaterCheckpoint(t *testing.T) { } last = seq } + mustFlush(t, s, sess) seg, err := s.SafetyValveSegments(t.Context(), sess) if err != nil { t.Fatalf("SafetyValveSegments: %v", err) @@ -735,6 +755,7 @@ func TestSafetyValveResupersededByLaterCheckpoint(t *testing.T) { if err := s.AppendTranscriptEntry(t.Context(), sess, last+1, true, `{"cp":2}`, "idem-cp2"); err != nil { t.Fatalf("checkpoint 2: %v", err) } + mustFlush(t, s, sess) after, err := s.SafetyValveSegments(t.Context(), sess) if err != nil { t.Fatalf("SafetyValveSegments after re-mark: %v", err) diff --git a/go/internal/store/db/agent_transcripts.sql.go b/go/internal/store/db/agent_transcripts.sql.go index a7fa0cb6b..1121e22c6 100644 --- a/go/internal/store/db/agent_transcripts.sql.go +++ b/go/internal/store/db/agent_transcripts.sql.go @@ -84,42 +84,41 @@ func (q *Queries) CollectSegment(ctx context.Context, arg CollectSegmentParams) } const hotTailBytes = `-- name: HotTailBytes :one -SELECT COALESCE(SUM(octet_length(entry_json)), 0)::BIGINT AS bytes - FROM agent_session_transcript_entries - WHERE session_id = $1 AND entry_seq > $2 +SELECT COALESCE(SUM(octet_length(e.entry_json)), 0)::BIGINT AS bytes + FROM agent_session_transcript_entries e + WHERE e.session_id = $1 + AND e.entry_seq > (SELECT COALESCE(MAX(cp.entry_seq), 0) + FROM agent_session_transcript_entries cp + WHERE cp.session_id = $1 AND cp.checkpoint) ` -type HotTailBytesParams struct { - SessionID string - EntrySeq int64 -} - -func (q *Queries) HotTailBytes(ctx context.Context, arg HotTailBytesParams) (int64, error) { - row := q.db.QueryRow(ctx, hotTailBytes, arg.SessionID, arg.EntrySeq) +// HotTailBytes and HotTailSizes derive the checkpoint inside the statement, so +// a checkpoint committed between the valve's reads cannot shift its window. +func (q *Queries) HotTailBytes(ctx context.Context, sessionID string) (int64, error) { + row := q.db.QueryRow(ctx, hotTailBytes, sessionID) var bytes int64 err := row.Scan(&bytes) return bytes, err } const hotTailSizes = `-- name: HotTailSizes :many -SELECT entry_seq, octet_length(entry_json)::BIGINT AS bytes - FROM agent_session_transcript_entries - WHERE session_id = $1 AND entry_seq > $2 - ORDER BY entry_seq +SELECT e.entry_seq, e.checkpoint, octet_length(e.entry_json)::BIGINT AS bytes + FROM agent_session_transcript_entries e + WHERE e.session_id = $1 + AND e.entry_seq > (SELECT COALESCE(MAX(cp.entry_seq), 0) + FROM agent_session_transcript_entries cp + WHERE cp.session_id = $1 AND cp.checkpoint) + ORDER BY e.entry_seq ` -type HotTailSizesParams struct { - SessionID string - EntrySeq int64 -} - type HotTailSizesRow struct { - EntrySeq int64 - Bytes int64 + EntrySeq int64 + Checkpoint bool + Bytes int64 } -func (q *Queries) HotTailSizes(ctx context.Context, arg HotTailSizesParams) ([]HotTailSizesRow, error) { - rows, err := q.db.Query(ctx, hotTailSizes, arg.SessionID, arg.EntrySeq) +func (q *Queries) HotTailSizes(ctx context.Context, sessionID string) ([]HotTailSizesRow, error) { + rows, err := q.db.Query(ctx, hotTailSizes, sessionID) if err != nil { return nil, err } @@ -127,7 +126,7 @@ func (q *Queries) HotTailSizes(ctx context.Context, arg HotTailSizesParams) ([]H var items []HotTailSizesRow for rows.Next() { var i HotTailSizesRow - if err := rows.Scan(&i.EntrySeq, &i.Bytes); err != nil { + if err := rows.Scan(&i.EntrySeq, &i.Checkpoint, &i.Bytes); err != nil { return nil, err } items = append(items, i) @@ -206,7 +205,7 @@ func (q *Queries) LatestCheckpointSeq(ctx context.Context, sessionID string) (in return seq, err } -const pruneTranscriptEntries = `-- name: PruneTranscriptEntries :exec +const pruneTranscriptEntries = `-- name: PruneTranscriptEntries :execrows DELETE FROM agent_session_transcript_entries WHERE session_id = $1 AND entry_seq >= $2 AND entry_seq <= $3 ` @@ -217,9 +216,12 @@ type PruneTranscriptEntriesParams struct { EntrySeq_2 int64 } -func (q *Queries) PruneTranscriptEntries(ctx context.Context, arg PruneTranscriptEntriesParams) error { - _, err := q.db.Exec(ctx, pruneTranscriptEntries, arg.SessionID, arg.EntrySeq, arg.EntrySeq_2) - return err +func (q *Queries) PruneTranscriptEntries(ctx context.Context, arg PruneTranscriptEntriesParams) (int64, error) { + result, err := q.db.Exec(ctx, pruneTranscriptEntries, arg.SessionID, arg.EntrySeq, arg.EntrySeq_2) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil } const remarkSafetyValveSuperseded = `-- name: RemarkSafetyValveSuperseded :exec @@ -239,10 +241,15 @@ func (q *Queries) RemarkSafetyValveSuperseded(ctx context.Context, arg RemarkSaf } const safetyValveSegments = `-- name: SafetyValveSegments :many -SELECT object_key, min_entry_seq, max_entry_seq, kind - FROM agent_session_archive_segments - WHERE session_id = $1 AND kind = $2 - ORDER BY min_entry_seq +SELECT a.object_key, a.min_entry_seq, a.max_entry_seq, a.kind + FROM agent_session_archive_segments a + WHERE a.session_id = $1 AND a.kind = $2 + -- A valve segment at or below the latest checkpoint is stale: resume must not + -- splice it after that checkpoint's body. + AND a.min_entry_seq > (SELECT COALESCE(MAX(cp.entry_seq), 0) + FROM agent_session_transcript_entries cp + WHERE cp.session_id = $1 AND cp.checkpoint) + ORDER BY a.min_entry_seq ` type SafetyValveSegmentsParams struct { diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 03ce70e5a..fc541f9ab 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -272,8 +272,10 @@ type Querier interface { // Feeds requireGroupCreateAuthz: owner, agent-owner, same namespace, or SHARED-visibility group. GroupCreateAuthorized(ctx context.Context, arg GroupCreateAuthorizedParams) (bool, error) HasForgeScope(ctx context.Context, arg HasForgeScopeParams) (bool, error) - HotTailBytes(ctx context.Context, arg HotTailBytesParams) (int64, error) - HotTailSizes(ctx context.Context, arg HotTailSizesParams) ([]HotTailSizesRow, error) + // HotTailBytes and HotTailSizes derive the checkpoint inside the statement, so + // a checkpoint committed between the valve's reads cannot shift its window. + HotTailBytes(ctx context.Context, sessionID string) (int64, error) + HotTailSizes(ctx context.Context, sessionID string) ([]HotTailSizesRow, error) InSweepSet(ctx context.Context, arg InSweepSetParams) (bool, error) // Account-domain queries (sqlc adoption T2, RIG-3034). These replace the inline // SQL literals that lived in internal/store/accounts.go; the hand-written Store @@ -528,7 +530,7 @@ type Querier interface { OwnersPeeredRowsForShare(ctx context.Context, arg OwnersPeeredRowsForShareParams) ([]OwnersPeeredRowsForShareRow, error) PinnedEntries(ctx context.Context, channelID string) ([]PinnedEntriesRow, error) PlacementForAgent(ctx context.Context, agentAccountID string) (PlacementForAgentRow, error) - PruneTranscriptEntries(ctx context.Context, arg PruneTranscriptEntriesParams) error + PruneTranscriptEntries(ctx context.Context, arg PruneTranscriptEntriesParams) (int64, error) PullRequestForgeUpdatedAt(ctx context.Context, arg PullRequestForgeUpdatedAtParams) (pgtype.Timestamptz, error) // Explicit links attach directly; a closing reference attaches only when none // of the PR's explicit targets is an issue on the board. diff --git a/go/internal/store/queries/agent_transcripts.sql b/go/internal/store/queries/agent_transcripts.sql index a4d8b990a..2e8f33255 100644 --- a/go/internal/store/queries/agent_transcripts.sql +++ b/go/internal/store/queries/agent_transcripts.sql @@ -41,21 +41,34 @@ SELECT e.entry_seq, e.checkpoint, e.entry_json ORDER BY entry_seq; -- name: SafetyValveSegments :many -SELECT object_key, min_entry_seq, max_entry_seq, kind - FROM agent_session_archive_segments - WHERE session_id = $1 AND kind = $2 - ORDER BY min_entry_seq; +SELECT a.object_key, a.min_entry_seq, a.max_entry_seq, a.kind + FROM agent_session_archive_segments a + WHERE a.session_id = $1 AND a.kind = $2 + -- A valve segment at or below the latest checkpoint is stale: resume must not + -- splice it after that checkpoint's body. + AND a.min_entry_seq > (SELECT COALESCE(MAX(cp.entry_seq), 0) + FROM agent_session_transcript_entries cp + WHERE cp.session_id = $1 AND cp.checkpoint) + ORDER BY a.min_entry_seq; +-- HotTailBytes and HotTailSizes derive the checkpoint inside the statement, so +-- a checkpoint committed between the valve's reads cannot shift its window. -- name: HotTailBytes :one -SELECT COALESCE(SUM(octet_length(entry_json)), 0)::BIGINT AS bytes - FROM agent_session_transcript_entries - WHERE session_id = $1 AND entry_seq > $2; +SELECT COALESCE(SUM(octet_length(e.entry_json)), 0)::BIGINT AS bytes + FROM agent_session_transcript_entries e + WHERE e.session_id = $1 + AND e.entry_seq > (SELECT COALESCE(MAX(cp.entry_seq), 0) + FROM agent_session_transcript_entries cp + WHERE cp.session_id = $1 AND cp.checkpoint); -- name: HotTailSizes :many -SELECT entry_seq, octet_length(entry_json)::BIGINT AS bytes - FROM agent_session_transcript_entries - WHERE session_id = $1 AND entry_seq > $2 - ORDER BY entry_seq; +SELECT e.entry_seq, e.checkpoint, octet_length(e.entry_json)::BIGINT AS bytes + FROM agent_session_transcript_entries e + WHERE e.session_id = $1 + AND e.entry_seq > (SELECT COALESCE(MAX(cp.entry_seq), 0) + FROM agent_session_transcript_entries cp + WHERE cp.session_id = $1 AND cp.checkpoint) + ORDER BY e.entry_seq; -- name: LatestCheckpointSeq :one SELECT COALESCE(MAX(entry_seq), 0)::BIGINT AS seq @@ -79,7 +92,7 @@ INSERT INTO agent_session_archive_segments VALUES ($1, $2, $3, $4, $5) ON CONFLICT (session_id, object_key) DO NOTHING; --- name: PruneTranscriptEntries :exec +-- name: PruneTranscriptEntries :execrows DELETE FROM agent_session_transcript_entries WHERE session_id = $1 AND entry_seq >= $2 AND entry_seq <= $3; diff --git a/go/internal/store/store.go b/go/internal/store/store.go index 59a57dd0b..121b02a82 100644 --- a/go/internal/store/store.go +++ b/go/internal/store/store.go @@ -8,9 +8,14 @@ import ( "errors" "fmt" "io/fs" + "log/slog" "sort" "strconv" "strings" + "time" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/metric" "github.com/RigelBuild/compass/go/internal/store/db" "github.com/jackc/pgx/v5" @@ -33,6 +38,8 @@ var migrationsFS embed.FS // the schema already current. const migrationLockKey int64 = 0x0C0A_5500_0000_0001 +const instrumentationScope = "github.com/RigelBuild/compass/go/internal/store" + // Store is the Postgres store of record. It owns a pgx connection pool and is // safe for concurrent use by every RPC handler; share it by pointer. Construct // it with Open and release it with Close. @@ -65,6 +72,11 @@ type Store struct { // resolved tenant (resolveTenant). Set once in Open before the store serves; // no lock, mirroring coordinationHook's set-once-before-serving discipline. bootstrapTenantID TenantID + // flushes runs transcript flushes after their commits, one pass at a time per + // session; Open starts it and Close joins it. + flushes *transcriptFlusher + // flushPassTimeout bounds one background flush pass; tests lower it. + flushPassTimeout time.Duration } // defaultPoolMaxConns caps the pool when the DSN sets no pool_max_conns. pgx's @@ -127,12 +139,29 @@ func Open(ctx context.Context, dsn string) (*Store, error) { return nil, err } s.bootstrapTenantID = bt + // A metric miss must never fail Open, matching the comms counter. + failures, err := otel.Meter(instrumentationScope).Int64Counter( + "compass.store.transcript_flush_failures", + metric.WithDescription("Count of background transcript flush passes that failed."), + ) + if err != nil { + slog.Warn("store: failed to create transcript flush failure counter; metric disabled", "err", err) + failures = nil + } + s.flushes = newTranscriptFlusher(failures) + s.flushPassTimeout = defaultFlushPassTimeout + s.startTranscriptFlusher(ctx) return s, nil } -// Close releases the connection pool. Idempotent-safe to call once; after it +// Close stops the transcript flusher, waits for every in-flight pass, then +// releases the pool, so no pass touches a closed pool. Call it once; after it // the Store must not be used. func (s *Store) Close() { + if s.flushes != nil && s.flushes.stop != nil { + s.flushes.stop() + <-s.flushes.stopped + } s.pool.Close() } diff --git a/go/internal/store/transcript_flusher.go b/go/internal/store/transcript_flusher.go new file mode 100644 index 000000000..e7eb7ed81 --- /dev/null +++ b/go/internal/store/transcript_flusher.go @@ -0,0 +1,236 @@ +package store + +import ( + "context" + "errors" + "fmt" + "log/slog" + "sync" + "time" + + "go.opentelemetry.io/otel/metric" + + "github.com/RigelBuild/compass/go/internal/store/db" +) + +const ( + // defaultFlushPassTimeout bounds one pass; it moves a cap-sized segment at + // about 9 Mbit/s, so a hung object store cannot hold a slot forever. + defaultFlushPassTimeout = 60 * time.Second + // maxConcurrentFlushes caps passes in flight, bounding the segment bodies + // held in memory while their PUTs run. + maxConcurrentFlushes = 4 +) + +// errFlushRaced marks a pruning flush whose delete count differed from the rows +// it archived: another pass or writer moved the range, so its tx rolled back. +var errFlushRaced = errors.New("store: flush prune count mismatch") + +// errNoObjectStore is the flush failure when no archive tier is wired. +var errNoObjectStore = errors.New("store: object store not configured for transcript flush") + +// errFlusherStopped is FlushTranscript's result once Close has stopped the worker. +var errFlusherStopped = errors.New("store: transcript flusher stopped") + +// flushSlot is one session with a pass owed or running. Guarded by +// transcriptFlusher.mu. +type flushSlot struct { + tenant TenantID + system, running, again bool + err error // last pass result + idle chan struct{} // closed when the slot is deleted +} + +// transcriptFlusher runs transcript flushes off the commit path, one pass at a +// time per session. The Store owns it: Open starts it and Close joins it. +type transcriptFlusher struct { + mu sync.Mutex + slots map[string]*flushSlot + wake, sem chan struct{} // cap 1; cap maxConcurrentFlushes + wg sync.WaitGroup + stop context.CancelFunc + stopped chan struct{} // closed after the loop and passes end + failures metric.Int64Counter // nil when the meter failed + passDone func(sessionID string, err error) // test hook; nil in prod + afterValveRead func() // test hook; nil in prod +} + +func newTranscriptFlusher(failures metric.Int64Counter) *transcriptFlusher { + return &transcriptFlusher{ + slots: map[string]*flushSlot{}, + wake: make(chan struct{}, 1), + sem: make(chan struct{}, maxConcurrentFlushes), + stopped: make(chan struct{}), + failures: failures, + } +} + +// startTranscriptFlusher starts the run loop. Its ctx keeps ctx's values but +// not its cancellation, so only Close stops it. +func (s *Store) startTranscriptFlusher(ctx context.Context) { + runCtx, stop := context.WithCancel(context.WithoutCancel(ctx)) + s.flushes.stop = stop + go s.runTranscriptFlusher(runCtx) +} + +// requestFlush owes sessionID a pass and never blocks. A running slot is marked +// to run once more, so triggers that land mid-pass coalesce into one pass. +func (s *Store) requestFlush(ctx context.Context, sessionID string) *flushSlot { + f := s.flushes + f.mu.Lock() + slot, ok := f.slots[sessionID] + switch { + case !ok: + // The commit ctx carries the session's tenant; the pass reuses its scope + // but not its deadline. + slot = &flushSlot{tenant: s.resolveTenant(ctx), system: isSystemRole(ctx), idle: make(chan struct{})} + f.slots[sessionID] = slot + case slot.running: + slot.again = true + } + f.mu.Unlock() + select { + case f.wake <- struct{}{}: + default: // A wake is already pending; the loop sees this slot on it. + } + return slot +} + +// runTranscriptFlusher starts one goroutine per owed slot until ctx ends, then +// waits for every pass before closing stopped. +func (s *Store) runTranscriptFlusher(ctx context.Context) { + f := s.flushes + defer close(f.stopped) + for { + select { + case <-ctx.Done(): + f.wg.Wait() + return + case <-f.wake: + } + f.mu.Lock() + for sessionID, slot := range f.slots { + if slot.running { + continue + } + slot.running = true + f.wg.Go(func() { s.runSlot(ctx, sessionID, slot) }) + } + f.mu.Unlock() + } +} + +// runSlot runs passes for one session until no trigger is owed, then deletes +// the slot. +func (s *Store) runSlot(ctx context.Context, sessionID string, slot *flushSlot) { + f := s.flushes + for { + err := s.runPass(ctx, sessionID, slot) + if f.passDone != nil { + f.passDone(sessionID, err) + } + f.mu.Lock() + slot.err = err + if slot.again && ctx.Err() == nil { + slot.again = false + f.mu.Unlock() + continue + } + delete(f.slots, sessionID) + close(slot.idle) + f.mu.Unlock() + return + } +} + +// runPass holds a semaphore slot and runs one pass under the slot's tenant +// scope and the pass deadline. A failure is logged, counted, and dropped: the +// next commit's pass covers the same rows. +func (s *Store) runPass(ctx context.Context, sessionID string, slot *flushSlot) error { + f := s.flushes + select { + case f.sem <- struct{}{}: + case <-ctx.Done(): + return fmt.Errorf("store: transcript flush: %w", ctx.Err()) + } + defer func() { <-f.sem }() + + passCtx := WithTenant(ctx, slot.tenant) + if slot.system { + passCtx = WithSystemRole(passCtx) + } + passCtx, cancel := context.WithTimeout(passCtx, s.flushPassTimeout) + defer cancel() + err := s.flushPass(passCtx, sessionID) + if err != nil { + slog.ErrorContext(ctx, "transcript flush failed", "session_id", sessionID, "err", err) + if f.failures != nil { + f.failures.Add(ctx, 1) + } + } + return err +} + +// flushPass works on the rows in PG now, not on its trigger: the primary flush +// below the latest checkpoint, then the safety valve. +func (s *Store) flushPass(ctx context.Context, sessionID string) error { + cp, err := s.latestCheckpointSeq(ctx, sessionID) + if err != nil { + return err + } + if cp > 0 { + // Archive every pre-checkpoint row as one superseded segment, and re-mark + // any safety_valve segment the checkpoint now supersedes, in one tx. + if err := s.flushUpto(ctx, sessionID, 0, cp-1, SegmentKindSuperseded, true, cp); err != nil { + return err + } + } + return s.maybeSafetyValve(ctx, sessionID) +} + +// FlushTranscript runs a pass for sessionID and waits for its slot to go idle, +// returning the last pass's result. It is the synchronous barrier for tests. +func (s *Store) FlushTranscript(ctx context.Context, sessionID string) error { + if s.objectStore == nil { + return errNoObjectStore + } + slot := s.requestFlush(ctx, sessionID) + select { + case <-slot.idle: + s.flushes.mu.Lock() + defer s.flushes.mu.Unlock() + return slot.err + case <-ctx.Done(): + return fmt.Errorf("store: wait for transcript flush: %w", ctx.Err()) + case <-s.flushes.stopped: + return errFlusherStopped + } +} + +// valveCut picks the safety valve's eviction range from the post-checkpoint +// rows: the oldest rows until the rest fit capBytes, never the newest row, and +// never a checkpoint row or anything after it. upto == 0 means no cut. +func valveCut(rows []db.HotTailSizesRow, capBytes int64) (from, upto uint64) { + if len(rows) == 0 { + return 0, 0 + } + var total int64 + for _, r := range rows { + total += r.Bytes + } + if total <= capBytes { + return 0, 0 + } + remaining := total + for i := range len(rows) - 1 { + if remaining <= capBytes || rows[i].Checkpoint { + break + } + upto = toUint64(rows[i].EntrySeq) + remaining -= rows[i].Bytes + } + if upto == 0 { + return 0, 0 + } + return toUint64(rows[0].EntrySeq), upto +} diff --git a/go/internal/store/transcript_flusher_test.go b/go/internal/store/transcript_flusher_test.go new file mode 100644 index 000000000..0e0c80910 --- /dev/null +++ b/go/internal/store/transcript_flusher_test.go @@ -0,0 +1,573 @@ +//go:build pgtest + +package store + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "testing" + "time" +) + +// blockingObjectStore holds every PUT until release closes or the pass ctx +// ends, and records how many PUTs ran at once. +type blockingObjectStore struct { + mu sync.Mutex + objects map[string][]byte + inFlight int + maxInFlight int + started chan string + release chan struct{} +} + +func newBlockingObjectStore() *blockingObjectStore { + return &blockingObjectStore{ + objects: map[string][]byte{}, + started: make(chan string, 16), + release: make(chan struct{}), + } +} + +func (b *blockingObjectStore) PutSegment(ctx context.Context, key string, body []byte) error { + b.mu.Lock() + b.inFlight++ + b.maxInFlight = max(b.maxInFlight, b.inFlight) + b.mu.Unlock() + defer func() { + b.mu.Lock() + b.inFlight-- + b.mu.Unlock() + }() + b.started <- key + select { + case <-b.release: + case <-ctx.Done(): + return fmt.Errorf("blocking fake: %w", ctx.Err()) + } + b.mu.Lock() + defer b.mu.Unlock() + b.objects[key] = append([]byte(nil), body...) + return nil +} + +func (b *blockingObjectStore) GetSegment(_ context.Context, key string) ([]byte, error) { + b.mu.Lock() + defer b.mu.Unlock() + body, ok := b.objects[key] + if !ok { + return nil, fmt.Errorf("blocking fake: no object at %q", key) + } + return body, nil +} + +func (b *blockingObjectStore) peak() int { + b.mu.Lock() + defer b.mu.Unlock() + return b.maxInFlight +} + +// recv waits for one value on ch, failing the test when ctx ends first. +func recv[T any](t *testing.T, ctx context.Context, ch <-chan T, what string) T { + t.Helper() + select { + case v := <-ch: + return v + case <-ctx.Done(): + t.Fatalf("waiting for %s: %v", what, ctx.Err()) + var zero T + return zero + } +} + +// awaitIdle waits for the session's current slot, if any, and returns its last +// pass result. +func awaitIdle(t *testing.T, s *Store, sessionID string) error { + t.Helper() + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + s.flushes.mu.Lock() + slot, ok := s.flushes.slots[sessionID] + s.flushes.mu.Unlock() + if !ok { + return nil + } + recv(t, ctx, slot.idle, "the slot to go idle") + s.flushes.mu.Lock() + defer s.flushes.mu.Unlock() + return slot.err +} + +// boundSession seeds a session and binds its lifetime, so appends land at seq. +func boundSession(t *testing.T, s *Store, handle string) string { + t.Helper() + sess := seedSession(t, s, handle, "sess-"+handle) + if _, err := s.BindLifetime(t.Context(), sess, sessionOwner(t, s, sess)); err != nil { + t.Fatalf("BindLifetime: %v", err) + } + return sess +} + +// pad40 is a 40-byte entry body, sized against the 40- and 80-byte test caps. +func pad40(c string) string { + return `{"pad":"` + strings.Repeat(c, 30) + `"}` +} + +// A commit returns once its row commits, while the flush's PUT is still blocked. +func TestAppendReturnsWhileFlushPutBlocked(t *testing.T) { + cases := []struct { + name string + cap int + seed []bool // checkpoint flag per seeded seq, from 1 + trigger bool // checkpoint flag of the triggering append + body string + wantKind SegmentKind + wantSeg [2]uint64 + wantTail []uint64 + }{ + { + name: "checkpoint", seed: []bool{false, false}, trigger: true, body: `{"cp":3}`, + wantKind: SegmentKindSuperseded, wantSeg: [2]uint64{1, 2}, wantTail: []uint64{3}, + }, + { + name: "safety valve", cap: 40, seed: []bool{true, false}, trigger: false, body: pad40("v"), + wantKind: SegmentKindSafetyValve, wantSeg: [2]uint64{2, 2}, wantTail: []uint64{1, 3}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + s := newTestStore(t) + if tc.cap > 0 { + s.SetSafetyValveCapBytes(tc.cap) + } + sess := boundSession(t, s, "nb-"+strings.ReplaceAll(tc.name, " ", "")) + // Seed with no object store, so the seed requests no pass. + for i, cp := range tc.seed { + appendOK(t, s, sess, uint64(i+1), cp, pad40("s")) + } + fake := newBlockingObjectStore() + s.SetObjectStore(fake) + + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + seq := uint64(len(tc.seed) + 1) + if err := s.AppendTranscriptEntry(ctx, sess, seq, tc.trigger, tc.body, fmt.Sprintf("idem-%d", seq)); err != nil { + t.Fatalf("append = %v, want nil while the PUT is blocked", err) + } + recv(t, ctx, fake.started, "the flush PUT") + if got := hotTailSeqs(t, s, sess); !seqsEqual(got, 1, 2, 3) { + t.Fatalf("hot tail while the PUT is blocked = %v, want [1 2 3]", got) + } + + close(fake.release) + if err := s.FlushTranscript(ctx, sess); err != nil { + t.Fatalf("FlushTranscript: %v", err) + } + man := manifestRows(t, s, sess) + if len(man) != 1 || man[0].Kind != tc.wantKind || man[0].MinEntrySeq != tc.wantSeg[0] || man[0].MaxEntrySeq != tc.wantSeg[1] { + t.Fatalf("manifest = %+v, want one %s %v", man, tc.wantKind, tc.wantSeg) + } + if got := hotTailSeqs(t, s, sess); !seqsEqual(got, tc.wantTail...) { + t.Fatalf("hot tail after flush = %v, want %v", got, tc.wantTail) + } + }) + } +} + +// Triggers that land during a running pass coalesce into one more pass, and one +// session's PUTs never overlap. +func TestFlushTriggersCoalesce(t *testing.T) { + s := newTestStore(t) + sess := boundSession(t, s, "coalesce") + appendOK(t, s, sess, 1, false, `{"e":1}`) + appendOK(t, s, sess, 2, false, `{"e":2}`) + fake := newBlockingObjectStore() + s.SetObjectStore(fake) + var ( + mu sync.Mutex + passes int + ) + s.flushes.passDone = func(string, error) { + mu.Lock() + passes++ + mu.Unlock() + } + + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + appendOK(t, s, sess, 3, true, `{"cp":3}`) + recv(t, ctx, fake.started, "pass 1's PUT") + s.flushes.mu.Lock() + slot := s.flushes.slots[sess] + s.flushes.mu.Unlock() + appendOK(t, s, sess, 4, false, `{"e":4}`) + appendOK(t, s, sess, 5, true, `{"cp":5}`) + + close(fake.release) + recv(t, ctx, slot.idle, "the slot to go idle") + mu.Lock() + got := passes + mu.Unlock() + if got != 2 { + t.Fatalf("passes = %d, want 2 (the two mid-pass triggers coalesce)", got) + } + if peak := fake.peak(); peak != 1 { + t.Fatalf("concurrent PUTs = %d, want 1", peak) + } + if tail := hotTailSeqs(t, s, sess); !seqsEqual(tail, 5) { + t.Fatalf("hot tail = %v, want [5]", tail) + } + man := manifestRows(t, s, sess) + if len(man) != 2 || man[0].MinEntrySeq != 1 || man[0].MaxEntrySeq != 2 || man[1].MinEntrySeq != 3 || man[1].MaxEntrySeq != 4 { + t.Fatalf("manifest = %+v, want superseded [1..2] and [3..4]", man) + } +} + +// A checkpoint that commits between the valve's reads is never evicted: resume +// heads with it, and no segment covers it. +func TestSafetyValveKeepsCheckpointCommittedMidPass(t *testing.T) { + s := newTestStore(t) + s.SetSafetyValveCapBytes(40) + sess := boundSession(t, s, "cpcut") + appendOK(t, s, sess, 1, true, `{"cp":1}`) + appendOK(t, s, sess, 2, false, pad40("a")) + appendOK(t, s, sess, 3, false, pad40("b")) + s.SetObjectStore(newFakeObjectStore()) + var once sync.Once + s.flushes.afterValveRead = func() { + once.Do(func() { + // A 40-byte checkpoint lands inside the range a stale cut would evict. + for _, a := range []struct { + seq uint64 + cp bool + body string + }{{4, true, pad40("c")}, {5, false, `{"e":5}`}} { + if err := s.AppendTranscriptEntry(t.Context(), sess, a.seq, a.cp, a.body, fmt.Sprintf("idem-%d", a.seq)); err != nil { + t.Errorf("mid-pass append %d: %v", a.seq, err) + } + } + }) + } + + mustFlush(t, s, sess) + tail := hotTailSeqs(t, s, sess) + if !seqsEqual(tail, 4, 5) { + t.Fatalf("hot tail = %v, want [4 5] (the mid-pass checkpoint kept)", tail) + } + for _, m := range manifestRows(t, s, sess) { + if m.MinEntrySeq <= 4 && m.MaxEntrySeq >= 4 { + t.Fatalf("segment %+v covers the mid-pass checkpoint 4", m) + } + } + snap, _, err := s.SessionResumeSnapshot(t.Context(), sess) + if err != nil { + t.Fatalf("SessionResumeSnapshot: %v", err) + } + if len(snap) == 0 || snap[0].EntrySeq != 4 || !snap[0].Checkpoint { + t.Fatalf("resume tail = %+v, want it headed by checkpoint 4", snap) + } +} + +// Two Stores whose valve ranges overlap under different keys: the database lets +// exactly one prune commit, and the other pass rolls back with errFlushRaced. +func TestOverlappingFlushesPruneOnce(t *testing.T) { + a, dsn := newTestStoreDSN(t) + b := reopenStore(t, dsn) + a.SetSafetyValveCapBytes(40) + b.SetSafetyValveCapBytes(80) + sess := boundSession(t, a, "raced") + appendOK(t, a, sess, 1, true, `{"cp":1}`) + appendOK(t, a, sess, 2, false, pad40("a")) + appendOK(t, a, sess, 3, false, pad40("b")) + fake := newBlockingObjectStore() + a.SetObjectStore(fake) + b.SetObjectStore(fake) + passA := make(chan error, 4) + a.flushes.passDone = func(_ string, err error) { passA <- err } + + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + // A (cap 40) cuts [2..3]; B (cap 80) cuts [2..2]. + appendOK(t, a, sess, 4, false, pad40("c")) + errB := make(chan error, 1) + go func() { errB <- b.FlushTranscript(ctx, sess) }() + recv(t, ctx, fake.started, "the first PUT") + recv(t, ctx, fake.started, "the second PUT") + close(fake.release) + gotA := recv(t, ctx, passA, "A's pass") + gotB := recv(t, ctx, errB, "B's flush") + + raced := 0 + for _, err := range []error{gotA, gotB} { + switch { + case errors.Is(err, errFlushRaced): + raced++ + case err != nil: + t.Fatalf("pass error = %v, want nil or errFlushRaced", err) + } + } + if raced != 1 { + t.Fatalf("passes reporting errFlushRaced = %d (A=%v, B=%v), want 1", raced, gotA, gotB) + } + man := manifestRows(t, a, sess) + if len(man) != 1 { + t.Fatalf("manifest = %+v, want exactly one row", man) + } + var want []uint64 + for seq := uint64(1); seq <= 4; seq++ { + if seq < man[0].MinEntrySeq || seq > man[0].MaxEntrySeq { + want = append(want, seq) + } + } + if got := hotTailSeqs(t, a, sess); !seqsEqual(got, want...) { + t.Fatalf("hot tail = %v, want %v (only the winning segment's rows pruned)", got, want) + } +} + +// A pass that outlives its deadline fails with DeadlineExceeded and leaves PG +// untouched. +func TestFlushPassDeadline(t *testing.T) { + s := newTestStore(t) + s.flushPassTimeout = 50 * time.Millisecond + sess := boundSession(t, s, "deadline") + appendOK(t, s, sess, 1, false, `{"e":1}`) + appendOK(t, s, sess, 2, false, `{"e":2}`) + appendOK(t, s, sess, 3, true, `{"cp":3}`) + s.SetObjectStore(newBlockingObjectStore()) // never released + + err := s.FlushTranscript(t.Context(), sess) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("FlushTranscript = %v, want context.DeadlineExceeded", err) + } + if got := hotTailSeqs(t, s, sess); !seqsEqual(got, 1, 2, 3) { + t.Fatalf("hot tail = %v, want [1 2 3]", got) + } + if man := manifestRows(t, s, sess); len(man) != 0 { + t.Fatalf("manifest = %+v, want empty", man) + } +} + +// shutdownObjectStore reports the pass ctx's cancellation, then holds the PUT +// until release so the test controls when the pass ends. +type shutdownObjectStore struct { + started, cancelled, release chan struct{} +} + +func (o *shutdownObjectStore) PutSegment(ctx context.Context, _ string, _ []byte) error { + close(o.started) + <-ctx.Done() + close(o.cancelled) + <-o.release + return fmt.Errorf("shutdown fake: %w", ctx.Err()) +} + +func (o *shutdownObjectStore) GetSegment(context.Context, string) ([]byte, error) { + return nil, errors.New("shutdown fake: no reads") +} + +// Close cancels the in-flight pass and waits for it before closing the pool. +func TestCloseJoinsInFlightFlush(t *testing.T) { + s, dsn := newTestStoreDSN(t) + sess := boundSession(t, s, "shutdown") + appendOK(t, s, sess, 1, false, `{"e":1}`) + appendOK(t, s, sess, 2, false, `{"e":2}`) + fake := &shutdownObjectStore{started: make(chan struct{}), cancelled: make(chan struct{}), release: make(chan struct{})} + s.SetObjectStore(fake) + closeDone := make(chan struct{}) + passDone := make(chan struct{}) + s.flushes.passDone = func(string, error) { + select { + case <-closeDone: + t.Errorf("pass ended after Close returned") + default: + } + close(passDone) + } + + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + appendOK(t, s, sess, 3, true, `{"cp":3}`) + recv(t, ctx, fake.started, "the flush PUT") + go func() { + s.Close() + close(closeDone) + }() + recv(t, ctx, fake.cancelled, "the pass ctx to cancel") + select { + case <-closeDone: + t.Fatal("Close returned while a pass was in flight") + default: + } + close(fake.release) + recv(t, ctx, closeDone, "Close") + select { + case <-passDone: + default: + t.Fatal("Close returned before the pass ended") + } + + other := reopenStore(t, dsn) + if got := hotTailSeqs(t, other, sess); !seqsEqual(got, 1, 2, 3) { + t.Fatalf("hot tail after Close = %v, want [1 2 3]", got) + } + if man := manifestRows(t, other, sess); len(man) != 0 { + t.Fatalf("manifest after Close = %+v, want empty", man) + } +} + +// A pass runs under the committing request's tenant, not the bootstrap tenant. +func TestFlushRunsUnderAppendTenant(t *testing.T) { + s, fake := storeWithFake(t) + tenantB := seedTenant(t, s, "flush-tenant-b") + ctxB := WithTenant(t.Context(), tenantB) + agent := seedTenantAgent(t, ctxB, s, "flush-b") + const sess = "sess-flush-b" + if err := s.RecordAgentSession(ctxB, sess, agent.ID); err != nil { + t.Fatalf("RecordAgentSession: %v", err) + } + for _, a := range []struct { + seq uint64 + cp bool + }{{1, false}, {2, false}, {3, true}} { + if err := s.AppendTranscriptEntry(ctxB, sess, a.seq, a.cp, fmt.Sprintf(`{"e":%d}`, a.seq), fmt.Sprintf("idem-b-%d", a.seq)); err != nil { + t.Fatalf("append %d: %v", a.seq, err) + } + } + // Wait on the append's own slot: a new FlushTranscript slot would take the + // waiter's tenant instead. + if err := awaitIdle(t, s, sess); err != nil { + t.Fatalf("flush pass: %v", err) + } + man := manifestRows(t, s, sess) + if len(man) != 1 || !fake.has(man[0].ObjectKey) { + t.Fatalf("manifest = %+v, want one stored segment", man) + } + var tenant string + if err := s.pool.QueryRow(t.Context(), + `SELECT tenant_id FROM agent_session_archive_segments WHERE session_id = $1`, sess).Scan(&tenant); err != nil { + t.Fatalf("read segment tenant: %v", err) + } + if tenant != string(tenantB) { + t.Fatalf("segment tenant = %q, want %q", tenant, tenantB) + } +} + +// A safety_valve row at or below the latest checkpoint is stale; both segment +// readers drop it and keep the one above. +func TestSafetyValveSegmentsHideRowsBelowCheckpoint(t *testing.T) { + s := newTestStore(t) + sess := boundSession(t, s, "svfilter") + for seq, cp := range []bool{true, false, false, true, false} { + appendOK(t, s, sess, uint64(seq+1), cp, fmt.Sprintf(`{"e":%d}`, seq+1)) + } + for _, r := range [][2]int64{{2, 3}, {5, 5}} { + if _, err := s.pool.Exec(t.Context(), + `INSERT INTO agent_session_archive_segments (session_id, object_key, min_entry_seq, max_entry_seq, kind, tenant_id) + VALUES ($1, $2, $3, $4, 'safety_valve', $5)`, + sess, segmentKey(sess, toUint64(r[0]), toUint64(r[1])), r[0], r[1], string(s.bootstrapTenantID)); err != nil { + t.Fatalf("seed manifest row %v: %v", r, err) + } + } + pool, err := s.SafetyValveSegments(t.Context(), sess) + if err != nil { + t.Fatalf("SafetyValveSegments: %v", err) + } + _, snap, err := s.SessionResumeSnapshot(t.Context(), sess) + if err != nil { + t.Fatalf("SessionResumeSnapshot: %v", err) + } + for name, got := range map[string][]ArchiveSegmentRow{"SafetyValveSegments": pool, "SessionResumeSnapshot": snap} { + if len(got) != 1 || got[0].MinEntrySeq != 5 || got[0].MaxEntrySeq != 5 { + t.Errorf("%s = %+v, want only the [5..5] segment above checkpoint 4", name, got) + } + } +} + +// gatedObjectStore counts PUTs in flight, signals each entry, and holds each PUT +// until one release token arrives or the pass ctx ends. +type gatedObjectStore struct { + mu sync.Mutex + active int + peak int + entered chan string + release chan struct{} +} + +func (g *gatedObjectStore) PutSegment(ctx context.Context, key string, _ []byte) error { + g.mu.Lock() + g.active++ + g.peak = max(g.peak, g.active) + g.mu.Unlock() + defer func() { + g.mu.Lock() + g.active-- + g.mu.Unlock() + }() + g.entered <- key + select { + case <-g.release: + return nil + case <-ctx.Done(): + return fmt.Errorf("gated fake: %w", ctx.Err()) + } +} + +func (g *gatedObjectStore) GetSegment(context.Context, string) ([]byte, error) { + return nil, errors.New("gated fake: no reads") +} + +func (g *gatedObjectStore) counts() (active, peak int) { + g.mu.Lock() + defer g.mu.Unlock() + return g.active, g.peak +} + +// Passes across sessions share one cap: a fifth session's pass waits until one +// of maxConcurrentFlushes running passes ends. +func TestFlushPassesShareGlobalCap(t *testing.T) { + s := newTestStore(t) + const sessions = maxConcurrentFlushes + 1 + ids := make([]string, sessions) + for i := range ids { + ids[i] = boundSession(t, s, fmt.Sprintf("cap%d", i)) + // Idempotency keys are global, so each session needs its own. + for seq, cp := range []bool{false, false, true} { + if err := s.AppendTranscriptEntry(t.Context(), ids[i], uint64(seq+1), cp, `{"e":1}`, fmt.Sprintf("cap%d-%d", i, seq)); err != nil { + t.Fatalf("append %s/%d: %v", ids[i], seq+1, err) + } + } + } + fake := &gatedObjectStore{entered: make(chan string, sessions), release: make(chan struct{})} + s.SetObjectStore(fake) + + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + for _, id := range ids[:maxConcurrentFlushes] { + s.requestFlush(ctx, id) + } + for range maxConcurrentFlushes { + recv(t, ctx, fake.entered, "a capped PUT") + } + + // The fifth pass is owed but must not reach its PUT while four hold the cap. + short, cancelShort := context.WithTimeout(ctx, 200*time.Millisecond) + defer cancelShort() + if err := s.FlushTranscript(short, ids[maxConcurrentFlushes]); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("FlushTranscript on the fifth session = %v, want context.DeadlineExceeded", err) + } + if active, peak := fake.counts(); active != maxConcurrentFlushes || peak != maxConcurrentFlushes { + t.Fatalf("PUTs active=%d peak=%d, want %d (the cap)", active, peak, maxConcurrentFlushes) + } + + fake.release <- struct{}{} + recv(t, ctx, fake.entered, "the queued fifth PUT") + close(fake.release) + for _, id := range ids { + mustFlush(t, s, id) + } + if _, peak := fake.counts(); peak != maxConcurrentFlushes { + t.Fatalf("peak PUTs = %d, want %d", peak, maxConcurrentFlushes) + } +} diff --git a/go/internal/store/valve_cut_test.go b/go/internal/store/valve_cut_test.go new file mode 100644 index 000000000..65ac9c043 --- /dev/null +++ b/go/internal/store/valve_cut_test.go @@ -0,0 +1,33 @@ +package store + +import ( + "testing" + + "github.com/RigelBuild/compass/go/internal/store/db" +) + +func TestValveCut(t *testing.T) { + row := func(seq int64, checkpoint bool, bytes int64) db.HotTailSizesRow { + return db.HotTailSizesRow{EntrySeq: seq, Checkpoint: checkpoint, Bytes: bytes} + } + cases := []struct { + name string + rows []db.HotTailSizesRow + capBytes int64 + wantFrom, want uint64 + }{ + {name: "under cap", rows: []db.HotTailSizesRow{row(2, false, 10), row(3, false, 10)}, capBytes: 40}, + {name: "evicts oldest until it fits", rows: []db.HotTailSizesRow{row(2, false, 40), row(3, false, 40), row(4, false, 40)}, capBytes: 40, wantFrom: 2, want: 3}, + {name: "keeps the newest row", rows: []db.HotTailSizesRow{row(2, false, 40), row(3, false, 100)}, capBytes: 40, wantFrom: 2, want: 2}, + {name: "stops before a mid-tail checkpoint", rows: []db.HotTailSizesRow{row(2, false, 40), row(3, false, 40), row(4, true, 40), row(5, false, 40)}, capBytes: 40, wantFrom: 2, want: 3}, + {name: "checkpoint first means no cut", rows: []db.HotTailSizesRow{row(2, true, 40), row(3, false, 40)}, capBytes: 40}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + from, upto := valveCut(tc.rows, tc.capBytes) + if from != tc.wantFrom || upto != tc.want { + t.Fatalf("valveCut = [%d..%d], want [%d..%d]", from, upto, tc.wantFrom, tc.want) + } + }) + } +} diff --git a/go/server/service_resume_pgtest_test.go b/go/server/service_resume_pgtest_test.go index 61be7fbaf..efba865ee 100644 --- a/go/server/service_resume_pgtest_test.go +++ b/go/server/service_resume_pgtest_test.go @@ -366,6 +366,10 @@ func TestStartAgentSessionResumeS3FallbackReconstructsEvictedEntries(t *testing. } } + // The valve runs on the store's background flusher; wait for its pass. + if err := f.store.FlushTranscript(ctx, logical); err != nil { + t.Fatalf("FlushTranscript: %v", err) + } // The valve must have fired: a safety_valve segment recorded and a body in // the object store. segs, err := f.store.SafetyValveSegments(ctx, logical)