fix(runnerhub): reap a Runner's session bindings in every tenant on enroll (RIG-4296)

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
 	}

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

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) {

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