diff --git a/go/internal/runner/gateway/gateway.go b/go/internal/runner/gateway/gateway.go index 1d45be4c9..04ed8ade7 100644 --- a/go/internal/runner/gateway/gateway.go +++ b/go/internal/runner/gateway/gateway.go @@ -174,13 +174,12 @@ type Gateway struct { // across the whole event stream and the hub only flags seq > lastSeq+1, so // replayed low seqs are accepted and loss in that range stops being detectable. - // Scope is per-Runner-link: relay.go's eventPublisher owns a SECOND counter and - // both feed the hub's one high-water mark, so gap detection is meaningful only - // while exactly one is live. This path replaced the stdout relay for gateway - // traffic, so that holds today; unifying the two is T9. + // In production every Gateway of one Runner link shares deps.Seq, so RunnerSeq + // stays monotonic across container sockets. The hub's stale-ERRORED guard + // relies on that: a resumed session's new socket must not restart at 1. pubMu sync.Mutex pub *sessionPublisher - seq seqCounter + seq *SeqCounter // publishMu fences telemetry admission while a lifecycle terminal state is sent. publishMu sync.Mutex @@ -222,6 +221,8 @@ type Deps struct { Events EventRelay // Committer forwards a durable conversation frame to the Server for commit (CommitConversationFrame). Committer ConversationCommitter + // Seq is the Runner link's RunnerSeq allocator; nil gives this Gateway its own. + Seq *SeqCounter } // NewGateway builds the AgentGateway handler for the container's socket: @@ -240,6 +241,10 @@ type Deps struct { // the stream outlives any one agent request. A caller with no distinct socket // scope (a hermetic test) passes context.Background(). func NewGateway(baseCtx context.Context, containerName string, deps Deps) *Gateway { + seq := deps.Seq + if seq == nil { + seq = &SeqCounter{} + } return &Gateway{ baseCtx: baseCtx, containerName: containerName, @@ -250,6 +255,7 @@ func NewGateway(baseCtx context.Context, containerName string, deps Deps) *Gatew board: deps.Board, events: deps.Events, committer: deps.Committer, + seq: seq, control: noopControlRouter{}, // ttl=0: no expiry, a pure size-bounded LRU (committedKeysMax). The cache // is advisory, so eviction is safe — it never drops the store's boundary. diff --git a/go/internal/runner/gateway/publisher.go b/go/internal/runner/gateway/publisher.go index de94adf12..9a00c9dcb 100644 --- a/go/internal/runner/gateway/publisher.go +++ b/go/internal/runner/gateway/publisher.go @@ -11,10 +11,9 @@ package gateway // emission order, else the hub's gap detector records a false gap. So a publisher // holds its stream mutex across BOTH the seq allocation AND the Send. -// Scope: the counter is Gateway-scoped (per socket / per Runner link) and survives -// a publisher replacement, since a per-publisher counter would restart the -// sequence on a swap. relay.go's eventPublisher owns a SECOND counter feeding the -// same high-water mark, so gap detection holds only while one is live; unifying is T9. +// Scope: the counter is Runner-link-wide (shared by every container's Gateway) and +// survives a publisher replacement, since a per-publisher counter would restart the +// sequence on a swap. import ( "context" @@ -32,9 +31,9 @@ type EventRelay interface { PublishEvents(ctx context.Context) *connect.ClientStreamForClient[compassv1internal.PublishEventsRequest, compassv1internal.PublishEventsResponse] } -// seqCounter is the RunnerSeq allocator shared by every publisher a Gateway -// builds for its socket. It carries ONLY the counter and the lock that guards -// the counter — deliberately not the publishers' stream lock. +// SeqCounter is the RunnerSeq allocator shared by every publisher of every +// Gateway on one Runner link. It carries ONLY the counter and the lock that +// guards the counter — deliberately not the publishers' stream lock. // // Sharing the counter is required: a publisher is replaceable within one session, // and a per-publisher counter restarts the sequence on that swap. Sharing the @@ -42,17 +41,16 @@ type EventRelay interface { // a replacement and then closes the stale publisher outside pubMu, so a // CloseAndReceive round-trip against an unresponsive-but-connected Server would // block every forward on the live replacement's separate upstream stream. Two -// distinct streams need no mutual ordering — the hub keeps one global high-water -// mark and cannot observe an interleaving between them — so the coupling would -// buy nothing and cost unbounded liveness. -type seqCounter struct { +// distinct streams may therefore deliver out of seq order; the hub tolerates a +// late lower seq, so the coupling would buy nothing and cost unbounded liveness. +type SeqCounter struct { mu sync.Mutex n uint64 } // next allocates and returns the next sequence. Called with the publisher's own // stream lock held, so allocation order still equals emission order. -func (c *seqCounter) next() uint64 { +func (c *SeqCounter) next() uint64 { c.mu.Lock() defer c.mu.Unlock() c.n++ @@ -64,9 +62,9 @@ func (c *seqCounter) next() uint64 { // permanent hole, and the hub flags a skipped number as in-transit loss // (runnerhub/hub.go:230). A durable frame erring back to the agent is correct, // expected behaviour — it must not make the Server report a loss that did not -// happen. Safe because the caller holds its stream lock across allocate-and-send, -// so no other goroutine can have sent past this value. -func (c *seqCounter) rollback(seq uint64) { +// happen. Only the latest seq is reclaimed: if another Gateway allocated since, +// the number stays burned and shows as a gap once the hub has a baseline. +func (c *SeqCounter) rollback(seq uint64) { c.mu.Lock() defer c.mu.Unlock() if c.n == seq { @@ -99,7 +97,7 @@ type sessionPublisher struct { // seq is the Gateway's shared allocator, carried across publishers so the // sequence survives a replacement. Only the counter is shared; its lock is // held just long enough to allocate. - seq *seqCounter + seq *SeqCounter // admit, when set, allocates under the Gateway's publish gate so a sealed // session refuses frames instead of sequencing them after its ERRORED report. admit func() (uint64, error) @@ -122,7 +120,7 @@ type sessionPublisher struct { // life: the Publish handler uses the socket-lifetime context, while lifecycle // reports use a bounded caller ctx. seq is the Gateway counter, carried across // publishers so the sequence never restarts. -func newSessionPublisher(ctx context.Context, relay EventRelay, sessionID string, seq *seqCounter) *sessionPublisher { +func newSessionPublisher(ctx context.Context, relay EventRelay, sessionID string, seq *SeqCounter) *sessionPublisher { ctx, cancel := context.WithCancel(ctx) return &sessionPublisher{ sessionID: sessionID, @@ -230,7 +228,7 @@ func (g *Gateway) acquirePublisher(sessionID string) *sessionPublisher { g.pub = nil } if g.pub == nil { - g.pub = newSessionPublisher(g.baseCtx, g.events, sessionID, &g.seq) + g.pub = newSessionPublisher(g.baseCtx, g.events, sessionID, g.seq) g.pub.admit = func() (uint64, error) { return g.admitFrame(sessionID) } g.pub.afterAdmit = g.afterAdmit } diff --git a/go/internal/runner/gateway/session_state.go b/go/internal/runner/gateway/session_state.go index 7febb5853..c1e4eaeba 100644 --- a/go/internal/runner/gateway/session_state.go +++ b/go/internal/runner/gateway/session_state.go @@ -52,7 +52,7 @@ func (l *SocketListener) PublishSessionState(ctx context.Context, sessionID stri } sendCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), stateSendTimeout) defer cancel() - pub := newSessionPublisher(sendCtx, g.events, sessionID, &g.seq) + pub := newSessionPublisher(sendCtx, g.events, sessionID, g.seq) forwardErr := pub.forward(frame) if forwardErr == nil { sharedErr = nil diff --git a/go/internal/runner/gateway/telemetry_ingest_test.go b/go/internal/runner/gateway/telemetry_ingest_test.go index ba957f7f3..da34fa0d5 100644 --- a/go/internal/runner/gateway/telemetry_ingest_test.go +++ b/go/internal/runner/gateway/telemetry_ingest_test.go @@ -1169,7 +1169,7 @@ func (s *seqSink) seqs() []uint64 { // a variant defect that restarts at a value which happens not to collide would // pass a duplicates-only check. [1 2] pins the contract. // -// RED: give newSessionPublisher its own &seqCounter{} instead of the Gateway's +// RED: give newSessionPublisher its own &SeqCounter{} instead of the Gateway's // -> seqs = [1 1], and this fails every run. func TestSequenceSurvivesPublisherReplacement(t *testing.T) { sink := &seqSink{} @@ -1204,6 +1204,28 @@ func TestSequenceSurvivesPublisherReplacement(t *testing.T) { } } +// A resumed session gets a new container socket, so a per-Gateway counter would +// restart at 1 and fall under the hub's stale-ERRORED boundary. +// +// RED: ignore deps.Seq in NewGateway -> seqs = [1 1]. +func TestSequenceSharedAcrossGateways(t *testing.T) { + sink := &seqSink{} + events := newRunnerServiceServer(t, sink) + shared := &SeqCounter{} + for _, name := range []string{"cont-1", "cont-2"} { + g := NewGateway(context.Background(), name, Deps{Sessions: boundSessions(), Events: events, Seq: shared}) + if err := g.acquirePublisher("sess-1").forward(traceFrame(name)); err != nil { + t.Fatalf("forward on %s = %v, want success", name, err) + } + if err := releaseCurrentPublisher(g); err != nil { + t.Fatalf("release on %s = %v", name, err) + } + } + if got, want := sink.seqs(), []uint64{1, 2}; !slices.Equal(got, want) { + t.Fatalf("RunnerSeq sequence = %v, want %v across Gateways sharing one counter", got, want) + } +} + // A failed forward must not burn a sequence number. The counter is // socket-lifetime, so an allocated-but-unsent number is a permanent hole, and // the hub reads a skipped number as in-transit loss (runnerhub/hub.go:230) — so diff --git a/go/internal/runner/host.go b/go/internal/runner/host.go index 6fb95e744..b226c7425 100644 --- a/go/internal/runner/host.go +++ b/go/internal/runner/host.go @@ -97,6 +97,9 @@ type agentHost struct { mu sync.Mutex sessions map[string]*liveSession sockets map[string]*gateway.SocketListener + // runnerSeq is shared by every container Gateway so RunnerSeq never restarts + // within one enrollment; the hub drops lifecycle frames at or below ERRORED's. + runnerSeq gateway.SeqCounter // afterExitCheck is a test seam between detecting exit and acquiring the // container lock; its returned func runs when retireOnExit returns. Nil in production. afterExitCheck func() func() @@ -871,7 +874,7 @@ func (h *agentHost) provisionVsockGateway(ctx context.Context, spec runtime.Agen h.teardownContainer(ctx, name) return "", fmt.Errorf("resolving vsock gateway endpoint for container %q: backend reports no session", name) } - deps := gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client} + deps := gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client, Seq: &h.runnerSeq} h.log.InfoContext(ctx, "serving agent gateway over vsock path", slog.String("container", name), slog.String("path", endpoint)) listener, err := gateway.Serve(ctx, endpoint, name, deps) @@ -1250,7 +1253,7 @@ func (h *agentHost) serveSocket(ctx context.Context, containerName string) (*gat // serveSocketAt is serveSocket with an explicit socket path: container tiers pass the // fixed RuntimeDir socket, the host tier a path in the handle's state dir (no mount). func (h *agentHost) serveSocketAt(ctx context.Context, containerName, path string) (*gateway.SocketListener, error) { - listener, err := gateway.Serve(ctx, path, containerName, gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client}) + listener, err := gateway.Serve(ctx, path, containerName, gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client, Seq: &h.runnerSeq}) if err != nil { return nil, fmt.Errorf("serving agent socket for container %q: %w", containerName, err) } diff --git a/go/internal/runnerhub/binding_cache_test.go b/go/internal/runnerhub/binding_cache_test.go index 4c1e17683..c1cac768e 100644 --- a/go/internal/runnerhub/binding_cache_test.go +++ b/go/internal/runnerhub/binding_cache_test.go @@ -54,6 +54,8 @@ type fakeBindingStore struct { mu sync.Mutex tenant store.TenantID bindings map[string]store.SessionBinding // session id -> binding + versions map[string]string // session id -> row version, new per write + written int resolveCalled bool resolveCtxSystemRole bool @@ -66,14 +68,14 @@ type fakeBindingStore struct { } func newFakeBindingStore() *fakeBindingStore { - return &fakeBindingStore{tenant: "tenant-a", bindings: map[string]store.SessionBinding{}} + return &fakeBindingStore{tenant: "tenant-a", bindings: map[string]store.SessionBinding{}, versions: map[string]string{}} } -func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, error) { +func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, string, error) { f.mu.Lock() defer f.mu.Unlock() if f.recordErr != nil { - return "", f.recordErr + return "", "", f.recordErr } // session_bindings_session_key: this session id already belongs to a // DIFFERENT account. The real store raises ErrConflict rather than @@ -81,7 +83,7 @@ func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID str // single-valued; a fake that silently overwrote would hide the whole // class (a Runner restart re-mints "sess-1", so id reuse is routine). if b, ok := f.bindings[sessionID]; ok && b.AccountID != accountID { - return "", fmt.Errorf("%w: session %q is already bound to a different agent", store.ErrConflict, sessionID) + return "", "", fmt.Errorf("%w: session %q is already bound to a different agent", store.ErrConflict, sessionID) } var displaced string // Account-keyed UPSERT: a prior binding for this account is displaced. @@ -89,49 +91,61 @@ func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID str if b.AccountID == accountID && sid != sessionID { displaced = sid delete(f.bindings, sid) + delete(f.versions, sid) break } } - f.bindings[sessionID] = store.SessionBinding{TenantID: f.tenant, SessionID: sessionID, AccountID: accountID, RunnerID: runnerID} - return displaced, nil + return displaced, f.putLocked(sessionID, accountID, runnerID), nil } -func (f *fakeBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) { +func (f *fakeBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, string, error) { f.mu.Lock() defer f.mu.Unlock() f.resolveCalled = true f.resolveCtxSystemRole = store.IsSystemRole(ctx) if f.resolveErr != nil { - return "", "", f.resolveErr + return "", "", "", f.resolveErr } b, ok := f.bindings[sessionID] if !ok { - return "", "", store.ErrNotFound + return "", "", "", store.ErrNotFound } - return b.AccountID, b.RunnerID, nil + return b.AccountID, b.RunnerID, f.versions[sessionID], nil } -func (f *fakeBindingStore) SessionForAccount(_ context.Context, accountID store.AccountID) (string, string, error) { +func (f *fakeBindingStore) SessionForAccount(_ context.Context, accountID store.AccountID) (string, string, string, error) { f.mu.Lock() defer f.mu.Unlock() if f.reverseErr != nil { - return "", "", f.reverseErr + return "", "", "", f.reverseErr } for sid, b := range f.bindings { if b.AccountID == accountID { - return sid, b.RunnerID, nil + return sid, b.RunnerID, f.versions[sid], nil } } - return "", "", store.ErrNotFound + return "", "", "", store.ErrNotFound } func (f *fakeBindingStore) DeleteSessionBinding(_ context.Context, sessionID string) error { f.mu.Lock() defer f.mu.Unlock() delete(f.bindings, sessionID) + delete(f.versions, sessionID) return nil } +func (f *fakeBindingStore) DeleteSessionBindingVersion(_ context.Context, sessionID, version string) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + if _, ok := f.bindings[sessionID]; !ok || f.versions[sessionID] != version { + return false, nil + } + delete(f.bindings, sessionID) + delete(f.versions, sessionID) + return true, nil +} + func (f *fakeBindingStore) DeleteSessionBindingsForRunner(_ context.Context, runnerID string) ([]store.SessionBinding, error) { f.mu.Lock() defer f.mu.Unlock() @@ -143,6 +157,7 @@ func (f *fakeBindingStore) DeleteSessionBindingsForRunner(_ context.Context, run if b.RunnerID == runnerID { removed = append(removed, b) delete(f.bindings, sid) + delete(f.versions, sid) } } return removed, nil @@ -173,6 +188,15 @@ func (f *fakeBindingStore) SessionBindingTenant(_ context.Context, sessionID, ru return "", store.ErrNotFound } +// putLocked writes a binding row with a fresh version, as every upsert does. +func (f *fakeBindingStore) putLocked(sessionID string, accountID store.AccountID, runnerID string) string { + f.written++ + version := fmt.Sprintf("v%d", f.written) + f.bindings[sessionID] = store.SessionBinding{TenantID: f.tenant, SessionID: sessionID, AccountID: accountID, RunnerID: runnerID} + f.versions[sessionID] = version + return version +} + func (f *fakeBindingStore) seed(sessionID string) { f.seedBinding(sessionID, testAgentAccount, testRunnerID) } @@ -180,7 +204,7 @@ func (f *fakeBindingStore) seed(sessionID string) { func (f *fakeBindingStore) seedBinding(sessionID string, accountID store.AccountID, runnerID string) { f.mu.Lock() defer f.mu.Unlock() - f.bindings[sessionID] = store.SessionBinding{TenantID: f.tenant, SessionID: sessionID, AccountID: accountID, RunnerID: runnerID} + f.putLocked(sessionID, accountID, runnerID) } // fakeRoutingFabric is an in-memory RoutingFabric double: PublishBindingChange @@ -843,18 +867,18 @@ func (b *reapAndReadBlockingBindingStore) DeleteSessionBindingsForRunner(ctx con return b.fakeBindingStore.DeleteSessionBindingsForRunner(ctx, runnerID) } -func (b *reapAndReadBlockingBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) { - account, runnerID, err := b.fakeBindingStore.ResolveSessionBinding(ctx, sessionID) +func (b *reapAndReadBlockingBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, string, error) { + account, runnerID, version, err := b.fakeBindingStore.ResolveSessionBinding(ctx, sessionID) close(b.readEntered) <-b.readRelease - return account, runnerID, err + return account, runnerID, version, err } -func (b *reapAndReadBlockingBindingStore) SessionForAccount(ctx context.Context, accountID store.AccountID) (string, string, error) { - sessionID, runnerID, err := b.fakeBindingStore.SessionForAccount(ctx, accountID) +func (b *reapAndReadBlockingBindingStore) SessionForAccount(ctx context.Context, accountID store.AccountID) (string, string, string, error) { + sessionID, runnerID, version, err := b.fakeBindingStore.SessionForAccount(ctx, accountID) close(b.readEntered) <-b.readRelease - return sessionID, runnerID, err + return sessionID, runnerID, version, err } // TestConcurrentResolveDuringAFaultingReapCannotResurrect fences the window @@ -951,7 +975,7 @@ func TestOlderEnrollReapCannotDeleteNewerPromotion(t *testing.T) { t.Fatal("promotion wrote its binding while the older reap's delete was in flight") } // Ordered after the reap, the live session's binding survives in the store and cache. - if sessionID, runnerID, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); err != nil || sessionID != "sess-new" || runnerID != testRunnerID { + if sessionID, runnerID, _, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); err != nil || sessionID != "sess-new" || runnerID != testRunnerID { t.Fatalf("durable SessionForAccount = (%q, %q, %v), want (sess-new, %s, nil)", sessionID, runnerID, err, testRunnerID) } if sessionID, ok := hub.CachedSessionForAccount(testAgentAccount); !ok || sessionID != "sess-new" { @@ -959,6 +983,48 @@ func TestOlderEnrollReapCannotDeleteNewerPromotion(t *testing.T) { } } +// A Stop's unbind racing an in-flight reap waits for the whole enroll, so the reap +// still finds the session and drives its OFFLINE edge. +func TestStopUnbindWaitsForEnrollReap(t *testing.T) { + hub := newHubOnly() + pres := &fakePresenceSink{} + hub.SetPresenceSink(pres) + plain := newFakeBindingStore() + hub.SetSessionBindingStore(plain) + ctx := store.WithTenant(context.Background(), "tenant-a") + hub.enroll(ctx, testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, testRunnerID) + hub.promoteSession(ctx, "cont-1", "sess-1") + + bindings := &blockingBindingStore{fakeBindingStore: plain, entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + reapDone := make(chan struct{}) + go func() { + defer close(reapDone) + hub.enroll(ctx, testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + }() + <-bindings.entered + unbindDone := make(chan struct{}) + go func() { + defer close(unbindDone) + hub.unbindSession(ctx, "sess-1") + }() + waitParkedOrDone(t, "(*Hub).releaseSession", unbindDone) + close(bindings.release) + <-reapDone + <-unbindDone + + offline := 0 + for _, rec := range pres.lifecycleSnapshot() { + if rec.sessionID == "sess-1" && rec.state == compassv1.AgentSessionState_AGENT_SESSION_STATE_DISCONNECTED { + offline++ + } + } + if offline == 0 { + t.Fatal("neither the reap nor the Stop drove sess-1 OFFLINE") + } +} + // overlapDetectingBindingStore parks the first reap and records whether a binding // write arrived while it was parked. type overlapDetectingBindingStore struct { @@ -985,7 +1051,7 @@ func (b *overlapDetectingBindingStore) DeleteSessionBindingsForRunner(ctx contex return b.fakeBindingStore.DeleteSessionBindingsForRunner(ctx, runnerID) } -func (b *overlapDetectingBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, error) { +func (b *overlapDetectingBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, string, error) { b.mu.Lock() if b.parked { b.overlap = true @@ -1069,7 +1135,7 @@ func TestEnrollDuringPromotionCannotResurrectBinding(t *testing.T) { if forwardCached || reverseCached { t.Fatalf("promotion repopulated cache after enroll: forward=%v reverse=%v", forwardCached, reverseCached) } - if sessionID, _, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); !errors.Is(err, store.ErrNotFound) { + if sessionID, _, _, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); !errors.Is(err, store.ErrNotFound) { t.Fatalf("durable SessionForAccount = (%q, %v), want ErrNotFound after enroll", sessionID, err) } } @@ -1134,7 +1200,7 @@ type blockingRecordBindingStore struct { recordRelease chan struct{} } -func (b *blockingRecordBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, error) { +func (b *blockingRecordBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, string, error) { close(b.recordEntered) <-b.recordRelease return b.fakeBindingStore.RecordSessionBinding(ctx, sessionID, accountID, runnerID) diff --git a/go/internal/runnerhub/handler.go b/go/internal/runnerhub/handler.go index 4df7d08e4..fcf5672d5 100644 --- a/go/internal/runnerhub/handler.go +++ b/go/internal/runnerhub/handler.go @@ -121,7 +121,7 @@ func (h *Handler) Sessions(ctx context.Context, stream *connect.BidiStream[compa if !ok { return errUnauthenticated } - router, _, err := h.hub.routerFor(subj.ID) + router, enrollGen, err := h.hub.routerForStream(subj.ID) if err != nil { // A Sessions stream with no enrolled Runner — the Runner must Enroll // before opening Sessions. @@ -129,7 +129,7 @@ func (h *Handler) Sessions(ctx context.Context, stream *connect.BidiStream[compa } router.setSessionUnknown(func(sessionID string) { - h.hub.dropLostSessionDetached(ctx, subj.ID, sessionID, false) + h.hub.dropLostSessionDetached(ctx, enrollGen, subj.ID, sessionID, nil, false) }) router.attach(stream.Send) defer router.detach(errStreamClosed) @@ -169,6 +169,7 @@ func (h *Handler) PublishEvents(ctx context.Context, stream *connect.ClientStrea if !ok { return nil, errUnauthenticated } + enrollGen := h.hub.EnrollGeneration() for stream.Receive() { msg := stream.Msg() if err := h.hub.Deliver(ctx, RunnerEvent{ @@ -176,6 +177,7 @@ func (h *Handler) PublishEvents(ctx context.Context, stream *connect.ClientStrea SessionID: msg.GetSessionId(), RunnerID: subj.ID, Frame: msg.GetFrame(), + EnrollGen: enrollGen, }); err != nil { return nil, connect.NewError(connect.CodeInternal, err) } diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 02c00f19d..01ca98d8c 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -21,6 +21,8 @@ import ( "sync" "sync/atomic" + lru "github.com/hashicorp/golang-lru/v2" + compassv1 "github.com/RigelBuild/compass/go/gen/compass/v1" "github.com/RigelBuild/compass/go/internal/fabric" compassv1internal "github.com/RigelBuild/compass/go/internal/gen/compass/v1" @@ -43,6 +45,9 @@ type RunnerEvent struct { RunnerID string // Frame is the relayed agent stdout frame, verbatim. Frame *compassv1internal.AgentFrame + // EnrollGen is the hub's enrollment generation when the frame's stream opened; + // a session frame from an older one is dropped. Zero means no enrollment fence. + EnrollGen uint64 } // LifecycleSink publishes an extracted agent-session lifecycle transition onto @@ -208,19 +213,23 @@ type SessionBindingStore interface { // returns the session id it DISPLACED (empty when the account held none) — // the prior session the hub must evict from both maps. It runs on the // request ctx (tenant-scoped), so the write lands under the acting tenant. - RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (displaced string, err error) + // version names the written row; a release conditioned on it cannot remove a re-bind. + RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (displaced, version string, err error) // ResolveSessionBinding resolves the agent account and owning Runner a live // session speaks for — the cache-miss read behind accountForSession. // store.ErrNotFound is the fail-closed miss (mapped to ok=false), - // byte-identical to today's CodeNotFound. - ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) + // byte-identical to today's CodeNotFound. version names the row read. + ResolveSessionBinding(ctx context.Context, sessionID string) (account store.AccountID, runnerID, version string, err error) // SessionForAccount resolves the live session bound to an account — the // cache-miss read behind SessionForAccount (the reverse direction). Same // fail-closed store.ErrNotFound contract. Returns the owning Runner id too. - SessionForAccount(ctx context.Context, accountID store.AccountID) (sessionID, runnerID string, err error) + SessionForAccount(ctx context.Context, accountID store.AccountID) (sessionID, runnerID, version string, err error) // DeleteSessionBinding releases one session's binding — the unbind write. // Idempotent: releasing an already-released session is a no-op success. DeleteSessionBinding(ctx context.Context, sessionID string) error + // DeleteSessionBindingVersion releases sessionID only while its row is still + // the write that returned version; removed reports whether it was. + DeleteSessionBindingVersion(ctx context.Context, sessionID, version string) (removed bool, err error) // DeleteSessionBindingsForRunner is the enroll sweep: it releases every // binding attached to runnerID and RETURNS the rows it removed, driving the // enroll reap (offline edges + held-deliver reap) from durable truth @@ -406,6 +415,8 @@ type Hub struct { reapStale map[string]uint64 // bindingEpoch fences read-throughs across reap start and completion. bindingEpoch uint64 + // lastLifetime numbers sessionAccounts inserts (sessionBinding.lifetime). mu-guarded. + lastLifetime uint64 // runnerEpoch entries are bounded by the number of distinct Runner ids. runnerEpoch map[string]uint64 // lifecycleCaller is the spawn/despawn execution seam RelayLifecycleCall delegates @@ -429,7 +440,14 @@ type Hub struct { // bindingWriteMu serializes whole enrolls (map-clear through reap) with promotion // writes and cache updates. Lock it before mu; never hold mu across a store call. bindingWriteMu sync.Mutex - mu sync.Mutex + // lifecycleMu guards lifecycleSeqs and sessionLocks. Held only for map access; + // a session's ordering is held by its own sessionLocks entry. + lifecycleMu sync.Mutex + // sessionLocks serialize one session's lifecycle frames from seq record through + // publication, so a slow binding lookup stalls only that session. Refcounted and + // removed when idle. lifecycleMu-guarded. + sessionLocks map[string]*sessionLock + mu sync.Mutex // runner is the single attached Runner (single-Runner MVP, OQ6). A second // enrollment re-attaches rather than registering a second entry. @@ -444,6 +462,17 @@ type Hub struct { // Stop removes, a Runner reconnect drops ALL, so a re-minted id fails closed // (CodeNotFound) not inheriting a stale account (OQ-2). sessionAccounts map[string]sessionBinding + // lifecycleSeqs maps a session to the highest RunnerSeq of an accepted lifecycle + // frame. One lifetime's frames arrive in seq order, so a lower one is from a + // cancelled stream or a dead lifetime. An LRU, so a boundary outlives unbind but + // not maxLifecycleSeqs newer sessions; purged on re-enroll. lifecycleMu-guarded. + lifecycleSeqs *lru.Cache[string, uint64] + // enrollMu fences session-frame delivery (read) against enroll (write), and + // guards enrollGen, which counts enrollments so a stream opened before a + // re-enroll cannot act on its sessions. Lock order: enrollMu, bindingWriteMu, a + // session lock, lifecycleMu, mu. + enrollMu sync.RWMutex + enrollGen uint64 // accountSessions is the REVERSE of sessionAccounts (account -> live session_id), // maintained wherever sessionAccounts is so the two never drift. The delivery // consumer (RIG-1569 T3) resolves a subscribed account to its live session to @@ -451,9 +480,13 @@ type Hub struct { accountSessions map[store.AccountID]string // lastSeq is the highest RunnerSeq Deliver has accepted, for gap detection. lastSeq uint64 - // seenGap records whether a sequence gap was ever observed (in-transit - // loss), surfaced for the board/diagnostics. - seenGap bool + // missingSeqs holds skipped RunnerSeqs not yet seen. Container Gateways share + // one counter but send on separate streams, so a skipped seq may arrive late; + // a gap is loss only while it stays open. Capped at maxMissingSeqs. + missingSeqs map[uint64]struct{} + // gapOverflow is set once more seqs went missing than the cap can track; it + // stays set for the enrollment, since untracked seqs can never be proven seen. + gapOverflow bool // unknownFrames counts frames whose oneof variant was unset or unrecognized // — logged and counted, never silently dropped (agent.proto:38-39). unknownFrames uint64 @@ -479,9 +512,14 @@ type attachedRunner struct { } // sessionBinding is one live session's principal and the Runner that owns it. +// version and lifetime name one binding of the session id, so a release from an +// older lifetime can skip a re-bind: version is the durable write ("" when none +// was made), lifetime is the cache entry, unique per insert and never zero. type sessionBinding struct { account store.AccountID runnerID string + version string + lifetime uint64 } // NewHub constructs a hub over the two write-through sinks and the agent-comms @@ -503,6 +541,10 @@ func NewHub(lifecycle LifecycleSink, tail SessionTailSink, comms CommsCaller, lo containerAccounts: make(map[string]sessionBinding), sessionAccounts: make(map[string]sessionBinding), accountSessions: make(map[store.AccountID]string), + lifecycleSeqs: newLifecycleSeqs(), + sessionLocks: make(map[string]*sessionLock), + enrollGen: 1, + missingSeqs: make(map[uint64]struct{}), reapStale: make(map[string]uint64), runnerEpoch: make(map[string]uint64), } @@ -676,7 +718,20 @@ func (h *Hub) SetBoardCaller(c BoardCaller) { // whose variant is unset or unrecognized is logged and counted, never silently // dropped (design.md:1427-1434, agent.proto:38-39). func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { - h.recordSeq(ev.RunnerSeq) + // Held for the whole delivery so enroll cannot reset sequence or lifecycle state + // between the generation check and their use. Never re-acquired below: a nested + // RLock deadlocks behind a pending enroll. + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + stale := h.staleEnrollmentLocked(ev) + // A lifecycle frame holds its session's lock from before its seq is recorded, + // so a higher seq the hub has already seen is never overtaken. + if !stale && sessionState(ev.Frame) != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED { + defer h.lockSession(ev.SessionID)() + } + if !stale { + h.recordSeq(ev.RunnerSeq) + } frame := ev.Frame oneof := frame.GetFrame() @@ -688,7 +743,10 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { } switch f := oneof.(type) { case *compassv1internal.AgentFrame_Session: - h.deliverSession(ctx, ev.RunnerID, ev.SessionID, f.Session) + // Acks stay unfenced: they are idempotent and record work already done. + if !stale { + h.deliverSession(ctx, ev, f.Session) + } return nil case *compassv1internal.AgentFrame_DeliveryAck: h.deliverAck(ctx, ev, f.DeliveryAck) @@ -705,12 +763,20 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { } } -// SeenGap reports whether Deliver ever observed a Runner-sequence gap. For the +// EnrollGeneration returns the current enrollment generation, for a stream to +// stamp on its RunnerEvents. +func (h *Hub) EnrollGeneration() uint64 { + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + return h.enrollGen +} + +// SeenGap reports whether a skipped Runner sequence is still unseen. For the // board/diagnostics; a gap means in-transit loss the Client bus resync recovers. func (h *Hub) SeenGap() bool { h.mu.Lock() defer h.mu.Unlock() - return h.seenGap + return h.seenGapLocked() } // UnknownFrames reports the count of unset/unrecognized frames Deliver has seen. @@ -736,7 +802,7 @@ func (h *Hub) DroppedAcks() uint64 { // one would interleave with Deliver and could report a gap without the drop that // accompanied it. type FrameDiagnostics struct { - // SeenGap is true once a Runner-sequence gap was observed (in-transit loss + // SeenGap is true while a skipped Runner sequence is unseen (in-transit loss // the Client bus resync recovers). SeenGap bool // UnknownFrames counts frames whose oneof variant was unset or unrecognized. @@ -755,7 +821,7 @@ func (h *Hub) FrameDiagnostics() FrameDiagnostics { h.mu.Lock() defer h.mu.Unlock() return FrameDiagnostics{ - SeenGap: h.seenGap, + SeenGap: h.seenGapLocked(), UnknownFrames: h.unknownFrames, DroppedAcks: h.droppedAcks, } @@ -790,7 +856,8 @@ func (h *Hub) fireRunnerReady() { // deliverSession routes session frames to the observation pane, publishes lifecycle // transitions, and retires an owned session when its Runner reports ERRORED. // UNSPECIFIED means "trace only, no transition". -func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, sf *compassv1internal.SessionFrame) { +func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1internal.SessionFrame) { + runnerID, sessionID, seq := ev.RunnerID, ev.SessionID, ev.RunnerSeq state := sf.GetState() lifecycle := state != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED // A frame the publishing Runner may not speak for is dropped whole: its trace, @@ -801,10 +868,24 @@ func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, sf slog.String("runner_id", runnerID), slog.String("session_id", sessionID)) return } - h.tail.RelaySessionFrame(sessionID, sf) if !lifecycle { + h.tail.RelaySessionFrame(sessionID, sf) + return + } + // Deliver holds the session's lock for a lifecycle frame. + h.lifecycleMu.Lock() + last, seen := h.lifecycleSeqs.Get(sessionID) + if !seen || seq > last { + h.lifecycleSeqs.Add(sessionID, seq) + } + h.lifecycleMu.Unlock() + if seen && seq <= last { + h.log.Debug("ignored lifecycle frame older than the session's latest", + slog.String("runner_id", runnerID), slog.String("session_id", sessionID), + slog.String("state", state.String()), slog.Uint64("runner_seq", seq), slog.Uint64("latest_seq", last)) return } + h.tail.RelaySessionFrame(sessionID, sf) // Resolve the session's agent account and stamp it onto the published status — the // DL-167 attribution join. A status published after a Runner reconnect cleared the // maps carries none (the residual gap). runnerRuntimeIdentity reads tier/posture in @@ -838,7 +919,7 @@ func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, sf } // The Runner can see the exit before any deliver is refused, so ERRORED is a loss too. if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED && hasAccount { - h.dropLostSessionDetached(ctx, runnerID, sessionID, true) + h.dropLostSessionDetached(ctx, h.enrollGen, runnerID, sessionID, &binding, true) } } @@ -940,20 +1021,91 @@ func (h *Hub) forgeNotificationAck(ctx context.Context, ev RunnerEvent, ack *com } } -// recordSeq advances the accepted-sequence high-water mark and flags a gap when -// the observed seq is not exactly one past the last (in-transit loss). The first -// event (lastSeq == 0) establishes the baseline without flagging. +// maxLifecycleSeqs bounds lifecycleSeqs. A buffered stale frame lands within +// seconds, long before this many newer sessions could evict its boundary. +const maxLifecycleSeqs = 4096 + +func newLifecycleSeqs() *lru.Cache[string, uint64] { + c, err := lru.New[string, uint64](maxLifecycleSeqs) + if err != nil { + panic(err) // only a non-positive size errors + } + return c +} + +// sessionLock is one session's lifecycle mutex; refs counts holders and waiters. +type sessionLock struct { + mu sync.Mutex + refs int +} + +// lockSession locks sessionID's lifecycle mutex and returns its unlock. +func (h *Hub) lockSession(sessionID string) func() { + h.lifecycleMu.Lock() + l := h.sessionLocks[sessionID] + if l == nil { + l = &sessionLock{} + h.sessionLocks[sessionID] = l + } + l.refs++ + h.lifecycleMu.Unlock() + l.mu.Lock() + return func() { + l.mu.Unlock() + h.lifecycleMu.Lock() + if l.refs--; l.refs == 0 { + delete(h.sessionLocks, sessionID) + } + h.lifecycleMu.Unlock() + } +} + +// sessionState is a session frame's lifecycle state, UNSPECIFIED for any other frame. +func sessionState(frame *compassv1internal.AgentFrame) compassv1.AgentSessionState { + return frame.GetSession().GetState() +} + +// staleEnrollmentLocked reports, and logs, an event whose stream predates the +// current enrollment. Zero EnrollGen means unfenced. Caller holds enrollMu. +func (h *Hub) staleEnrollmentLocked(ev RunnerEvent) bool { + if ev.EnrollGen == 0 || ev.EnrollGen == h.enrollGen { + return false + } + h.log.Debug("dropped event from a stream opened before re-enroll", + slog.String("runner_id", ev.RunnerID), slog.String("session_id", ev.SessionID)) + return true +} + +// maxMissingSeqs bounds missingSeqs against a corrupt or hostile seq jump. +const maxMissingSeqs = 4096 + +// recordSeq advances the accepted-sequence high-water mark and records the seqs +// a jump skipped; a late arrival fills its gap. The first event (lastSeq == 0) +// establishes the baseline without flagging. func (h *Hub) recordSeq(seq uint64) { h.mu.Lock() defer h.mu.Unlock() + if seq <= h.lastSeq { + delete(h.missingSeqs, seq) + return + } if h.lastSeq != 0 && seq > h.lastSeq+1 { - h.seenGap = true h.log.Warn("runner event sequence gap", slog.Uint64("expected", h.lastSeq+1), slog.Uint64("got", seq)) + for missing := h.lastSeq + 1; missing < seq; missing++ { + if len(h.missingSeqs) >= maxMissingSeqs { + h.gapOverflow = true + break + } + h.missingSeqs[missing] = struct{}{} + } } - if seq > h.lastSeq { - h.lastSeq = seq - } + h.lastSeq = seq +} + +// seenGapLocked reports whether any skipped seq is still unseen. Caller holds mu. +func (h *Hub) seenGapLocked() bool { + return h.gapOverflow || len(h.missingSeqs) > 0 } // countUnknown records and logs an unknown frame. @@ -1001,6 +1153,13 @@ type promotedPair struct { // so none of its pre-enroll sessions live. A failed durable reap still runs the // in-RAM fallback, then returns the error so the Runner retries enrollment. func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier compassv1.RuntimeTier, egressPosture compassv1.EgressPosture) (reattached bool, err error) { + // Held until the runner and router are replaced, so a stream that snapshots + // both under it sees one enrollment's pair. + h.enrollMu.Lock() + h.lifecycleMu.Lock() + h.lifecycleSeqs.Purge() + h.enrollGen++ + h.lifecycleMu.Unlock() // Held from the map-clear through the reap, so no promotion lands in between. h.bindingWriteMu.Lock() h.mu.Lock() @@ -1030,6 +1189,10 @@ func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier clear(h.containerAccounts) clear(h.sessionAccounts) clear(h.accountSessions) + // A new Runner process restarts RunnerSeq, so gap tracking restarts with it. + h.lastSeq = 0 + clear(h.missingSeqs) + h.gapOverflow = false // Refuse read-through for this Runner from the instant the maps are cleared: // a concurrent lookup could otherwise resurrect rows while the reap is in flight. h.bindingEpoch++ @@ -1039,6 +1202,7 @@ func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier h.reapStale[id] = epoch } h.mu.Unlock() + h.enrollMu.Unlock() // Choose the reap set. With a durable store the ROWS are reaped and the edges // driven from what was removed (the authoritative set); without one, the in-RAM @@ -1113,6 +1277,15 @@ func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier // no Runner is enrolled. The id travels out with the router so a caller that must // attribute the call to a Runner names the one that served it, rather than re-reading // the registry and racing a re-enroll onto the wrong id. +// routerForStream returns the Runner's router and the enrollment generation it +// belongs to, read as one pair. +func (h *Hub) routerForStream(sessionID string) (*commandRouter, uint64, error) { + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + router, _, err := h.routerFor(sessionID) + return router, h.enrollGen, err +} + func (h *Hub) routerFor(sessionID string) (*commandRouter, string, error) { h.mu.Lock() defer h.mu.Unlock() diff --git a/go/internal/runnerhub/hub_test.go b/go/internal/runnerhub/hub_test.go index 044cd530a..0f125aaf8 100644 --- a/go/internal/runnerhub/hub_test.go +++ b/go/internal/runnerhub/hub_test.go @@ -149,6 +149,63 @@ func TestDeliverSequenceGapDetection(t *testing.T) { t.Fatal("SeenGap true on a contiguous 1,2,3; a gap-free stream must not flag") } }) + + // A restarted Runner's counter starts over, so the tracker must start over too. + t.Run("re-enroll restarts tracking", func(t *testing.T) { + hub := newHubOnly() + enroll := func() { + hub.enroll(context.Background(), testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + } + enroll() + deliverSeq(t, hub, 10) + deliverSeq(t, hub, 12) + enroll() + deliverSeq(t, hub, 1) + if hub.SeenGap() { + t.Fatal("SeenGap true after re-enroll; the old process's missing 11 must not carry over") + } + deliverSeq(t, hub, 3) + if !hub.SeenGap() { + t.Fatal("SeenGap false after the new process's 1 then 3") + } + }) + + // Untracked seqs can never be proven seen, so overflow holds for the enrollment. + t.Run("overflow stays a gap until re-enroll", func(t *testing.T) { + hub := newHubOnly() + enroll := func() { + hub.enroll(context.Background(), testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + } + enroll() + deliverSeq(t, hub, 1) + deliverSeq(t, hub, maxMissingSeqs+3) + for seq := uint64(2); seq < maxMissingSeqs+3; seq++ { + deliverSeq(t, hub, seq) + } + if !hub.SeenGap() { + t.Fatal("SeenGap false after an overflowed jump; untracked seqs cannot be proven seen") + } + enroll() + deliverSeq(t, hub, 1) + if hub.SeenGap() { + t.Fatal("SeenGap true after re-enroll; overflow must reset with the enrollment") + } + }) + + // Container Gateways share one counter on separate streams, so 2 can trail 3. + t.Run("late arrival closes the gap", func(t *testing.T) { + hub := newHubOnly() + deliverSeq(t, hub, 1) + deliverSeq(t, hub, 3) + deliverSeq(t, hub, 4) + if !hub.SeenGap() { + t.Fatal("SeenGap false while seq 2 is still missing") + } + deliverSeq(t, hub, 2) + if hub.SeenGap() { + t.Fatal("SeenGap true after the late seq 2 arrived; nothing was lost") + } + }) } // deliverSeq delivers one trace frame at seq n (a frame that touches no sink of diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 8fe38f17d..4a54409bd 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -5,11 +5,15 @@ package runnerhub import ( "context" "errors" + "runtime" + "slices" + "sync" "testing" "testing/synctest" "time" compassv1 "github.com/RigelBuild/compass/go/gen/compass/v1" + compassv1internal "github.com/RigelBuild/compass/go/internal/gen/compass/v1" "github.com/RigelBuild/compass/go/internal/store" ) @@ -63,13 +67,13 @@ func TestDropLostSessionOnlyForOwningRunner(t *testing.T) { hub.bindContainer("cont-1", testAgentAccount, "runner-1") hub.promoteSession(ctx, "cont-1", "sess-1") - hub.dropLostSession(ctx, "runner-other", "sess-1", false) + hub.dropLostSession(ctx, "runner-other", "sess-1", nil, false) if _, ok := hub.accountForSession(ctx, "sess-1"); !ok { t.Fatal("foreign refusal unbound sess-1") } sink.none(t, "foreign refusal") - hub.dropLostSession(ctx, "runner-1", "sess-1", false) + hub.dropLostSession(ctx, "runner-1", "sess-1", nil, false) if _, ok := hub.accountForSession(ctx, "sess-1"); ok { t.Fatal("owning refusal left sess-1 bound") } @@ -122,6 +126,798 @@ func TestErroredSessionDropsBindingAndReportsLoss(t *testing.T) { } } +func TestStaleStateAfterErroredIsIgnored(t *testing.T) { + ctx := t.Context() + hub, lifecycle, tail := newHub() + bindings := newFakeBindingStore() + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + settle := &fakeSettleSink{} + hub.SetSettleSink(settle) + presence := &fakePresenceSink{} + hub.SetPresenceSink(presence) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { + t.Fatalf("lost report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) + } + // The fallback ERRORED overtook a frame buffered on the cancelled shared stream. + deliverState(t, hub, 4, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + frames := tail.snapshot() + if len(frames) != 1 || frames[0].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("relayed session frames = %+v, want only ERRORED; a stale state must not reach the tail", frames) + } + if got := settle.snapshot(); len(got) != 1 || got[0].state != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("settle edges = %+v, want only ERRORED", got) + } + // The ERRORED loss path adds its own DISCONNECTED edge; only the stale WORKING is wrong. + for _, rec := range presence.lifecycleSnapshot() { + if rec.state == compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { + t.Fatalf("presence edges = %+v, want no edge for the stale frame", presence.lifecycleSnapshot()) + } + } +} + +func TestNewLifetimeStateAfterErroredPublishes(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + deliverState(t, hub, 6, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + // A dead-lifetime frame that arrives after recovery began must still not publish. + deliverState(t, hub, 4, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + deliverState(t, hub, 7, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + + assertPublished(t, lifecycle, + compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED, + compassv1.AgentSessionState_AGENT_SESSION_STATE_READY, + compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) +} + +func TestErroredOlderThanNewLifetimeIsIgnored(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + // Separate streams can deliver out of seq order: the resumed READY overtakes ERRORED. + deliverState(t, hub, 6, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + lost.none(t, "an older ERRORED must not retire the resumed lifetime") +} + +func TestErroredCleanupKeepsSameEnrollmentRebind(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingTenantBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old, ok := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + if !ok { + t.Fatal("sess-1 not bound after promote") + } + + // The old lifetime's cleanup pauses resolving the tenant; a resume re-binds + // the same session in the same enrollment before it continues. + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("ERRORED cleanup did not start") + } + hub.promoteSession(ctx, "cont-1", "sess-1") + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("ERRORED cleanup did not finish") + } + + lost.none(t, "the old lifetime's cleanup must not report the resumed session lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } + if account, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil || account != testAgentAccount { + t.Fatalf("durable binding after old cleanup = (%s, %v), want (%s, nil)", account, err, testAgentAccount) + } +} + +// pausingTenantBindingStore holds the first SessionBindingTenant until released. +type pausingTenantBindingStore struct { + *fakeBindingStore + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (b *pausingTenantBindingStore) SessionBindingTenant(ctx context.Context, sessionID, runnerID string) (store.TenantID, error) { + b.once.Do(func() { + close(b.entered) + <-b.release + }) + return b.fakeBindingStore.SessionBindingTenant(ctx, sessionID, runnerID) +} + +func TestErroredCleanupKeepsRebindDuringDurableDelete(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingDeleteBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + // The cleanup has matched the old binding and pauses in its durable delete; + // a resume re-binds the session before the delete runs. + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not reach the durable delete") + } + hub.promoteSession(ctx, "cont-1", "sess-1") + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not finish") + } + + lost.none(t, "a cleanup whose row was re-bound must not report the session lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } + if account, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil || account != testAgentAccount { + t.Fatalf("durable binding after old cleanup = (%s, %v), want (%s, nil)", account, err, testAgentAccount) + } +} + +func TestErroredCleanupKeepsPeerRebind(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingDeleteBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not reach the durable delete") + } + // Another Server re-binds the row; this hub's cache still holds the old binding. + if _, _, err := bindings.fakeBindingStore.RecordSessionBinding(ctx, "sess-1", testAgentAccount, "runner-1"); err != nil { + t.Fatalf("peer RecordSessionBinding: %v", err) + } + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not finish") + } + + lost.none(t, "a cleanup whose row a peer re-bound must not report the session lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } +} + +func TestErroredCleanupDropsStaleAccountForPeerRebind(t *testing.T) { + const otherAccount store.AccountID = "acct-other" + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingDeleteBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + hub.SetSessionLostSink(newRecordingLostSink()) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not reach the durable delete") + } + // A peer re-binds the session id to another account. + if err := bindings.fakeBindingStore.DeleteSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("peer DeleteSessionBinding: %v", err) + } + if _, _, err := bindings.fakeBindingStore.RecordSessionBinding(ctx, "sess-1", otherAccount, "runner-1"); err != nil { + t.Fatalf("peer RecordSessionBinding: %v", err) + } + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not finish") + } + + if sess, ok := hub.SessionForAccount(ctx, testAgentAccount); ok { + t.Fatalf("old account resolves to %q after the peer took the session, want a miss", sess) + } + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != otherAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, otherAccount) + } +} + +func TestErroredCleanupWithoutDurableWriteKeepsPeerRow(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := newFakeBindingStore() + bindings.recordErr = errors.New("store down") + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + // The durable write fails, so this hub caches a binding with no version. + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + if old.version != "" { + t.Fatalf("cached version after failed write = %q, want empty", old.version) + } + // A peer Server then writes the session's row. + bindings.mu.Lock() + bindings.recordErr = nil + bindings.mu.Unlock() + if _, _, err := bindings.RecordSessionBinding(ctx, "sess-1", testAgentAccount, "runner-1"); err != nil { + t.Fatalf("peer RecordSessionBinding: %v", err) + } + + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + if _, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("peer's durable binding after old cleanup: %v, want it kept", err) + } + lost.none(t, "a session the peer re-bound is not lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want the peer's (%s, true)", account, ok, testAgentAccount) + } +} + +func TestErroredCleanupWithoutDurableWriteKeepsLegacyRow(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := newFakeBindingStore() + bindings.recordErr = errors.New("store down") + hub.SetSessionBindingStore(bindings) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + // A row written before binding_version existed carries the column default "". + bindings.mu.Lock() + bindings.bindings["sess-1"] = store.SessionBinding{TenantID: bindings.tenant, SessionID: "sess-1", AccountID: testAgentAccount, RunnerID: "runner-1"} + bindings.versions["sess-1"] = "" + bindings.mu.Unlock() + + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + if _, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("legacy durable binding after version-less cleanup: %v, want it kept", err) + } +} + +func TestErroredCleanupReleasesReverseReadThroughBinding(t *testing.T) { + ctx := store.WithTenant(t.Context(), "tenant-a") + hub, _, _ := newHub() + bindings := newFakeBindingStore() + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + // Another instance promoted the session; this hub learns it from a delivery's reverse read. + bindings.seed("sess-1") + if sess, ok := hub.SessionForAccount(ctx, testAgentAccount); !ok || sess != "sess-1" { + t.Fatalf("SessionForAccount = (%q, %v), want (sess-1, true)", sess, ok) + } + gen := hub.EnrollGeneration() + hub.mu.Lock() + seen := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + hub.dropLostSessionIfCurrent(ctx, gen, testRunnerID, "sess-1", &seen, true) + if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { + t.Fatalf("loss report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) + } + if _, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, store.ErrNotFound) { + t.Fatalf("durable binding after ERRORED cleanup: %v, want ErrNotFound", err) + } +} + +// pausingDeleteBindingStore holds the first DeleteSessionBindingVersion until released. +type pausingDeleteBindingStore struct { + *fakeBindingStore + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (b *pausingDeleteBindingStore) DeleteSessionBindingVersion(ctx context.Context, sessionID, version string) (bool, error) { + b.once.Do(func() { + close(b.entered) + <-b.release + }) + return b.fakeBindingStore.DeleteSessionBindingVersion(ctx, sessionID, version) +} + +func TestLostSessionCleanupFromOldEnrollmentKeepsRebinding(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + oldGen := hub.EnrollGeneration() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + hub.dropLostSessionIfCurrent(ctx, oldGen, "runner-1", "sess-1", nil, true) + lost.none(t, "cleanup from the old enrollment must not unbind the re-bound session") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after stale cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } + + hub.dropLostSessionIfCurrent(ctx, hub.EnrollGeneration(), "runner-1", "sess-1", nil, true) + if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { + t.Fatalf("current-enrollment cleanup report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) + } +} + +func TestLowerSeqCannotOvertakeHigherSeqMidDelivery(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + store := &pausingResolveBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(store) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(store.release) }) } + t.Cleanup(release) + + // READY seq 6 has recorded its seq and pauses resolving the (uncached) binding. + readyDone := make(chan error, 1) + go func() { + readyDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 6, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_READY), + }) + }() + select { + case <-store.entered: + case <-time.After(10 * time.Second): + t.Fatal("READY did not reach the binding lookup") + } + if !sessionLockHeld(hub, "sess-1") { + t.Fatal("lifecycle lock free while READY seq 6 is mid-delivery; seq 5 could overtake it") + } + erroredDone := make(chan error, 1) + go func() { + erroredDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 5, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }) + }() + release() + for name, done := range map[string]chan error{"READY": readyDone, "ERRORED": erroredDone} { + select { + case err := <-done: + if err != nil { + t.Fatalf("Deliver(%s) = %v, want nil", name, err) + } + case <-time.After(10 * time.Second): + t.Fatalf("Deliver(%s) did not complete", name) + } + } + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) +} + +// pausingResolveBindingStore holds the first ResolveSessionBinding until released. +type pausingResolveBindingStore struct { + *fakeBindingStore + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (b *pausingResolveBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, string, error) { + b.once.Do(func() { + close(b.entered) + <-b.release + }) + return b.fakeBindingStore.ResolveSessionBinding(ctx, sessionID) +} + +// sessionLockHeld reports whether sessionID's lifecycle lock is held right now. +func sessionLockHeld(hub *Hub, sessionID string) bool { + hub.lifecycleMu.Lock() + l := hub.sessionLocks[sessionID] + hub.lifecycleMu.Unlock() + if l == nil { + return false + } + if l.mu.TryLock() { + l.mu.Unlock() + return false + } + return true +} + +func TestSlowBindingLookupStallsOnlyItsSession(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + store := &pausingResolveBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(store) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(store.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-2", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-2", "sess-2") + + slowDone := make(chan error, 1) + go func() { + slowDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_READY), + }) + }() + select { + case <-store.entered: + case <-time.After(10 * time.Second): + t.Fatal("sess-1 did not reach the binding lookup") + } + fastDone := make(chan error, 1) + go func() { + fastDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-2", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }) + }() + select { + case err := <-fastDone: + if err != nil { + t.Fatalf("Deliver(sess-2) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("sess-2 lifecycle delivery stalled behind sess-1's binding lookup") + } + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + release() + select { + case err := <-slowDone: + if err != nil { + t.Fatalf("Deliver(sess-1) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("sess-1 did not complete") + } +} + +func TestReenrollClearsErroredBoundary(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + + // A restarted Runner's counter starts over, so its low seqs are a new lifetime. + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + deliverState(t, hub, 1, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + + assertPublished(t, lifecycle, + compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED, + compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) +} + +func TestReenrollWaitsForInFlightErrored(t *testing.T) { + ctx := t.Context() + lifecycle := &fakeLifecycleSink{} + tail := &pausingTailSink{entered: make(chan struct{}), release: make(chan struct{})} + hub := NewHub(lifecycle, tail, nil, discardLogger()) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(tail.release) }) } + t.Cleanup(release) + + oldDone := make(chan error, 1) + go func() { + oldDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 5, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }) + }() + select { + case <-tail.entered: + case <-time.After(10 * time.Second): + t.Fatal("old ERRORED did not reach the tail sink") + } + if hub.enrollMu.TryLock() { + hub.enrollMu.Unlock() + t.Fatal("enroll lock free while a session frame is mid-delivery") + } + enrolled := make(chan struct{}) + go func() { + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + close(enrolled) + }() + // A pending writer refuses new readers, so the ERRORED's cleanup queues behind enroll. + for deadline := time.Now().Add(10 * time.Second); hub.enrollMu.TryRLock(); { + hub.enrollMu.RUnlock() + if time.Now().After(deadline) { + t.Fatal("re-enroll never queued behind the in-flight delivery") + } + runtime.Gosched() + } + release() + select { + case err := <-oldDone: + if err != nil { + t.Fatalf("Deliver(old ERRORED) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("old ERRORED did not complete after release") + } + select { + case <-enrolled: + case <-time.After(10 * time.Second): + t.Fatal("re-enroll did not complete") + } + if _, _, err := hub.routerFor("sess-1"); err != nil { + t.Fatalf("routerFor(sess-1) after re-bind = %v, want the binding kept", err) + } + deliverState(t, hub, 1, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + statuses := lifecycle.snapshot() + if n := len(statuses); n == 0 || statuses[n-1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_READY { + t.Fatalf("published statuses = %+v, want the new enrollment's READY last", statuses) + } +} + +// pausingTailSink holds the first ERRORED frame in RelaySessionFrame until released. +type pausingTailSink struct { + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (s *pausingTailSink) RelaySessionFrame(_ string, frame *compassv1internal.SessionFrame) { + if frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + return + } + s.once.Do(func() { + close(s.entered) + <-s.release + }) +} + +func TestFramesFromStreamBeforeReenrollAreDropped(t *testing.T) { + ctx := t.Context() + hub, lifecycle, tail := newHub() + settle := &fakeSettleSink{} + hub.SetSettleSink(settle) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + oldGen := hub.EnrollGeneration() + + // The Runner restarts and resumes the same session id before the old stream drains. + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 9, SessionID: "sess-1", EnrollGen: oldGen, + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }); err != nil { + t.Fatalf("Deliver(old-stream ERRORED) = %v, want nil", err) + } + + if got := lifecycle.snapshot(); len(got) != 0 { + t.Fatalf("published statuses = %+v, want none from the old stream", got) + } + if got := tail.snapshot(); len(got) != 0 { + t.Fatalf("relayed frames = %+v, want none from the old stream", got) + } + if got := settle.snapshot(); len(got) != 0 { + t.Fatalf("settle edges = %+v, want none from the old stream", got) + } + lost.none(t, "an old-stream ERRORED must not unbind the resumed session") + deliverState(t, hub, 1, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) +} + +func deliverState(t *testing.T, hub *Hub, seq uint64, state compassv1.AgentSessionState) { + t.Helper() + if err := hub.Deliver(t.Context(), RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: seq, SessionID: "sess-1", Frame: sessionStateFrame(state), + }); err != nil { + t.Fatalf("Deliver(%s, seq %d) = %v, want nil", state, seq, err) + } +} + +func assertPublished(t *testing.T, lifecycle *fakeLifecycleSink, want ...compassv1.AgentSessionState) { + t.Helper() + statuses := lifecycle.snapshot() + got := make([]compassv1.AgentSessionState, 0, len(statuses)) + for _, s := range statuses { + got = append(got, s.GetState()) + } + if !slices.Equal(got, want) { + t.Fatalf("published states = %v, want %v", got, want) + } +} + +func TestConcurrentErroredPublishesAfterInFlightLifecycle(t *testing.T) { + ctx := t.Context() + lifecycle := &blockingLifecycleSink{workingEntered: make(chan struct{}), releaseWorking: make(chan struct{})} + tail := &blockingTailSink{errored: make(chan struct{})} + var releaseOnce sync.Once + releaseWorking := func() { releaseOnce.Do(func() { close(lifecycle.releaseWorking) }) } + t.Cleanup(releaseWorking) + + hub := NewHub(lifecycle, tail, nil, discardLogger()) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + workingDone := make(chan error, 1) + go func() { + workingDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }) + }() + select { + case <-lifecycle.workingEntered: + case <-time.After(10 * time.Second): + t.Fatal("WORKING lifecycle publish did not block") + } + if !sessionLockHeld(hub, "sess-1") { + t.Fatal("lifecycle lock released while WORKING is still being published") + } + + erroredDone := make(chan error, 1) + go func() { + erroredDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }) + }() + releaseWorking() + select { + case err := <-workingDone: + if err != nil { + t.Fatalf("Deliver(WORKING) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("Deliver(WORKING) did not complete after release") + } + select { + case err := <-erroredDone: + if err != nil { + t.Fatalf("Deliver(ERRORED) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("Deliver(ERRORED) did not complete after WORKING release") + } + select { + case <-tail.errored: + default: + t.Fatal("ERRORED never reached the tail sink") + } + + statuses := lifecycle.snapshot() + if len(statuses) != 2 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING || statuses[1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("published lifecycle order = %+v, want WORKING then ERRORED", statuses) + } +} + +type blockingLifecycleSink struct { + mu sync.Mutex + statuses []*compassv1.AgentSessionStatus + workingEntered chan struct{} + releaseWorking chan struct{} +} + +func (s *blockingLifecycleSink) PublishSessionStatus(status *compassv1.AgentSessionStatus) { + if status.GetState() == compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { + close(s.workingEntered) + <-s.releaseWorking + } + s.mu.Lock() + s.statuses = append(s.statuses, status) + s.mu.Unlock() +} + +func (s *blockingLifecycleSink) snapshot() []*compassv1.AgentSessionStatus { + s.mu.Lock() + defer s.mu.Unlock() + return append([]*compassv1.AgentSessionStatus(nil), s.statuses...) +} + +type blockingTailSink struct { + errored chan struct{} +} + +func (s *blockingTailSink) RelaySessionFrame(_ string, frame *compassv1internal.SessionFrame) { + if frame.GetState() == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + close(s.errored) + } +} + type chanEndSink chan string func (c chanEndSink) OnSessionEnded(_ context.Context, sessionID string) { c <- sessionID } @@ -223,8 +1019,8 @@ func TestSessionEndedWithoutStopIsArchived(t *testing.T) { hub.bindContainer("cont-1", testAgentAccount, "runner-1") hub.promoteSession(ctx, "cont-1", "sess-lost") - hub.dropLostSession(ctx, "runner-other", "sess-lost", false) - hub.dropLostSession(ctx, "runner-1", "sess-lost", false) + hub.dropLostSession(ctx, "runner-other", "sess-lost", nil, false) + hub.dropLostSession(ctx, "runner-1", "sess-lost", nil, false) if got := recvEnded(t, ended); got != "sess-lost" { t.Fatalf("archived %q after the lost drop, want sess-lost", got) } diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index 382d9e3d6..c921e1595 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -84,10 +84,10 @@ func (h *Hub) promoteSession(ctx context.Context, containerName, sessionID strin // Write the durable binding first (h.mu released). The store upsert is keyed // on the account, so it lands whether or not the account held a prior // session, and it returns the displaced session id. - var displaced string + var displaced, version string tenant := "" if bindings != nil && runnerID != "" { - d, err := bindings.RecordSessionBinding(ctx, sessionID, account, runnerID) + d, v, err := bindings.RecordSessionBinding(ctx, sessionID, account, runnerID) if err != nil { // A durable-write fault must not fail the Start that already succeeded // on the Runner: log it and fall back to the in-RAM cache so the session @@ -96,7 +96,7 @@ func (h *Hub) promoteSession(ctx context.Context, containerName, sessionID strin h.log.Error("record session binding failed; falling back to in-RAM cache", "session_id", sessionID, "account", string(account), "error", err) } else { - displaced = d + displaced, version = d, v tenant = string(bindings.EffectiveTenant(ctx)) } } @@ -105,7 +105,7 @@ func (h *Hub) promoteSession(ctx context.Context, containerName, sessionID strin h.mu.Lock() // Keep the container->account entry: a resume Start on this container fetches // secrets by container_name before exec. Remove and re-enroll clear it. - h.sessionAccounts[sessionID] = sessionBinding{account: account, runnerID: runnerID} + h.sessionAccounts[sessionID] = h.newBindingLocked(account, runnerID, version) h.accountSessions[account] = sessionID // Evict the displaced session from the forward map: the account moved off it, // so it now resolves nowhere. Guard displaced != sessionID for the rebind @@ -180,24 +180,66 @@ func (h *Hub) publishBindingChange(ctx context.Context, routing RoutingFabric, t // lock-then-release-then-fire discipline promoteSession uses, so the sink (which // enqueues into the presence loop) never runs under h.mu. func (h *Hub) unbindSession(ctx context.Context, sessionID string) { + h.releaseSession(ctx, sessionID, nil) +} + +// releaseSession is unbindSession limited, when only is non-nil, to that one +// binding of sessionID, so a release from an older lifetime leaves a re-bind +// alone. It reports whether the binding was released. +func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessionBinding) bool { // The maps are a cache, so the durable row is deleted FIRST (on the request ctx // so it stays tenant-scoped), then the maps are evicted under h.mu. Delete is // by session id and idempotent — a session already displaced has no row, so a // stale release matches nothing and leaves the live binding alone. + // A Stop's release takes bindingWriteMu so it orders wholly before or after an + // enroll's map-clear through reap. A limited release skips it: mid-reap no binding + // is cached (promotion waits here, read-through is refused), so it releases nothing. + unlockWrite := func() {} + if only == nil { + h.bindingWriteMu.Lock() + unlockWrite = h.bindingWriteMu.Unlock + } h.mu.Lock() bindings := h.bindings routing := h.routing + if only != nil { + if live, ok := h.sessionAccounts[sessionID]; ok && !sameBinding(live, *only) { + h.mu.Unlock() + unlockWrite() + return false + } + } h.mu.Unlock() tenant := "" if bindings != nil { - if err := bindings.DeleteSessionBinding(ctx, sessionID); err != nil { + // A limited release deletes only its own write. A version-less entry + // wrote no row, so it deletes none: a legacy or peer row is not its own. + removed := true + var err error + switch { + case only == nil: + err = bindings.DeleteSessionBinding(ctx, sessionID) + case only.version != "": + removed, err = bindings.DeleteSessionBindingVersion(ctx, sessionID, only.version) + default: + // Any row is someone else's write, so the session is still bound there. + _, _, _, rerr := bindings.ResolveSessionBinding(ctx, sessionID) + removed = errors.Is(rerr, store.ErrNotFound) + } + switch { + case err != nil: // A durable-delete fault must not fail the Stop that already // succeeded on the Runner: log and continue to evict the cache. The // next re-enroll sweep retires any surviving row. h.log.Error("delete session binding failed; evicting cache anyway", "session_id", sessionID, "error", err) - } else { + case !removed: + // The row was re-bound since only was read: the session is not lost. + h.evictStaleBinding(sessionID, only) + unlockWrite() + return false + default: tenant = string(bindings.EffectiveTenant(ctx)) } } @@ -207,7 +249,14 @@ func (h *Hub) unbindSession(ctx context.Context, sessionID string) { account store.AccountID wentOffline bool ) - if binding, ok := h.sessionAccounts[sessionID]; ok { + binding, ok := h.sessionAccounts[sessionID] + if ok && only != nil && !sameBinding(binding, *only) { + // Re-bound in the cache while the durable delete ran. + h.mu.Unlock() + unlockWrite() + return false + } + if ok { // Drop the reverse entry only if it still points at THIS session — a // promoteSession for the account onto a newer session would have already // repointed it, and a stale delete would then unbind the live one. @@ -220,6 +269,7 @@ func (h *Hub) unbindSession(ctx context.Context, sessionID string) { delete(h.sessionAccounts, sessionID) presence := h.presence h.mu.Unlock() + unlockWrite() // Invalidate peer instances' caches (h.mu released, nil-safe, best-effort): // this session has no binding any more (BindingUnbound), so a peer can drop @@ -238,6 +288,29 @@ func (h *Hub) unbindSession(ctx context.Context, sessionID string) { if wentOffline && presence != nil { presence.OnSessionLifecycle(account, sessionID, compassv1.AgentSessionState_AGENT_SESSION_STATE_DISCONNECTED) } + return true +} + +// evictStaleBinding drops sessionID's cache entry while it is still stale, so the +// next read resolves the row that replaced it. +func (h *Hub) evictStaleBinding(sessionID string, stale *sessionBinding) { + if stale == nil { + return + } + h.mu.Lock() + defer h.mu.Unlock() + if live, ok := h.sessionAccounts[sessionID]; ok && sameBinding(live, *stale) { + delete(h.sessionAccounts, sessionID) + if h.accountSessions[stale.account] == sessionID { + delete(h.accountSessions, stale.account) + } + } +} + +// newBindingLocked builds a sessionAccounts entry with a fresh lifetime. Caller holds mu. +func (h *Hub) newBindingLocked(account store.AccountID, runnerID, version string) sessionBinding { + h.lastLifetime++ + return sessionBinding{account: account, runnerID: runnerID, version: version, lifetime: h.lastLifetime} } // unbindContainer drops a container's provisioned account binding. Remove is the @@ -357,7 +430,7 @@ func (h *Hub) lookupSessionBinding(ctx context.Context, sessionID string) (sessi if !h.readThroughAllowed(ctx, bindings, enrolled) { return sessionBinding{}, bindingUnverifiable } - account, runnerID, err := bindings.ResolveSessionBinding(ctx, sessionID) + account, runnerID, version, err := bindings.ResolveSessionBinding(ctx, sessionID) if errors.Is(err, store.ErrNotFound) { return sessionBinding{}, bindingNotFound } @@ -368,7 +441,6 @@ func (h *Hub) lookupSessionBinding(ctx context.Context, sessionID string) (sessi // session hits without a table round-trip. Re-check under the lock: a // concurrent promote/unbind may have run, so a live map entry wins over the // row just read (avoids clobbering a fresher binding with a staler one). - resolved := sessionBinding{account: account, runnerID: runnerID} h.mu.Lock() if live, ok := h.sessionAccounts[sessionID]; ok { h.mu.Unlock() @@ -380,6 +452,7 @@ func (h *Hub) lookupSessionBinding(ctx context.Context, sessionID string) (sessi h.mu.Unlock() return sessionBinding{}, bindingUnverifiable } + resolved := h.newBindingLocked(account, runnerID, version) h.sessionAccounts[sessionID] = resolved h.mu.Unlock() return resolved, bindingFound @@ -430,7 +503,7 @@ func (h *Hub) SessionForAccount(ctx context.Context, account store.AccountID) (s if !h.readThroughAllowed(ctx, bindings, enrolled) { return "", false } - sessionID, ownerID, err := bindings.SessionForAccount(ctx, account) + sessionID, ownerID, version, err := bindings.SessionForAccount(ctx, account) if err != nil || ownerID != runnerID { return "", false } @@ -448,7 +521,7 @@ func (h *Hub) SessionForAccount(ctx context.Context, account store.AccountID) (s return "", false } h.accountSessions[account] = sessionID - h.sessionAccounts[sessionID] = sessionBinding{account: account, runnerID: ownerID} + h.sessionAccounts[sessionID] = h.newBindingLocked(account, ownerID, version) h.mu.Unlock() return sessionID, true } @@ -890,38 +963,62 @@ func commsCallError(err error) *compassv1internal.CommsCallError { } // dropLostSession unbinds, archives and reports a session its Runner lost. Only -// the owning Runner may unbind it. -func (h *Hub) dropLostSession(ctx context.Context, runnerID, sessionID string, errored bool) { +// the owning Runner may unbind it. A non-nil seen limits it to that binding: a +// resume may have re-bound the session id since the loss was observed. +func (h *Hub) dropLostSession(ctx context.Context, runnerID, sessionID string, seen *sessionBinding, errored bool) { ctx, scoped := h.runnerSessionCtx(ctx, runnerID, sessionID) if !scoped { return } - account, ok := h.accountForRunnerSession(ctx, runnerID, sessionID) - if !ok { + binding, ok := h.resolveSessionBinding(ctx, sessionID) + if !ok || runnerID == "" || binding.runnerID != runnerID || (seen != nil && !sameBinding(binding, *seen)) { + return + } + if !h.releaseSession(ctx, sessionID, &binding) { return } - h.unbindSession(ctx, sessionID) h.mu.Lock() lost := h.lost h.mu.Unlock() h.archiveEnded(ctx, sessionID) h.log.Warn("runner reports bound session lost; released binding", - "session_id", sessionID, "agent_account_id", account, "errored", errored) + "session_id", sessionID, "agent_account_id", binding.account, "errored", errored) if lost != nil { - lost.OnSessionLost(sessionID, account, errored) + lost.OnSessionLost(sessionID, binding.account, errored) } } // dropLostSessionDetached runs dropLostSession off the caller's receive loop: it // does store work, and the stream ctx dies with the stream. -func (h *Hub) dropLostSessionDetached(ctx context.Context, runnerID, sessionID string, errored bool) { +// gen is the enrollment the loss was observed under. +func (h *Hub) dropLostSessionDetached(ctx context.Context, gen uint64, runnerID, sessionID string, seen *sessionBinding, errored bool) { go func() { dctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), lostSessionTimeout) defer cancel() - h.dropLostSession(dctx, runnerID, sessionID, errored) + h.dropLostSessionIfCurrent(dctx, gen, runnerID, sessionID, seen, errored) }() } +// dropLostSessionIfCurrent runs dropLostSession only while enrollment gen is +// current: a re-enroll since the loss may have re-bound sessionID. +func (h *Hub) dropLostSessionIfCurrent(ctx context.Context, gen uint64, runnerID, sessionID string, seen *sessionBinding, errored bool) { + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + if h.enrollGen != gen { + return + } + h.dropLostSession(ctx, runnerID, sessionID, seen, errored) +} + +// sameBinding reports whether a and b are one binding of a session id: the same +// durable row when either has one, else the same cache entry. +func sameBinding(a, b sessionBinding) bool { + if a.version != "" || b.version != "" { + return a.version == b.version + } + return a.lifetime == b.lifetime +} + // runnerSessionCtx scopes a Runner-originated ctx to the tenant that binds // sessionID to runnerID. Runners are shared across tenants, so the door sets none; // without this, binding reads and deletes run under the bootstrap tenant. The diff --git a/go/internal/runnerhub/runner_tenant_pgtest_test.go b/go/internal/runnerhub/runner_tenant_pgtest_test.go index ba00b6e7a..6025cd49e 100644 --- a/go/internal/runnerhub/runner_tenant_pgtest_test.go +++ b/go/internal/runnerhub/runner_tenant_pgtest_test.go @@ -28,16 +28,16 @@ func TestDropLostSessionScopesToTheSessionTenant(t *testing.T) { hub.SetSessionLostSink(sink) hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) // Enroll reaps every binding on this Runner; re-record so the drop reads a live row cold. - if _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { + if _, _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { t.Fatalf("RecordSessionBinding after enroll: %v", err) } - hub.dropLostSession(ctx, "runner-1", "sess-b", false) + hub.dropLostSession(ctx, "runner-1", "sess-b", nil, false) if lost, _ := sink.waitOne(t); lost != agent.ID { t.Fatalf("lost = %s, want %s: the tenant-B session was not resolved, so no wake", lost, agent.ID) } - if _, _, err := st.ResolveSessionBinding(ctxB, "sess-b"); err == nil { + if _, _, _, err := st.ResolveSessionBinding(ctxB, "sess-b"); err == nil { t.Fatal("tenant B's durable binding survived the drop; a cache miss would resurrect it") } } @@ -101,7 +101,7 @@ func openTenantBSession(t *testing.T, ctx context.Context) (*store.Store, store. if err != nil { t.Fatalf("CreateAgent: %v", err) } - if _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { + if _, _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { t.Fatalf("RecordSessionBinding: %v", err) } return st, agent, ctxB diff --git a/go/internal/runnerhub/seam_test.go b/go/internal/runnerhub/seam_test.go index ac8154d71..908f26034 100644 --- a/go/internal/runnerhub/seam_test.go +++ b/go/internal/runnerhub/seam_test.go @@ -23,7 +23,9 @@ import ( "context" "errors" "io" + "runtime" "testing" + "time" "connectrpc.com/connect" @@ -121,6 +123,47 @@ func TestSeamEnrollSessionsRoundTripAndPublishEvents(t *testing.T) { } } +func TestSeamPublishEventsOpenedBeforeEnrollIsFenced(t *testing.T) { + hub, _, tail := newHub() + resolver := &fakeResolver{tokens: map[string]resolverEntry{ + "runner-tok": {subj: store.Subject{Kind: store.SubjectRunner, ID: "runner-1"}}, + }} + url := newMountedH2CServer(t, hub, resolver.resolve) + client := newRawRunnerClient(t, url, "runner-tok") + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + pub := client.PublishEvents(ctx) + if err := pub.Send(&compassv1internal.PublishEventsRequest{RunnerSeq: 1, SessionId: "sess-wire", Frame: sessionTraceFrame("before enroll")}); err != nil { + t.Fatalf("PublishEvents.Send = %v", err) + } + // Gate on the server handler having taken the frame, so the stream opened pre-enroll. + for deadline := time.Now().Add(10 * time.Second); ; runtime.Gosched() { + hub.mu.Lock() + seen := hub.lastSeq == 1 + hub.mu.Unlock() + if seen { + break + } + if time.Now().After(deadline) { + t.Fatal("pre-enroll frame never reached Deliver") + } + } + if _, err := client.Enroll(ctx, connect.NewRequest(&compassv1internal.EnrollRequest{RunnerId: "runner-1"})); err != nil { + t.Fatalf("Enroll = %v, want success", err) + } + bindSession(hub, "sess-wire") + if err := pub.Send(&compassv1internal.PublishEventsRequest{RunnerSeq: 2, SessionId: "sess-wire", Frame: sessionTraceFrame("after enroll")}); err != nil { + t.Fatalf("PublishEvents.Send = %v", err) + } + if _, err := pub.CloseAndReceive(); err != nil { + t.Fatalf("PublishEvents.CloseAndReceive = %v", err) + } + if calls := tail.snapshot(); len(calls) != 0 { + t.Fatalf("tail sink saw %d frames, want none from a stream opened before enroll", len(calls)) + } +} + // startRoundTrip is the outcome of Hub.Start driven over the wire. type startRoundTrip struct { resp *compassv1.StartAgentSessionResponse diff --git a/go/internal/store/compute_usage_pgtest_test.go b/go/internal/store/compute_usage_pgtest_test.go index 72a24cff1..3c29011a6 100644 --- a/go/internal/store/compute_usage_pgtest_test.go +++ b/go/internal/store/compute_usage_pgtest_test.go @@ -159,7 +159,7 @@ func TestComputeUsageEventsStampedAfterLockWait(t *testing.T) { } done := make(chan error, 1) go func() { - _, err := s.RecordSessionBinding(ctx, "wait-after", agent.ID, "runner-1") + _, _, err := s.RecordSessionBinding(ctx, "wait-after", agent.ID, "runner-1") done <- err }() deadline := time.After(5 * time.Second) @@ -223,7 +223,7 @@ func TestComputeUsageConflictRollsBackIntervalEvents(t *testing.T) { mustBind(t, ctx, s, "sess-2", accountY.ID, "runner-1") before := computeEvents(t, s, tenant, accountX.ID) - if _, err := s.RecordSessionBinding(ctx, "sess-2", accountX.ID, "runner-1"); !errors.Is(err, ErrConflict) { + if _, _, err := s.RecordSessionBinding(ctx, "sess-2", accountX.ID, "runner-1"); !errors.Is(err, ErrConflict) { t.Fatalf("RecordSessionBinding(X, sess-2) error = %v, want ErrConflict", err) } after := computeEvents(t, s, tenant, accountX.ID) @@ -231,7 +231,7 @@ func TestComputeUsageConflictRollsBackIntervalEvents(t *testing.T) { after[0].IntervalID != before[0].IntervalID || !after[0].OccurredAt.Equal(before[0].OccurredAt) { t.Fatalf("account X events before=%+v after=%+v, want its unchanged open start only", before, after) } - if sessionID, runnerID, err := s.SessionForAccount(ctx, accountX.ID); err != nil || sessionID != "sess-1" || runnerID != "runner-1" { + if sessionID, runnerID, _, err := s.SessionForAccount(ctx, accountX.ID); err != nil || sessionID != "sess-1" || runnerID != "runner-1" { t.Fatalf("account X resolves to (%q, %q, %v), want (sess-1, runner-1)", sessionID, runnerID, err) } } diff --git a/go/internal/store/db/models.go b/go/internal/store/db/models.go index d85b737e4..75e215d4a 100644 --- a/go/internal/store/db/models.go +++ b/go/internal/store/db/models.go @@ -380,6 +380,7 @@ type SessionBinding struct { CreatedAt pgtype.Timestamptz UpdatedAt pgtype.Timestamptz UsageIntervalID string + BindingVersion string } type SystemAccount struct { diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index adde0c45d..3f8e0a787 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -140,6 +140,10 @@ type Querier interface { // Both deletes also write an estimated start for a binding an older server made // without one; ON CONFLICT keeps any real start. DeleteSessionBinding(ctx context.Context, sessionID string) error + // DeleteSessionBinding limited to one write of the row: a re-bind since that + // write set a new binding_version, so it is left alone. Returns rows removed; + // Postgres runs every data-modifying CTE to completion. + DeleteSessionBindingVersion(ctx context.Context, arg DeleteSessionBindingVersionParams) (int64, error) // The reconnect sweep, run by Hub.enroll under the system role because a Runner // is shared across tenants. :many with RETURNING: each removed row drives a // presence DISCONNECTED edge, a held-deliver reap, and a tenant-scoped archive. diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go index b320e0a24..ee760b341 100644 --- a/go/internal/store/db/session_bindings.sql.go +++ b/go/internal/store/db/session_bindings.sql.go @@ -48,6 +48,53 @@ func (q *Queries) DeleteSessionBinding(ctx context.Context, sessionID string) er return err } +const deleteSessionBindingVersion = `-- name: DeleteSessionBindingVersion :one +WITH d AS ( + DELETE FROM session_bindings AS b + WHERE b.session_id = $1 AND b.binding_version = $2 + RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, + b.session_id, b.runner_id, b.created_at +), starts AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id, estimated + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'start', d.created_at, + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id, TRUE + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 +), ends AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 +) +SELECT count(*) FROM d +` + +type DeleteSessionBindingVersionParams struct { + SessionID string + BindingVersion string +} + +// DeleteSessionBinding limited to one write of the row: a re-bind since that +// write set a new binding_version, so it is left alone. Returns rows removed; +// Postgres runs every data-modifying CTE to completion. +func (q *Queries) DeleteSessionBindingVersion(ctx context.Context, arg DeleteSessionBindingVersionParams) (int64, error) { + row := q.db.QueryRow(ctx, deleteSessionBindingVersion, arg.SessionID, arg.BindingVersion) + var count int64 + err := row.Scan(&count) + return count, err +} + const deleteSessionBindingsForRunner = `-- name: DeleteSessionBindingsForRunner :many WITH d AS ( DELETE FROM session_bindings AS b @@ -226,12 +273,13 @@ func (q *Queries) LockSessionBindingAccount(ctx context.Context, arg LockSession } const recordSessionBinding = `-- name: RecordSessionBinding :exec -INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id) -VALUES ($1, $2, $3, $4) +INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id, binding_version) +VALUES ($1, $2, $3, $4, $5) ON CONFLICT (tenant_id, agent_account_id) DO UPDATE SET session_id = EXCLUDED.session_id, runner_id = EXCLUDED.runner_id, - usage_interval_id = EXCLUDED.usage_interval_id + usage_interval_id = EXCLUDED.usage_interval_id, + binding_version = EXCLUDED.binding_version ` type RecordSessionBindingParams struct { @@ -239,6 +287,7 @@ type RecordSessionBindingParams struct { SessionID string RunnerID string UsageIntervalID string + BindingVersion string } // What it DISPLACED comes from SessionBindingForUpdate above, not from a @@ -249,39 +298,42 @@ func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBin arg.SessionID, arg.RunnerID, arg.UsageIntervalID, + arg.BindingVersion, ) return err } const sessionBinding = `-- name: SessionBinding :one -SELECT agent_account_id, runner_id FROM session_bindings WHERE session_id = $1 +SELECT agent_account_id, runner_id, binding_version FROM session_bindings WHERE session_id = $1 ` type SessionBindingRow struct { AgentAccountID string RunnerID string + BindingVersion string } func (q *Queries) SessionBinding(ctx context.Context, sessionID string) (SessionBindingRow, error) { row := q.db.QueryRow(ctx, sessionBinding, sessionID) var i SessionBindingRow - err := row.Scan(&i.AgentAccountID, &i.RunnerID) + err := row.Scan(&i.AgentAccountID, &i.RunnerID, &i.BindingVersion) return i, err } const sessionBindingForAccount = `-- name: SessionBindingForAccount :one -SELECT session_id, runner_id FROM session_bindings WHERE agent_account_id = $1 +SELECT session_id, runner_id, binding_version FROM session_bindings WHERE agent_account_id = $1 ` type SessionBindingForAccountRow struct { - SessionID string - RunnerID string + SessionID string + RunnerID string + BindingVersion string } func (q *Queries) SessionBindingForAccount(ctx context.Context, agentAccountID string) (SessionBindingForAccountRow, error) { row := q.db.QueryRow(ctx, sessionBindingForAccount, agentAccountID) var i SessionBindingForAccountRow - err := row.Scan(&i.SessionID, &i.RunnerID) + err := row.Scan(&i.SessionID, &i.RunnerID, &i.BindingVersion) return i, err } diff --git a/go/internal/store/migrations/0003_session_binding_version.sql b/go/internal/store/migrations/0003_session_binding_version.sql new file mode 100644 index 000000000..f9574b32a --- /dev/null +++ b/go/internal/store/migrations/0003_session_binding_version.sql @@ -0,0 +1,9 @@ +-- binding_version names one write of a session_bindings row. Every upsert sets a +-- fresh UUID, so a release conditioned on it cannot remove a later re-bind of the +-- same session id. Not xmin: transaction ids wrap around and repeat. +ALTER TABLE session_bindings + ADD COLUMN binding_version TEXT NOT NULL DEFAULT ''; + +-- Give rows written before this column a real version too. +UPDATE session_bindings SET binding_version = gen_random_uuid()::TEXT +WHERE binding_version = ''; diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql index 45d6cf932..0850e2f1f 100644 --- a/go/internal/store/queries/session_bindings.sql +++ b/go/internal/store/queries/session_bindings.sql @@ -115,18 +115,19 @@ ON CONFLICT (tenant_id, interval_id, kind) DO NOTHING; -- What it DISPLACED comes from SessionBindingForUpdate above, not from a -- RETURNING here. The binding update and event writes share the Store tx. -- name: RecordSessionBinding :exec -INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id) -VALUES ($1, $2, $3, $4) +INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id, binding_version) +VALUES ($1, $2, $3, $4, $5) ON CONFLICT (tenant_id, agent_account_id) DO UPDATE SET session_id = EXCLUDED.session_id, runner_id = EXCLUDED.runner_id, - usage_interval_id = EXCLUDED.usage_interval_id; + usage_interval_id = EXCLUDED.usage_interval_id, + binding_version = EXCLUDED.binding_version; -- name: SessionBinding :one -SELECT agent_account_id, runner_id FROM session_bindings WHERE session_id = $1; +SELECT agent_account_id, runner_id, binding_version FROM session_bindings WHERE session_id = $1; -- name: SessionBindingForAccount :one -SELECT session_id, runner_id FROM session_bindings WHERE agent_account_id = $1; +SELECT session_id, runner_id, binding_version FROM session_bindings WHERE agent_account_id = $1; -- Both deletes also write an estimated start for a binding an older server made -- without one; ON CONFLICT keeps any real start. @@ -159,6 +160,40 @@ SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_t ORDER BY d.tenant_id, d.usage_interval_id ON CONFLICT DO NOTHING; +-- DeleteSessionBinding limited to one write of the row: a re-bind since that +-- write set a new binding_version, so it is left alone. Returns rows removed; +-- Postgres runs every data-modifying CTE to completion. +-- name: DeleteSessionBindingVersion :one +WITH d AS ( + DELETE FROM session_bindings AS b + WHERE b.session_id = $1 AND b.binding_version = $2 + RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, + b.session_id, b.runner_id, b.created_at +), starts AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id, estimated + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'start', d.created_at, + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id, TRUE + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 +), ends AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 +) +SELECT count(*) FROM d; + -- The reconnect sweep, run by Hub.enroll under the system role because a Runner -- is shared across tenants. :many with RETURNING: each removed row drives a -- presence DISCONNECTED edge, a held-deliver reap, and a tenant-scoped archive. diff --git a/go/internal/store/session_bindings.go b/go/internal/store/session_bindings.go index ea236788f..625a1aaad 100644 --- a/go/internal/store/session_bindings.go +++ b/go/internal/store/session_bindings.go @@ -94,19 +94,19 @@ type SessionBinding struct { // resolves the principal a comms call runs under. Any OTHER unique violation is // unexpected, so it falls through to the generic wrap with its own message // intact rather than being relabelled as a session collision. -func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) (string, error) { +func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) (displaced, version string, err error) { if sessionID == "" { - return "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) + return "", "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) } if accountID == "" { - return "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + return "", "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) } // Unlike agent_placements.runner_id, '' is NOT an accepted unknown sentinel: // a placement outlives its Runner, but a binding exists only while a Runner // is attached, and runner_id is the sweep key that retires it. A '' binding // no real Runner's re-enroll could sweep would linger as a stale session. if runnerID == "" { - return "", fmt.Errorf("%w: runner id is required", ErrInvalidArgument) + return "", "", fmt.Errorf("%w: runner id is required", ErrInvalidArgument) } // beginTenantTx, not s.pool.Begin: it arms SET LOCAL ROLE + the @@ -115,7 +115,7 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco // would run as the owner with no GUC and silently disable tenant isolation. tx, err := s.beginTenantTx(ctx) if err != nil { - return "", fmt.Errorf("store: begin record session binding: %w", err) + return "", "", fmt.Errorf("store: begin record session binding: %w", err) } // No-op after a successful commit; the rollback that matters is on every // error path below, where there is nothing further to report about it. @@ -129,23 +129,23 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco Column1: pgtype.Text{String: string(s.resolveTenant(ctx)), Valid: true}, Column2: pgtype.Text{String: string(accountID), Valid: true}, }); err != nil { - return "", fmt.Errorf("store: lock session binding account: %w", err) + return "", "", fmt.Errorf("store: lock session binding account: %w", err) } prior, err := qtx.SessionBindingForUpdate(ctx, string(accountID)) if err != nil && !noRows(err) { - return "", fmt.Errorf("store: read prior session binding: %w", err) + return "", "", fmt.Errorf("store: read prior session binding: %w", err) } if noRows(err) { prior = db.SessionBindingForUpdateRow{} } - displaced := prior.SessionID + displaced = prior.SessionID intervalID := prior.UsageIntervalID if intervalID != "" { // An older server may have written the prior row without a start event. if err := qtx.EnsureComputeUsageIntervalStart(ctx, string(accountID)); err != nil { - return "", fmt.Errorf("store: ensure compute usage interval start: %w", err) + return "", "", fmt.Errorf("store: ensure compute usage interval start: %w", err) } } if prior.SessionID != sessionID { @@ -156,7 +156,7 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco SessionID: prior.SessionID, RunnerID: prior.RunnerID, }); err != nil { - return "", fmt.Errorf("store: end displaced compute usage interval: %w", err) + return "", "", fmt.Errorf("store: end displaced compute usage interval: %w", err) } } intervalID = uuid.NewString() @@ -166,28 +166,30 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco SessionID: sessionID, RunnerID: runnerID, }); err != nil { - return "", fmt.Errorf("store: start compute usage interval: %w", err) + return "", "", fmt.Errorf("store: start compute usage interval: %w", err) } } + version = uuid.NewString() if err := qtx.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ SessionID: sessionID, AgentAccountID: string(accountID), RunnerID: runnerID, UsageIntervalID: intervalID, + BindingVersion: version, }); err != nil { if pgErrIs(err, pgForeignKeyViolation) { - return "", fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) + return "", "", fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) } if pgErrIs(err, pgUniqueViolation) && pgConstraintName(err) == "session_bindings_session_key" { - return "", fmt.Errorf("%w: session %q is already bound to a different agent", ErrConflict, sessionID) + return "", "", fmt.Errorf("%w: session %q is already bound to a different agent", ErrConflict, sessionID) } - return "", fmt.Errorf("store: record session binding: %w", err) + return "", "", fmt.Errorf("store: record session binding: %w", err) } if err := tx.Commit(ctx); err != nil { - return "", fmt.Errorf("store: commit record session binding: %w", err) + return "", "", fmt.Errorf("store: commit record session binding: %w", err) } - return displaced, nil + return displaced, version, nil } // ResolveSessionBinding resolves the account and Runner a live session speaks @@ -197,18 +199,18 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco // resolves the scope a call runs under, so a miss must never surface as an // empty AccountID with a nil error. A zero-value account id would flow onward as // a real (wrong) principal instead of stopping the call. -func (s *Store) ResolveSessionBinding(ctx context.Context, sessionID string) (AccountID, string, error) { +func (s *Store) ResolveSessionBinding(ctx context.Context, sessionID string) (account AccountID, runnerID, version string, err error) { if sessionID == "" { - return "", "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) + return "", "", "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) } row, err := s.q.SessionBinding(ctx, sessionID) if err != nil { if noRows(err) { - return "", "", fmt.Errorf("%w: session %q is not bound", ErrNotFound, sessionID) + return "", "", "", fmt.Errorf("%w: session %q is not bound", ErrNotFound, sessionID) } - return "", "", fmt.Errorf("store: resolve session binding: %w", err) + return "", "", "", fmt.Errorf("store: resolve session binding: %w", err) } - return AccountID(row.AgentAccountID), row.RunnerID, nil + return AccountID(row.AgentAccountID), row.RunnerID, row.BindingVersion, nil } // SessionForAccount resolves the live session bound to an agent account — the @@ -228,19 +230,20 @@ func (s *Store) ResolveSessionBinding(ctx context.Context, sessionID string) (Ac // dropped on a Runner reconnect. Fail-closed for the same reason as above: an // empty session id with a nil error would be dispatched to as if it were a live // session. The consumer's own contract turns this into "push nothing now, let -// the cursor sweep deliver on the recipient's next start". -func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (string, string, error) { +// the cursor sweep deliver on the recipient's next start". version names the +// row's write, as ResolveSessionBinding returns it. +func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (sessionID, runnerID, version string, err error) { if accountID == "" { - return "", "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + return "", "", "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) } row, err := s.q.SessionBindingForAccount(ctx, string(accountID)) if err != nil { if noRows(err) { - return "", "", fmt.Errorf("%w: agent %q has no live session", ErrNotFound, accountID) + return "", "", "", fmt.Errorf("%w: agent %q has no live session", ErrNotFound, accountID) } - return "", "", fmt.Errorf("store: resolve session for account: %w", err) + return "", "", "", fmt.Errorf("store: resolve session for account: %w", err) } - return row.SessionID, row.RunnerID, nil + return row.SessionID, row.RunnerID, row.BindingVersion, nil } // DeleteSessionBinding releases the binding for sessionID — the single-session @@ -263,6 +266,20 @@ func (s *Store) DeleteSessionBinding(ctx context.Context, sessionID string) erro return nil } +// DeleteSessionBindingVersion is DeleteSessionBinding limited to the write that +// returned version (from Record or Resolve): a re-bind of the same session id +// since then has a new version and stays. removed reports whether a row went. +func (s *Store) DeleteSessionBindingVersion(ctx context.Context, sessionID, version string) (removed bool, err error) { + if sessionID == "" { + return false, fmt.Errorf("%w: session id is required", ErrInvalidArgument) + } + n, err := s.q.DeleteSessionBindingVersion(ctx, db.DeleteSessionBindingVersionParams{SessionID: sessionID, BindingVersion: version}) + if err != nil { + return false, fmt.Errorf("store: delete session binding version: %w", err) + } + return n > 0, nil +} + // DeleteSessionBindingsForRunner is the enroll sweep: it releases every binding // attached to runnerID and RETURNS the bindings it removed. Hub.enroll clears all // bindings on every Runner enroll, because an enrolling Runner has no live sessions diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index 65613e9b6..6c26fc997 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -35,7 +35,7 @@ import ( // ctx is the SECOND parameter, after t, matching mustTopic (messages_test.go). func mustBind(t *testing.T, ctx context.Context, s *Store, sessionID string, accountID AccountID, runnerID string) string { t.Helper() - displaced, err := s.RecordSessionBinding(ctx, sessionID, accountID, runnerID) + displaced, _, err := s.RecordSessionBinding(ctx, sessionID, accountID, runnerID) if err != nil { t.Fatalf("RecordSessionBinding(%q, %q, %q): %v", sessionID, accountID, runnerID, err) } @@ -91,7 +91,7 @@ func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - displaced, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") + displaced, _, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") if err != nil { t.Fatalf("RecordSessionBinding: %v", err) } @@ -99,7 +99,7 @@ func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { t.Fatalf("first bind displaced %q, want \"\" — the account held no prior session, and a caller reaping a phantom id would clear a live registry entry", displaced) } - gotAccount, gotRunner, err := s.ResolveSessionBinding(ctx, "sess-1") + gotAccount, gotRunner, _, err := s.ResolveSessionBinding(ctx, "sess-1") if err != nil { t.Fatalf("ResolveSessionBinding: %v", err) } @@ -110,7 +110,7 @@ func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { t.Fatalf("ResolveSessionBinding runner = %q, want runner-1", gotRunner) } - gotSession, gotRunner, err := s.SessionForAccount(ctx, agent.ID) + gotSession, gotRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("SessionForAccount: %v", err) } @@ -128,7 +128,7 @@ func TestSessionBindingLookupsFailClosed(t *testing.T) { ctx := context.Background() s := newTestStore(t) - account, runner, err := s.ResolveSessionBinding(ctx, "never-bound") + account, runner, _, err := s.ResolveSessionBinding(ctx, "never-bound") if !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionBinding(never-bound) err = %v, want errors.Is(_, ErrNotFound)", err) } @@ -140,7 +140,7 @@ func TestSessionBindingLookupsFailClosed(t *testing.T) { // started a session resolves nothing to dispatch to. owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - session, runner, err := s.SessionForAccount(ctx, agent.ID) + session, runner, _, err := s.SessionForAccount(ctx, agent.ID) if !errors.Is(err, ErrNotFound) { t.Fatalf("SessionForAccount(unbound agent) err = %v, want errors.Is(_, ErrNotFound)", err) } @@ -167,7 +167,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { mustBind(t, ctx, s, "sess-old", agent.ID, "runner-1") - displaced, err := s.RecordSessionBinding(ctx, "sess-new", agent.ID, "runner-2") + displaced, _, err := s.RecordSessionBinding(ctx, "sess-new", agent.ID, "runner-2") if err != nil { t.Fatalf("re-point onto sess-new: %v (a re-point must never be refused — promoteSession has nowhere to put an error)", err) } @@ -176,7 +176,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { } // The account now resolves to the NEW session (the session_id assignment). - gotSession, gotRunner, err := s.SessionForAccount(ctx, agent.ID) + gotSession, gotRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("SessionForAccount after the re-point: %v", err) } @@ -192,7 +192,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { // The displaced session id no longer resolves: it named the same row, which // now carries the new session. - if _, _, err := s.ResolveSessionBinding(ctx, "sess-old"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-old"); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionBinding(sess-old) err = %v, want ErrNotFound — the displaced session must stop resolving", err) } @@ -260,7 +260,7 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T mustBind(t, ctx, s, "sess-1", agentA.ID, "runner-1") - displaced, err := s.RecordSessionBinding(ctx, "sess-1", agentB.ID, "runner-1") + displaced, _, err := s.RecordSessionBinding(ctx, "sess-1", agentB.ID, "runner-1") sentinelIs(t, err, ErrConflict, "a session id already bound to a different agent") if displaced != "" { t.Fatalf("the refused bind returned displaced = %q, want \"\" — a failed write displaced nothing", displaced) @@ -268,14 +268,14 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T // The refused write changed nothing: the session still speaks for agent A, // and agent B still has no live session. - gotAccount, _, err := s.ResolveSessionBinding(ctx, "sess-1") + gotAccount, _, _, err := s.ResolveSessionBinding(ctx, "sess-1") if err != nil { t.Fatalf("ResolveSessionAccount after the refused bind: %v", err) } if gotAccount != agentA.ID { t.Fatalf("sess-1 resolves to %q, want the original owner %q", gotAccount, agentA.ID) } - if _, _, err := s.SessionForAccount(ctx, agentB.ID); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.SessionForAccount(ctx, agentB.ID); !errors.Is(err, ErrNotFound) { t.Fatalf("SessionForAccount(agent-b) err = %v, want ErrNotFound — the refused bind must not have landed", err) } } @@ -287,7 +287,7 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T func TestRecordSessionBindingUnknownAgentIsInvalidArgument(t *testing.T) { s := newTestStore(t) - _, err := s.RecordSessionBinding(context.Background(), "sess-1", "no-such-agent", "runner-1") + _, _, err := s.RecordSessionBinding(context.Background(), "sess-1", "no-such-agent", "runner-1") sentinelIs(t, err, ErrInvalidArgument, "binding for an unknown agent") } @@ -307,10 +307,10 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { t.Fatalf("DeleteSessionBinding: %v", err) } - if _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionAccount after delete err = %v, want ErrNotFound", err) } - if _, _, err := s.SessionForAccount(ctx, agent.ID); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.SessionForAccount(ctx, agent.ID); !errors.Is(err, ErrNotFound) { t.Fatalf("SessionForAccount after delete err = %v, want ErrNotFound", err) } @@ -330,6 +330,53 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { } } +// A release conditioned on a version removes only that row version: a re-bind of +// the same session id writes a new version, so the stale release leaves it bound. +func TestDeleteSessionBindingByVersionSkipsRebind(t *testing.T) { + ctx := context.Background() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + _, oldVersion, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") + if err != nil { + t.Fatalf("RecordSessionBinding: %v", err) + } + if _, _, read, err := s.ResolveSessionBinding(ctx, "sess-1"); err != nil || read != oldVersion { + t.Fatalf("ResolveSessionBinding version = (%q, %v), want (%q, nil)", read, err, oldVersion) + } + _, newVersion, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") + if err != nil { + t.Fatalf("re-bind RecordSessionBinding: %v", err) + } + if newVersion == oldVersion || oldVersion == "" { + t.Fatalf("re-bind kept version %q, want a new one", oldVersion) + } + + removed, err := s.DeleteSessionBindingVersion(ctx, "sess-1", oldVersion) + if err != nil || removed { + t.Fatalf("stale DeleteSessionBinding = (%v, %v), want (false, nil)", removed, err) + } + if got, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); err != nil || got != agent.ID { + t.Fatalf("binding after stale delete = (%q, %v), want (%q, nil)", got, err, agent.ID) + } + tenant := s.EffectiveTenant(ctx) + if events := computeEvents(t, s, tenant, agent.ID); len(events) != 1 || events[0].Kind != "start" { + t.Fatalf("usage events after stale delete = %+v, want only the open start", events) + } + removed, err = s.DeleteSessionBindingVersion(ctx, "sess-1", newVersion) + if err != nil || !removed { + t.Fatalf("current DeleteSessionBinding = (%v, %v), want (true, nil)", removed, err) + } + if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { + t.Fatalf("ResolveSessionBinding after delete err = %v, want ErrNotFound", err) + } + events := computeEvents(t, s, tenant, agent.ID) + if len(events) != 2 || events[0].Kind != "start" || events[1].Kind != "end" || events[0].IntervalID != events[1].IntervalID { + t.Fatalf("usage events after current delete = %+v, want one start and its end", events) + } +} + // TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding is the reconnect // sweep itself, and the test that would catch a `:exec` regression. The returned // rows are NOT diagnostic: each drives a presence DISCONNECTED edge for its @@ -390,14 +437,14 @@ func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { // The swept bindings are actually gone in both directions. for _, sessionID := range []string{"sess-a", "sess-b", "sess-z"} { - if _, _, err := s.ResolveSessionBinding(ctx, sessionID); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, sessionID); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionBinding(%s) after the sweep err = %v, want ErrNotFound", sessionID, err) } } // The other Runner's binding SURVIVES: a re-enroll must not drop live // sessions on a Runner that never reconnected. - gotAccount, _, err := s.ResolveSessionBinding(ctx, "sess-elsewhere") + gotAccount, _, _, err := s.ResolveSessionBinding(ctx, "sess-elsewhere") if err != nil { t.Fatalf("runner-2's binding did not survive the runner-1 sweep: %v", err) } @@ -470,7 +517,7 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { // Tenant B resolves A's session id: the row is not in B's view, so this must // fail closed rather than hand B tenant A's account. - gotAccount, _, err := s.ResolveSessionBinding(ctxB, "sess-a") + gotAccount, _, _, err := s.ResolveSessionBinding(ctxB, "sess-a") if !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B ResolveSessionBinding(sess-a) err = %v, want ErrNotFound — cross-tenant read leak", err) } @@ -479,7 +526,7 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { } // The reverse direction leaks nothing either. - if gotSession, _, err := s.SessionForAccount(ctxB, agentA.ID); !errors.Is(err, ErrNotFound) { + if gotSession, _, _, err := s.SessionForAccount(ctxB, agentA.ID); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B SessionForAccount(A's agent) = (%q, %v), want ErrNotFound — cross-tenant read leak", gotSession, err) } @@ -495,7 +542,7 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { // Control: tenant A still sees its OWN binding, proving the policy is not // simply hiding everything. - if gotAccount, _, err := s.ResolveSessionBinding(ctxA, "sess-a"); err != nil { + if gotAccount, _, _, err := s.ResolveSessionBinding(ctxA, "sess-a"); err != nil { t.Fatalf("tenant A cannot see its OWN binding — policy over-blocks: %v", err) } else if gotAccount != agentA.ID { t.Fatalf("tenant A's own binding resolves to %q, want %q", gotAccount, agentA.ID) @@ -554,14 +601,14 @@ func TestSessionBindingSameSessionIDInTwoTenantsCoexist(t *testing.T) { // Each tenant resolves its OWN account from the shared session id. A single // surviving row would make one of these two answer with the other tenant's // account — the relay resolving a foreign principal. - gotA, _, err := s.ResolveSessionBinding(ctxA, "sess-shared") + gotA, _, _, err := s.ResolveSessionBinding(ctxA, "sess-shared") if err != nil { t.Fatalf("tenant A ResolveSessionBinding(sess-shared) after B's bind: %v (B's write clobbered A's binding)", err) } if gotA != agentA.ID { t.Fatalf("tenant A resolves sess-shared to %q, want its own agent %q", gotA, agentA.ID) } - gotB, _, err := s.ResolveSessionBinding(ctxB, "sess-shared") + gotB, _, _, err := s.ResolveSessionBinding(ctxB, "sess-shared") if err != nil { t.Fatalf("tenant B ResolveSessionBinding(sess-shared): %v", err) } @@ -570,10 +617,10 @@ func TestSessionBindingSameSessionIDInTwoTenantsCoexist(t *testing.T) { } // And the reverse direction, per tenant. - if got, runner, err := s.SessionForAccount(ctxA, agentA.ID); err != nil || got != "sess-shared" || runner != "runner-a" { + if got, runner, _, err := s.SessionForAccount(ctxA, agentA.ID); err != nil || got != "sess-shared" || runner != "runner-a" { t.Fatalf("tenant A SessionForAccount = (%q, %q, %v), want (sess-shared, runner-a, nil)", got, runner, err) } - if got, runner, err := s.SessionForAccount(ctxB, agentB.ID); err != nil || got != "sess-shared" || runner != "runner-b" { + if got, runner, _, err := s.SessionForAccount(ctxB, agentB.ID); err != nil || got != "sess-shared" || runner != "runner-b" { t.Fatalf("tenant B SessionForAccount = (%q, %q, %v), want (sess-shared, runner-b, nil)", got, runner, err) } } @@ -608,10 +655,10 @@ func TestSessionBindingSameAccountIDInTwoTenantsCoexist(t *testing.T) { } // Both rows survive, one per tenant, each resolving its own session. - if got, runner, err := s.SessionForAccount(ctxA, shared.ID); err != nil || got != "sess-a" || runner != "runner-a" { + if got, runner, _, err := s.SessionForAccount(ctxA, shared.ID); err != nil || got != "sess-a" || runner != "runner-a" { t.Fatalf("tenant A SessionForAccount = (%q, %q, %v), want (sess-a, runner-a, nil) — B's write reached A's row", got, runner, err) } - if got, runner, err := s.SessionForAccount(ctxB, shared.ID); err != nil || got != "sess-b" || runner != "runner-b" { + if got, runner, _, err := s.SessionForAccount(ctxB, shared.ID); err != nil || got != "sess-b" || runner != "runner-b" { t.Fatalf("tenant B SessionForAccount = (%q, %q, %v), want (sess-b, runner-b, nil)", got, runner, err) } @@ -625,10 +672,10 @@ func TestSessionBindingSameAccountIDInTwoTenantsCoexist(t *testing.T) { } // Each tenant's session resolves only in its own tenant. - if _, _, err := s.ResolveSessionBinding(ctxA, "sess-b"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxA, "sess-b"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant A resolved B's session err = %v, want ErrNotFound", err) } - if _, _, err := s.ResolveSessionBinding(ctxB, "sess-a"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxB, "sess-a"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B resolved A's session err = %v, want ErrNotFound", err) } } @@ -663,7 +710,7 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { // Under the system role BOTH rows are visible, so the :one read is ambiguous. // It does NOT error — that is the hazard. It returns one of the two, and // which one is not something the caller can control or detect. - got, gotRunner, err := s.SessionForAccount(WithSystemRole(context.Background()), shared.ID) + got, gotRunner, _, err := s.SessionForAccount(WithSystemRole(context.Background()), shared.ID) if err != nil { t.Fatalf("SessionForAccount under the system role: %v — the current behaviour is a SILENT pick, not an error; if this now errors, PR3 changed the contract and this test must be updated deliberately", err) } @@ -674,10 +721,10 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { // The same read on the REQUEST path is exact in both tenants: the hazard is // the system role's missing scoping, NOT anything about the data. - if v, runner, err := s.SessionForAccount(ctxA, shared.ID); err != nil || v != "sess-a" || runner != "runner-a" { + if v, runner, _, err := s.SessionForAccount(ctxA, shared.ID); err != nil || v != "sess-a" || runner != "runner-a" { t.Fatalf("tenant A request-path SessionForAccount = (%q, %q, %v), want (sess-a, runner-a, nil)", v, runner, err) } - if v, runner, err := s.SessionForAccount(ctxB, shared.ID); err != nil || v != "sess-b" || runner != "runner-b" { + if v, runner, _, err := s.SessionForAccount(ctxB, shared.ID); err != nil || v != "sess-b" || runner != "runner-b" { t.Fatalf("tenant B request-path SessionForAccount = (%q, %q, %v), want (sess-b, runner-b, nil)", v, runner, err) } @@ -700,7 +747,7 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { // pooled connection it draws, so this asserts the DISJUNCTION rather than // pretending either branch is the contract. Both are defects; PR3 owns the // fix, and this test is what will fail when it lands one. - sysDisplaced, sysErr := s.RecordSessionBinding(WithSystemRole(context.Background()), "sess-sys", shared.ID, "runner-sys") + sysDisplaced, _, sysErr := s.RecordSessionBinding(WithSystemRole(context.Background()), "sess-sys", shared.ID, "runner-sys") switch { case sysErr != nil: t.Logf("system-role RecordSessionBinding failed (connection had no prior armed statement, so the GUC was undefined and the DEFAULT was NULL): %v", sysErr) @@ -716,10 +763,10 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { t.Fatalf("system-role write left %d rows stamped tenant_id = '', want 1 — the observed failure mode is an ORPHAN row, so if it is now stamped with a real tenant the write path grew scoping and this test must be updated deliberately", orphans) } // And it is invisible to every tenant: no policy matches ''. - if _, _, err := s.ResolveSessionBinding(ctxA, "sess-sys"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxA, "sess-sys"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant A sees the untenanted binding err = %v, want ErrNotFound", err) } - if _, _, err := s.ResolveSessionBinding(ctxB, "sess-sys"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxB, "sess-sys"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B sees the untenanted binding err = %v, want ErrNotFound", err) } } @@ -786,7 +833,7 @@ func TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced(t *testin results := make(chan outcome, 2) launch := func(sessionID, runnerID string) { go func() { - d, err := s.RecordSessionBinding(ctx, sessionID, agent.ID, runnerID) + d, _, err := s.RecordSessionBinding(ctx, sessionID, agent.ID, runnerID) results <- outcome{sessionID, d, err} }() } @@ -836,7 +883,7 @@ func TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced(t *testin if n := countBindings(t, ctx, s, agent.ID); n != 1 { t.Fatalf("bindings after two concurrent re-points = %d, want 1", n) } - live, liveRunner, err := s.SessionForAccount(ctx, agent.ID) + live, liveRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("SessionForAccount after the race: %v", err) } @@ -848,7 +895,7 @@ func TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced(t *testin // Both displaced sessions are gone, and between them they were reported to // exactly the two callers: nothing was destroyed unreported. for _, dead := range []string{"sess-A", winner} { - if _, _, err := s.ResolveSessionBinding(ctx, dead); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, dead); !errors.Is(err, ErrNotFound) { t.Fatalf("displaced session %q still resolves (err = %v), want ErrNotFound", dead, err) } } @@ -914,7 +961,7 @@ func TestRecordSessionBindingConcurrentFirstBindsReportTheDestroyedSession(t *te {"sess-2", "runner-2"}, } { go func() { - d, err := s.RecordSessionBinding(ctx, b.session, agent.ID, b.runner) + d, _, err := s.RecordSessionBinding(ctx, b.session, agent.ID, b.runner) results <- outcome{b.session, d, err} }() } @@ -957,7 +1004,7 @@ func TestRecordSessionBindingConcurrentFirstBindsReportTheDestroyedSession(t *te t.Fatalf("bindings after two concurrent first binds = %d, want 1 — the PK must fold them into one row", n) } // The winner's session is gone, and it WAS reported, so it can be reaped. - if _, _, err := s.ResolveSessionBinding(ctx, winner); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, winner); !errors.Is(err, ErrNotFound) { t.Fatalf("the overwritten session %q still resolves (err = %v), want ErrNotFound", winner, err) } } @@ -1046,7 +1093,7 @@ func TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep(t *t go func() { ready.Done() <-start - displaced, err := s.RecordSessionBinding(ctx, fresh, agent.ID, "runner-2") + displaced, _, err := s.RecordSessionBinding(ctx, fresh, agent.ID, "runner-2") binds <- bindResult{displaced, err} }() go func() { @@ -1098,7 +1145,7 @@ func TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep(t *t if n := countBindings(t, ctx, s, agent.ID); n != 1 { t.Fatalf("iteration %d: bindings after the race = %d, want 1 — the bind must land whether or not the sweep removed the prior row", i, n) } - live, liveRunner, err := s.SessionForAccount(ctx, agent.ID) + live, liveRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("iteration %d: SessionForAccount after the race: %v", i, err) } @@ -1107,7 +1154,7 @@ func TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep(t *t } // And the displaced session is gone in every ordering, which is what // makes "reported to nobody" a genuine strand rather than a deferral. - if _, _, err := s.ResolveSessionBinding(ctx, old); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, old); !errors.Is(err, ErrNotFound) { t.Fatalf("iteration %d: displaced session %q still resolves (err = %v), want ErrNotFound", i, old, err) } } diff --git a/go/server/lifecycle_wake_pgtest_test.go b/go/server/lifecycle_wake_pgtest_test.go index fe8511882..db1b20da7 100644 --- a/go/server/lifecycle_wake_pgtest_test.go +++ b/go/server/lifecycle_wake_pgtest_test.go @@ -134,11 +134,11 @@ func TestWakeAgentStaleBindingRowStillResumes(t *testing.T) { if err := f.store.AppendTranscriptEntry(ctx, logical, 1, true, `{"header":true}`, "k1"); err != nil { t.Fatalf("append checkpoint: %v", err) } - if _, err := f.store.RecordSessionBinding(ctx, "sess-dead", f.agentID, fakeRunnerID); err != nil { + if _, _, err := f.store.RecordSessionBinding(ctx, "sess-dead", f.agentID, fakeRunnerID); err != nil { t.Fatalf("RecordSessionBinding: %v", err) } // Checked on the store, not the hub: a hub read-through would warm the cache. - if got, runnerID, err := f.store.SessionForAccount(ctx, f.agentID); err != nil || got != "sess-dead" || runnerID != fakeRunnerID { + if got, runnerID, _, err := f.store.SessionForAccount(ctx, f.agentID); err != nil || got != "sess-dead" || runnerID != fakeRunnerID { t.Fatalf("precondition: durable row = (%q, %q, %v), want (sess-dead, %s, nil)", got, runnerID, err, fakeRunnerID) } f.runner.forget()