Protect HealthCheckLoop from same nil-pointer crash, fix test nits

- Add recoverWorkerPanic to the doHealthCheck call in HealthCheckLoop.
  Same race: w.client.Close() from a concurrent crash/retire nils out
  FlightServiceClient while the health check goroutine calls DoAction.

- Replace nil contexts with context.Background() in tests.

- Add TestDestroySessionAfterOnWorkerCrash to verify the exact
  production sequence: OnWorkerCrash cleans up, then the deferred
  DestroySession is a safe no-op.

- Add comment explaining intentional double-close of TCP connection
  (OnWorkerCrash + handleConnection defer).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: fuziontech

controlplane/session_mgr.go

@@ -171,7 +171,10 @@ func (sm *SessionManager) OnWorkerCrash(workerID int, errorFn func(pid int32)) {
 			}
 			// Close the TCP connection to unblock the message loop's read.
 			// This causes the session goroutine to exit instead of looping
-			// with ErrWorkerDead on every query.
+			// with ErrWorkerDead on every query. The deferred close in
+			// handleConnection will also call Close() on the same conn;
+			// that's harmless (net.Conn.Close on a closed socket returns
+			// an error which is discarded).
 			if session.connCloser != nil {
 				_ = session.connCloser.Close()
 			}

controlplane/session_mgr_test.go

@@ -211,3 +211,43 @@ func TestRecoverWorkerPanic_RuntimeErrorRePanics(t *testing.T) {
 
 	t.Fatal("should not reach here")
 }
+
+func TestDestroySessionAfterOnWorkerCrash(t *testing.T) {
+	// Verify that DestroySession is a safe no-op when OnWorkerCrash already
+	// cleaned up the session. This is the exact production sequence:
+	// OnWorkerCrash runs from the health check, then the deferred
+	// DestroySession runs when handleConnection returns.
+	pool := &FlightWorkerPool{
+		workers: make(map[int]*ManagedWorker),
+	}
+	sm := NewSessionManager(pool, nil)
+
+	conn := &mockCloser{}
+	executor := &server.FlightExecutor{}
+	pid := int32(1010)
+
+	sm.mu.Lock()
+	sm.sessions[pid] = &ManagedSession{
+		PID:        pid,
+		WorkerID:   9,
+		Executor:   executor,
+		connCloser: conn,
+	}
+	sm.byWorker[9] = []int32{pid}
+	sm.mu.Unlock()
+
+	// Simulate crash cleanup
+	sm.OnWorkerCrash(9, func(pid int32) {})
+
+	if sm.SessionCount() != 0 {
+		t.Fatal("expected 0 sessions after OnWorkerCrash")
+	}
+
+	// Now DestroySession runs (from deferred call in handleConnection).
+	// Should be a no-op — no panic, no double-close of worker resources.
+	sm.DestroySession(pid)
+
+	if sm.SessionCount() != 0 {
+		t.Fatal("expected 0 sessions after DestroySession")
+	}
+}

controlplane/worker_mgr.go

@@ -484,10 +484,18 @@ func (p *FlightWorkerPool) HealthCheckLoop(ctx context.Context, interval time.Du
 						}
 						_ = os.Remove(w.socketPath)
 					default:
-						// Worker is alive, do a health check
-						hctx, cancel := context.WithTimeout(ctx, 3*time.Second)
-						err := doHealthCheck(hctx, w.client)
-						cancel()
+						// Worker is alive, do a health check.
+						// Recover nil-pointer panics: w.client.Close() (from a
+						// concurrent crash/retire) nils out FlightServiceClient,
+						// racing with the DoAction call inside doHealthCheck.
+						var healthErr error
+						func() {
+							defer recoverWorkerPanic(&healthErr)
+							hctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+							healthErr = doHealthCheck(hctx, w.client)
+							cancel()
+						}()
+						err := healthErr
 
 						if err != nil {
 							mu.Lock()

server/flight_executor_test.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"errors"
 	"runtime"
 	"strings"
@@ -12,7 +13,7 @@ func TestFlightExecutorMarkDead_QueryContext(t *testing.T) {
 	e := &FlightExecutor{} // client is nil — would panic if accessed
 	e.MarkDead()
 
-	_, err := e.QueryContext(nil, "SELECT 1")
+	_, err := e.QueryContext(context.Background(), "SELECT 1")
 	if !errors.Is(err, ErrWorkerDead) {
 		t.Fatalf("expected ErrWorkerDead, got %v", err)
 	}
@@ -22,7 +23,7 @@ func TestFlightExecutorMarkDead_ExecContext(t *testing.T) {
 	e := &FlightExecutor{}
 	e.MarkDead()
 
-	_, err := e.ExecContext(nil, "SET x = 1")
+	_, err := e.ExecContext(context.Background(), "SET x = 1")
 	if !errors.Is(err, ErrWorkerDead) {
 		t.Fatalf("expected ErrWorkerDead, got %v", err)
 	}
@@ -33,7 +34,7 @@ func TestFlightExecutorMarkDeadIdempotent(t *testing.T) {
 	e.MarkDead()
 	e.MarkDead() // should not panic
 
-	_, err := e.QueryContext(nil, "SELECT 1")
+	_, err := e.QueryContext(context.Background(), "SELECT 1")
 	if !errors.Is(err, ErrWorkerDead) {
 		t.Fatalf("expected ErrWorkerDead after double MarkDead, got %v", err)
 	}
@@ -109,7 +110,7 @@ func TestFlightExecutorNilClient_QueryContextRecovers(t *testing.T) {
 		client: nil, // simulates closed client
 	}
 
-	_, err := e.QueryContext(nil, "SELECT 1")
+	_, err := e.QueryContext(context.Background(), "SELECT 1")
 	if err == nil {
 		t.Fatal("expected error from nil client")
 	}
@@ -123,7 +124,7 @@ func TestFlightExecutorNilClient_ExecContextRecovers(t *testing.T) {
 		client: nil,
 	}
 
-	_, err := e.ExecContext(nil, "SET x = 1")
+	_, err := e.ExecContext(context.Background(), "SET x = 1")
 	if err == nil {
 		t.Fatal("expected error from nil client")
 	}