fix(telemetry): resolve shutdown marker after async-op drain (Spec 080, review round 3)

Author: Dumbris

internal/runtime/prechurn_lifecycle_test.go

@@ -138,3 +138,42 @@ func TestRuntimeCloseCleanupBranchStillResolvesAndClosesStorage(t *testing.T) {
 		t.Fatalf("close 2: %v", err)
 	}
 }
+
+// TestRuntimeCloseAfterExternalStopAsyncStillResolvesClean guards the split
+// Close sequence (Spec 080 FR-010, review round 3): Close now runs
+// storage.StopAsync (stop + drain queued async DB ops) BEFORE resolving the
+// shutdown marker, then closes the DB — so the marker is truly the last DB
+// write. Driving StopAsync externally first exercises the double-stop path
+// (StopAsync inside Close, then again inside storageManager.Close): it must
+// not panic, and the marker must still resolve to clean.
+func TestRuntimeCloseAfterExternalStopAsyncStillResolvesClean(t *testing.T) {
+	t.Setenv("MCPPROXY_LAUNCHED_BY", "")
+	dataDir := t.TempDir()
+
+	rt, _ := previousShutdownVia(t, dataDir)
+	db := rt.StorageManager().GetDB()
+	if db == nil {
+		t.Fatal("precondition: storage DB must be open")
+	}
+
+	// Worst-case double stop: the async manager is already stopped and
+	// drained before Close runs the same sequence again.
+	rt.StorageManager().StopAsync()
+	if err := rt.Close(); err != nil {
+		t.Fatalf("close after external StopAsync: %v", err)
+	}
+
+	// Storage fully closed…
+	if err := db.View(func(*bbolt.Tx) error { return nil }); !errors.Is(err, berrors.ErrDatabaseNotOpen) {
+		t.Fatalf("expected storage closed after Close, View err = %v", err)
+	}
+
+	// …and the marker resolved: the next instance reports a clean shutdown.
+	rt2, prev := previousShutdownVia(t, dataDir)
+	if prev != "clean" {
+		t.Fatalf("after Close with external StopAsync: expected clean, got %q", prev)
+	}
+	if err := rt2.Close(); err != nil {
+		t.Fatalf("close 2: %v", err)
+	}
+}

internal/runtime/runtime.go

@@ -654,27 +654,36 @@ func (r *Runtime) Close() error {
 	}
 
 	// Spec 080 (US3, FR-010): resolve the shutdown marker to "clean" at the
-	// LAST point the DB is still open — after cache/index cleanup, immediately
-	// before storage closes (cache.Close only stops a goroutine and index.Close
-	// only touches Bleve; neither closes the BBolt DB). Every branch above —
-	// including the container-cleanup verification, whose early exits live
-	// inside verifyContainerCleanup — reaches this point, so a graceful Close
-	// always resolves; conversely a hang/SIGKILL/panic ANYWHERE earlier in
-	// shutdown leaves the marker armed and the next instance honestly reports
-	// previous_shutdown="crash". Idempotent on double Close; on an
-	// already-closed DB the write fails harmlessly (logged at debug) and the
-	// marker is left untouched.
-	if r.prechurnStore != nil && r.storageManager != nil {
-		if db := r.storageManager.GetDB(); db != nil {
-			if err := r.prechurnStore.ResolveCleanShutdown(db); err != nil {
-				if r.logger != nil {
-					r.logger.Debug("Failed to resolve shutdown marker to clean", zap.Error(err))
+	// LAST point the DB is still open — i.e. after the async storage manager
+	// has stopped AND drained its queue (those queued operations perform BBolt
+	// writes; StopAsync below runs that drain), immediately before the BBolt
+	// handle closes. Every branch above — including the container-cleanup
+	// verification, whose early exits live inside verifyContainerCleanup —
+	// reaches this point, so a graceful Close always resolves; conversely a
+	// hang/SIGKILL/panic ANYWHERE earlier in shutdown — including mid-drain —
+	// leaves the marker armed and the next instance honestly reports
+	// previous_shutdown="crash". Idempotent on double Close (StopAsync no-ops,
+	// and storageManager.Close's internal async stop no-ops after StopAsync);
+	// on an already-closed DB the marker write fails harmlessly (logged at
+	// debug) and the marker is left untouched.
+	if r.storageManager != nil {
+		// (1) Stop + drain queued async DB operations — the last DB writes
+		// other than the marker resolve itself.
+		r.storageManager.StopAsync()
+
+		// (2) Resolve the marker to clean, now that no other DB work remains.
+		if r.prechurnStore != nil {
+			if db := r.storageManager.GetDB(); db != nil {
+				if err := r.prechurnStore.ResolveCleanShutdown(db); err != nil {
+					if r.logger != nil {
+						r.logger.Debug("Failed to resolve shutdown marker to clean", zap.Error(err))
+					}
 				}
 			}
 		}
-	}
 
-	if r.storageManager != nil {
+		// (3) Close the BBolt handle; the async stop inside Close is a no-op,
+		// so no DB work intervenes between the marker resolve and db.Close.
 		if err := r.storageManager.Close(); err != nil {
 			errs = append(errs, fmt.Errorf("close storage manager: %w", err))
 		}

internal/storage/async_ops.go

@@ -30,6 +30,7 @@ type AsyncManager struct {
 	opQueue chan Operation
 	ctx     context.Context
 	cancel  context.CancelFunc
+	stateMu sync.Mutex // guards started; makes Start/Stop idempotent
 	started bool
 	wg      sync.WaitGroup
 }
@@ -48,6 +49,9 @@ func NewAsyncManager(db *BoltDB, logger *zap.SugaredLogger) *AsyncManager {
 
 // Start begins processing storage operations in a dedicated goroutine
 func (am *AsyncManager) Start() {
+	am.stateMu.Lock()
+	defer am.stateMu.Unlock()
+
 	if am.started {
 		return
 	}
@@ -56,13 +60,18 @@ func (am *AsyncManager) Start() {
 	go am.processOperations()
 }
 
-// Stop gracefully shuts down the async manager
+// Stop gracefully shuts down the async manager, draining any queued
+// operations. Idempotent: subsequent calls are no-ops (no double-cancel,
+// no DB work), so Manager.StopAsync followed by Manager.Close is safe.
 func (am *AsyncManager) Stop() {
+	am.stateMu.Lock()
+	defer am.stateMu.Unlock()
+
 	if !am.started {
 		return
 	}
 	am.cancel()
-	am.wg.Wait() // Wait for the goroutine to finish
+	am.wg.Wait() // Wait for the goroutine to finish (drains the queue)
 	am.started = false
 }
 

internal/storage/async_ops_test.go

@@ -399,3 +399,56 @@ func TestSaveServerSyncFieldCoverage(t *testing.T) {
 
 	t.Logf("ServerConfig has %d fields, all mapped to UpstreamRecord", serverConfigType.NumField())
 }
+
+// TestManagerStopAsyncDrainsThenCloseIsIdempotent guards the split shutdown
+// sequence behind Spec 080 FR-010: StopAsync must stop the async manager AND
+// drain queued operations to the DB (so a caller can perform a final write —
+// the telemetry shutdown marker — strictly after all async DB work), and the
+// subsequent Close must not re-run the drain, double-cancel, or panic. A
+// StopAsync after Close must also be a harmless no-op.
+func TestManagerStopAsyncDrainsThenCloseIsIdempotent(t *testing.T) {
+	logger := zaptest.NewLogger(t).Sugar()
+	manager, err := NewManager(t.TempDir(), logger)
+	if err != nil {
+		t.Fatalf("NewManager: %v", err)
+	}
+
+	// Seed a record, then queue an async write against it.
+	if err := manager.SaveUpstreamServer(&config.ServerConfig{
+		Name:     "drain-test",
+		Protocol: "stdio",
+		Command:  "true",
+		Enabled:  false,
+	}); err != nil {
+		t.Fatalf("SaveUpstreamServer: %v", err)
+	}
+	manager.asyncMgr.EnableServerAsync("drain-test", true)
+
+	// StopAsync must flush the queued write before returning.
+	manager.StopAsync()
+	got, err := manager.GetUpstreamServer("drain-test")
+	if err != nil {
+		t.Fatalf("GetUpstreamServer after StopAsync: %v", err)
+	}
+	if !got.Enabled {
+		t.Fatal("queued async write not drained by StopAsync")
+	}
+
+	// The DB is still open between StopAsync and Close — this is the window
+	// where the shutdown marker is resolved.
+	if manager.GetDB() == nil {
+		t.Fatal("DB must remain open after StopAsync")
+	}
+
+	// Double StopAsync, then Close (whose internal async stop must no-op).
+	manager.StopAsync()
+	if err := manager.Close(); err != nil {
+		t.Fatalf("Close after StopAsync: %v", err)
+	}
+
+	// StopAsync and Close after Close must not panic.
+	manager.StopAsync()
+	if err := manager.Close(); err != nil {
+		t.Fatalf("second Close: %v", err)
+	}
+}

internal/storage/manager.go

@@ -43,7 +43,26 @@ func NewManager(dataDir string, logger *zap.SugaredLogger) (*Manager, error) {
 	}, nil
 }
 
-// Close closes the storage manager
+// StopAsync stops the async operation manager, draining any queued
+// operations to the database. It is idempotent: a second call (including
+// the one inside Close) is a no-op, so callers may StopAsync then Close.
+//
+// Use this when a final DB write must happen strictly AFTER all queued
+// async operations have flushed but BEFORE the DB handle closes — e.g.
+// the telemetry shutdown marker (Spec 080 FR-010): StopAsync, write the
+// marker, then Close, which then closes the DB with no intervening work.
+func (m *Manager) StopAsync() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if m.asyncMgr != nil {
+		m.asyncMgr.Stop()
+	}
+}
+
+// Close closes the storage manager: it stops the async manager (draining
+// queued operations to the DB — a no-op if StopAsync already ran), then
+// closes the underlying BBolt database.
 func (m *Manager) Close() error {
 	m.mu.Lock()
 	defer m.mu.Unlock()