diff --git a/go/internal/runnerhub/binding_cache_test.go b/go/internal/runnerhub/binding_cache_test.go index fec4fb581..0feecff98 100644 --- a/go/internal/runnerhub/binding_cache_test.go +++ b/go/internal/runnerhub/binding_cache_test.go @@ -58,6 +58,7 @@ type fakeBindingStore struct { resolveErr error reverseErr error deleteForRunnerErr error + tenantErr error } func newFakeBindingStore() *fakeBindingStore { @@ -156,6 +157,9 @@ const testRunnerID = "runner-1" func (f *fakeBindingStore) SessionBindingTenant(_ context.Context, sessionID, runnerID string) (store.TenantID, error) { f.mu.Lock() defer f.mu.Unlock() + if f.tenantErr != nil { + return "", f.tenantErr + } if b, ok := f.bindings[sessionID]; ok && b.RunnerID == runnerID { return f.tenant, nil } diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index fc2ada08b..2b8bee947 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -1073,8 +1073,12 @@ func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier h.archiveEnded(store.WithTenant(ctx, b.TenantID), b.SessionID) } } else { + // A failed sweep leaves the rows: resolve each session's tenant from them when + // the store still answers, else archive unscoped as before. An ambiguous id is skipped. for _, sessionID := range reapedSessions { - h.archiveEnded(ctx, sessionID) + if sctx, ok := h.runnerSessionCtx(ctx, id, sessionID); ok { + h.archiveEnded(sctx, sessionID) + } } } return reattached diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 806e1c540..63030dace 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -4,7 +4,9 @@ package runnerhub import ( "context" + "errors" "testing" + "testing/synctest" "time" compassv1 "github.com/RigelBuild/compass/go/gen/compass/v1" @@ -62,6 +64,77 @@ func recvEnded(t *testing.T, c chanEndSink) string { } } +type endedArchive struct { + tenant store.TenantID + sessionID string +} + +type tenantRecordingEndSink struct{ ended chan endedArchive } + +func (s tenantRecordingEndSink) OnSessionEnded(ctx context.Context, sessionID string) { + tenant, _ := store.TenantFromContext(ctx) + s.ended <- endedArchive{tenant: tenant, sessionID: sessionID} +} + +// When the durable sweep faults, the RAM-snapshot fallback still archives each +// session under its own tenant, read from the binding rows the fault left behind. +func TestEnrollReapFaultArchivesUnderSessionTenant(t *testing.T) { + ctx := t.Context() + hub := newHubOnly() + bindings := newFakeBindingStore() + hub.SetSessionBindingStore(bindings) + ended := tenantRecordingEndSink{ended: make(chan endedArchive, 1)} + hub.SetSessionEndSink(ended) + subj := runnerSubject() + tier, egress := compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED + hub.enroll(ctx, "runner-1", subj, tier, egress) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + bindings.mu.Lock() + bindings.deleteForRunnerErr = errors.New("sweep fault") + bindings.mu.Unlock() + hub.enroll(ctx, "runner-1", subj, tier, egress) + select { + case got := <-ended.ended: + if got != (endedArchive{tenant: bindings.tenant, sessionID: "sess-1"}) { + t.Fatalf("archived %+v, want sess-1 under %q", got, bindings.tenant) + } + case <-time.After(10 * time.Second): + t.Fatal("no session-end archive within 10s") + } +} + +// A session id bound in several tenants resolves to no single tenant, so the +// fallback skips its archive rather than reading an arbitrary tenant's copy. +func TestEnrollReapFaultSkipsAmbiguousSessionArchive(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := t.Context() + hub := newHubOnly() + bindings := newFakeBindingStore() + hub.SetSessionBindingStore(bindings) + ended := tenantRecordingEndSink{ended: make(chan endedArchive, 1)} + hub.SetSessionEndSink(ended) + subj := runnerSubject() + tier, egress := compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED + hub.enroll(ctx, "runner-1", subj, tier, egress) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + bindings.mu.Lock() + bindings.deleteForRunnerErr = errors.New("sweep fault") + bindings.tenantErr = store.ErrConflict + bindings.mu.Unlock() + hub.enroll(ctx, "runner-1", subj, tier, egress) + synctest.Wait() + select { + case got := <-ended.ended: + t.Fatalf("archived %+v for an ambiguous session id, want no archive", got) + default: + } + }) +} + // A session that ends without Stop is archived: on a lost-session drop by its owning // Runner, and on each binding the re-enroll reap removes. A foreign refusal is not. func TestSessionEndedWithoutStopIsArchived(t *testing.T) { diff --git a/go/internal/runnerhub/runner_tenant_pgtest_test.go b/go/internal/runnerhub/runner_tenant_pgtest_test.go index e326caa65..e36b14455 100644 --- a/go/internal/runnerhub/runner_tenant_pgtest_test.go +++ b/go/internal/runnerhub/runner_tenant_pgtest_test.go @@ -43,13 +43,16 @@ func TestDropLostSessionScopesToTheSessionTenant(t *testing.T) { } // An enrolling Runner spans tenants, so its system-role sweep must reap tenant B's -// binding and record its end event with tenant B's id. +// binding, record its end event with tenant B's id, and archive it under tenant B. func TestEnrollReapsSessionBindingAcrossTenants(t *testing.T) { ctx := context.Background() st, agent, ctxB := openTenantBSession(t, ctx) + tenantB, _ := store.TenantFromContext(ctxB) hub := newHubOnly() hub.SetSessionBindingStore(st) + ended := tenantRecordingEndSink{ended: make(chan endedArchive, 1)} + hub.SetSessionEndSink(ended) hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) var remaining, ends int @@ -65,6 +68,14 @@ func TestEnrollReapsSessionBindingAcrossTenants(t *testing.T) { if remaining != 0 || ends != 1 { t.Fatalf("tenant B binding remaining = %d, end events = %d; want reaped binding and one end", remaining, ends) } + select { + case got := <-ended.ended: + if got != (endedArchive{tenant: tenantB, sessionID: "sess-b"}) { + t.Fatalf("archived %+v, want sess-b under %q", got, tenantB) + } + case <-time.After(10 * time.Second): + t.Fatal("no session-end archive within 10s") + } } // openTenantBSession opens a store with a non-bootstrap tenant B whose agent holds