fix(telemetry): close startup-vs-shutdown races in marker write barriers (Spec 080, review round 5)

Author: Dumbris

internal/runtime/activity_service.go

@@ -48,10 +48,19 @@ type ActivityService struct {
 	// background loops (retention, usage flush) and the per-event async
 	// detection goroutines, all of which write to BBolt. startMu/started make
 	// Stop return immediately when Start never ran (done would never close).
+	// stopped is the terminal state (Spec 080, review round 5): production
+	// launches Start via `go` (lifecycle.go), so a fast shutdown can run Stop
+	// BEFORE the Start goroutine is scheduled — Stop marks stopped under
+	// startMu and a later Start becomes a no-op instead of launching BBolt
+	// writers after the shutdown-marker path began. Start's registration
+	// (subscribe + every workersWG.Add) happens entirely under startMu, so a
+	// Stop that loses the race blocks until registration is complete and its
+	// Wait cannot miss a late worker.
 	done      chan struct{}
 	workersWG sync.WaitGroup
 	startMu   sync.Mutex
 	started   bool
+	stopped   bool
 
 	// Retention configuration
 	maxAge        time.Duration
@@ -127,20 +136,29 @@ func (s *ActivityService) SetRetentionConfig(maxAge time.Duration, maxRecords in
 // Start begins listening for activity events and persisting them.
 // It should be called as a goroutine: go svc.Start(ctx, runtime)
 func (s *ActivityService) Start(ctx context.Context, rt *Runtime) {
-	// Mark started so Stop knows the done channel WILL close; refuse a second
-	// Start (the done/WaitGroup bookkeeping is single-shot).
+	// Registration runs entirely under startMu (Spec 080 FR-010, review round
+	// 5). If Stop already ran (fast shutdown beat this goroutine — production
+	// launches Start via `go`), the service is terminally stopped: return
+	// without subscribing or launching any BBolt-writing worker. Otherwise
+	// mark started so Stop knows the done channel WILL close, and refuse a
+	// second Start (the done/WaitGroup bookkeeping is single-shot). Holding
+	// startMu through every workersWG.Add below means a concurrent Stop
+	// blocks until registration is complete — its Wait cannot miss a worker.
 	s.startMu.Lock()
+	if s.stopped {
+		s.startMu.Unlock()
+		s.logger.Debug("Activity service Start called after Stop; not starting")
+		return
+	}
 	if s.started {
 		s.startMu.Unlock()
 		s.logger.Warn("Activity service Start called twice; ignoring")
 		return
 	}
 	s.started = true
-	s.startMu.Unlock()
 
 	// Subscribe to runtime events
 	eventCh := rt.SubscribeEvents()
-	defer rt.UnsubscribeEvents(eventCh)
 
 	// Start retention loop in a separate goroutine. Tracked in workersWG: it
 	// prunes activity records (BBolt writes), so Stop must await it.
@@ -158,6 +176,9 @@ func (s *ActivityService) Start(ctx context.Context, rt *Runtime) {
 		defer s.workersWG.Done()
 		s.runUsageFlushLoop(ctx)
 	}()
+	s.startMu.Unlock()
+
+	defer rt.UnsubscribeEvents(eventCh)
 
 	s.logger.Info("Activity service started")
 
@@ -255,8 +276,16 @@ func (s *ActivityService) runRetentionCleanup() {
 // flush runs on ctx.Done inside the event loop, and Stop waits for it, so the
 // flush is captured before the marker resolves. Idempotent, and returns
 // immediately when Start never ran.
+//
+// Stop is also terminal (Spec 080, review round 5): it marks stopped under
+// startMu, so a Start that has not yet registered (production starts the
+// service via `go` in lifecycle.go) becomes a no-op instead of launching
+// retention/usage/persist loops after the shutdown-marker path began. If
+// Start is mid-registration, acquiring startMu here blocks until every
+// workersWG.Add has happened, so the Wait below cannot miss a worker.
 func (s *ActivityService) Stop() {
 	s.startMu.Lock()
+	s.stopped = true
 	started := s.started
 	s.startMu.Unlock()
 	if !started {

internal/runtime/activity_service_test.go

@@ -1,6 +1,7 @@
 package runtime
 
 import (
+	"context"
 	"os"
 	"testing"
 	"time"
@@ -919,3 +920,42 @@ func TestHandleToolCallCompleted_ByteCapture(t *testing.T) {
 	assert.Equal(t, 1500, record.RequestBytes, "RequestBytes must reflect pre-truncation request size")
 	assert.Equal(t, 98304, record.ResponseBytes, "ResponseBytes must reflect pre-truncation response size")
 }
+
+// TestActivityServiceStartAfterStopIsNoOp (Spec 080 FR-010, review round 5):
+// production launches Start via `go r.activityService.Start(...)` in
+// lifecycle.go, so a fast shutdown can run Stop BEFORE the Start goroutine is
+// ever scheduled. Stop must leave a terminal stopped state that turns the
+// late Start into a no-op — no retention/usage/persist loops launched, no
+// storage access. Storage and runtime are nil on purpose: any registration
+// step (SubscribeEvents, initUsageFromStorage, the worker loops) would panic,
+// so a regression fails loudly.
+func TestActivityServiceStartAfterStopIsNoOp(t *testing.T) {
+	svc := NewActivityService(nil, zap.NewNop())
+
+	// Fast shutdown wins the race: Stop runs before Start ever does. It must
+	// return immediately (Start never ran, done never closes).
+	svc.Stop()
+
+	// The delayed Start must return immediately without registering anything.
+	// A non-no-op Start would either panic (nil rt/storage) or block forever
+	// in the event loop (ctx is never cancelled) and trip the timeout below.
+	returned := make(chan struct{})
+	go func() {
+		defer close(returned)
+		svc.Start(context.Background(), nil)
+	}()
+	select {
+	case <-returned:
+	case <-time.After(2 * time.Second):
+		t.Fatal("Start after Stop did not return; it must be a no-op")
+	}
+
+	svc.startMu.Lock()
+	started := svc.started
+	svc.startMu.Unlock()
+	assert.False(t, started, "Start after Stop must not mark the service started")
+
+	// Stop stays idempotent after the no-op Start.
+	svc.Stop()
+	svc.Stop()
+}

internal/runtime/supervisor/supervisor.go

@@ -240,9 +240,25 @@ func (s *Supervisor) Start() {
 	s.wg.Add(1)
 	go s.exemptionCleanupLoop()
 
-	// Phase 7.1: Trigger initial reconciliation to populate StateView
+	// Phase 7.1: Trigger initial reconciliation to populate StateView.
+	// Registered in s.wg with a ctx-aware timer (Spec 080 FR-010/FR-011,
+	// review round 5): Stop() cancels s.ctx and then waits on s.wg, so it
+	// either cancels this goroutine inside the 500ms warm-up window or waits
+	// for the reconcile to finish. A bare time.Sleep goroutine here could wake
+	// AFTER Stop() returned and write last_error_code/diagnostics via
+	// reconcile()/updateStateView()/notifyErrorCode() — after the clean-
+	// shutdown marker resolved, or against a closed DB.
+	s.wg.Add(1)
 	go func() {
-		time.Sleep(500 * time.Millisecond) // Give servers time to connect
+		defer s.wg.Done()
+		timer := time.NewTimer(500 * time.Millisecond) // Give servers time to connect
+		defer timer.Stop()
+		select {
+		case <-s.ctx.Done():
+			s.logger.Debug("Supervisor stopping before initial reconciliation")
+			return
+		case <-timer.C:
+		}
 		currentConfig := s.configSvc.Current()
 		if err := s.reconcile(currentConfig); err != nil {
 			s.logger.Error("Initial reconciliation failed", zap.Error(err))

internal/runtime/supervisor/supervisor_test.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -896,3 +897,113 @@ func TestSupervisor_ErrorCodeNotifierSynchronous(t *testing.T) {
 		t.Fatalf("notifier received a non-MCPX code: %q", got)
 	}
 }
+
+// barrierUpstreamAdapter wraps MockUpstreamAdapter and counts any upstream
+// call that happens after the barrier is armed. Used to prove Stop() is a
+// write barrier for the delayed initial-reconciliation goroutine (Spec 080).
+type barrierUpstreamAdapter struct {
+	*MockUpstreamAdapter
+	barrier    atomic.Bool
+	violations atomic.Int32
+}
+
+func (b *barrierUpstreamAdapter) check() {
+	if b.barrier.Load() {
+		b.violations.Add(1)
+	}
+}
+
+func (b *barrierUpstreamAdapter) AddServer(name string, cfg *config.ServerConfig) error {
+	b.check()
+	return b.MockUpstreamAdapter.AddServer(name, cfg)
+}
+
+func (b *barrierUpstreamAdapter) RemoveServer(name string) error {
+	b.check()
+	return b.MockUpstreamAdapter.RemoveServer(name)
+}
+
+func (b *barrierUpstreamAdapter) ConnectServer(ctx context.Context, name string) error {
+	b.check()
+	return b.MockUpstreamAdapter.ConnectServer(ctx, name)
+}
+
+func (b *barrierUpstreamAdapter) DisconnectServer(name string) error {
+	b.check()
+	return b.MockUpstreamAdapter.DisconnectServer(name)
+}
+
+func (b *barrierUpstreamAdapter) ConnectAll(ctx context.Context) error {
+	b.check()
+	return b.MockUpstreamAdapter.ConnectAll(ctx)
+}
+
+func (b *barrierUpstreamAdapter) GetServerState(name string) (*ServerState, error) {
+	b.check()
+	return b.MockUpstreamAdapter.GetServerState(name)
+}
+
+func (b *barrierUpstreamAdapter) GetAllStates() map[string]*ServerState {
+	b.check()
+	return b.MockUpstreamAdapter.GetAllStates()
+}
+
+func (b *barrierUpstreamAdapter) IsUserLoggedOut(name string) bool {
+	b.check()
+	return b.MockUpstreamAdapter.IsUserLoggedOut(name)
+}
+
+// TestSupervisor_StopBeforeInitialReconcileIsBarrier (Spec 080 FR-010/FR-011,
+// review round 5): Start() arms a delayed (500ms) initial reconciliation.
+// Stop() must be a barrier for it — the goroutine is registered in s.wg and
+// waits on a ctx-aware timer, so Stop() (cancel + wg.Wait) deterministically
+// either cancels it inside the window or waits for the reconcile to finish.
+// Before the fix it was a bare `go func() { time.Sleep(500ms); reconcile() }`:
+// Stop() returned with nothing to wait for, Runtime.Close resolved the clean-
+// shutdown marker, and the goroutine then woke and wrote diagnostics/state
+// (reconcile -> updateStateView -> notifyErrorCode) after the marker — or
+// against a closed DB.
+func TestSupervisor_StopBeforeInitialReconcileIsBarrier(t *testing.T) {
+	cfg := &config.Config{
+		Listen: "127.0.0.1:8080",
+		Servers: []*config.ServerConfig{
+			{Name: "late-server", Enabled: true},
+		},
+	}
+
+	configSvc := configsvc.NewService(cfg, "/tmp/config.json", zap.NewNop())
+	defer configSvc.Close()
+
+	upstream := &barrierUpstreamAdapter{MockUpstreamAdapter: NewMockUpstreamAdapter()}
+	sup := New(configSvc, upstream, zap.NewNop())
+
+	var stopReturned atomic.Bool
+	var notifierAfterStop atomic.Int32
+	sup.SetErrorCodeNotifier(func(string) {
+		if stopReturned.Load() {
+			notifierAfterStop.Add(1)
+		}
+	})
+
+	sup.Start()
+	sup.Stop() // well inside the 500ms initial-reconcile delay
+
+	// Everything the supervisor owns (reconciliation loop, event forwarding,
+	// exemption cleanup, the initial-reconcile goroutine, in-flight actions)
+	// is joined by Stop() via s.wg / drainActions, so from this point on no
+	// upstream call and no notifier call may ever happen again.
+	upstream.barrier.Store(true)
+	stopReturned.Store(true)
+
+	// Regression net: the pre-fix bare goroutine would wake ~500ms after
+	// Start() and run reconcile() against the stopped supervisor. Wait out the
+	// full window plus slack, then assert the barrier held. With the fix this
+	// sleep is pure idle time — the wg-joined goroutine already exited before
+	// Stop() returned, via the s.ctx.Done() arm of its select.
+	time.Sleep(700 * time.Millisecond)
+
+	require.Zero(t, upstream.violations.Load(),
+		"upstream adapter was called after Supervisor.Stop() returned")
+	require.Zero(t, notifierAfterStop.Load(),
+		"error-code notifier fired after Supervisor.Stop() returned")
+}