Recheck takeover liveness and classify reconnect failures

Author: bill-ph

controlplane/flight_ingress.go

@@ -3,6 +3,7 @@ package controlplane
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"log/slog"
 	"sync"
@@ -136,6 +137,9 @@ func (p *orgRoutedSessionProvider) ReconnectSession(ctx context.Context, record
 
 	pid, executor, err := sessions.ReconnectFlightSession(ctx, record.Username, record.WorkerID, record.OwnerEpoch)
 	if err != nil {
+		if errors.Is(err, configstore.ErrWorkerOwnerEpochMismatch) {
+			return 0, nil, flightsqlingress.MarkDurableReconnectTerminal(err)
+		}
 		return 0, nil, err
 	}
 

controlplane/k8s_pool.go

@@ -1180,6 +1180,11 @@ func (p *K8sWorkerPool) reserveClaimedWorker(ctx context.Context, claimed *confi
 	if reservedRecord != nil {
 		p.persistWorkerRecord(reservedRecord)
 	}
+	if err := p.checkReservedWorkerLiveness(ctx, worker); err != nil {
+		slog.Warn("Claimed worker failed liveness recheck.", "worker", worker.ID, "pod", worker.PodName(), "error", err)
+		p.retireWorkerWithReason(worker.ID, RetireReasonCrash)
+		return nil, err
+	}
 	return worker, nil
 }
 

controlplane/k8s_pool_test.go

@@ -929,6 +929,11 @@ func TestK8sPoolClaimSpecificWorkerTakesOverRuntimeWorker(t *testing.T) {
 	pool.runtimeStore = store
 	worker := &ManagedWorker{ID: 44, done: make(chan struct{})}
 	pool.workers[worker.ID] = worker
+	livenessChecked := false
+	pool.healthCheckFunc = func(ctx context.Context, worker *ManagedWorker) error {
+		livenessChecked = true
+		return nil
+	}
 
 	claimed, err := pool.claimSpecificWorker(context.Background(), 44, 7, &WorkerAssignment{
 		OrgID:          "analytics",
@@ -966,6 +971,9 @@ func TestK8sPoolClaimSpecificWorkerTakesOverRuntimeWorker(t *testing.T) {
 	if state.Assignment == nil || state.Assignment.OrgID != "analytics" {
 		t.Fatalf("expected analytics assignment, got %#v", state.Assignment)
 	}
+	if !livenessChecked {
+		t.Fatal("expected claimSpecificWorker to recheck worker liveness")
+	}
 }
 
 func TestK8sPoolClaimSpecificWorkerReturnsEpochMismatchError(t *testing.T) {
@@ -992,6 +1000,42 @@ func TestK8sPoolClaimSpecificWorkerReturnsEpochMismatchError(t *testing.T) {
 	}
 }
 
+func TestK8sPoolClaimSpecificWorkerRetiresUnhealthyWorker(t *testing.T) {
+	pool, _ := newTestK8sPool(t, 5)
+	leaseExpiry := time.Date(2026, time.March, 20, 19, 30, 0, 0, time.UTC)
+	store := &captureRuntimeWorkerStore{
+		takenOver: &configstore.WorkerRecord{
+			WorkerID:          44,
+			PodName:           "duckgres-worker-test-cp-44",
+			State:             configstore.WorkerStateReserved,
+			OrgID:             "analytics",
+			OwnerCPInstanceID: pool.cpInstanceID,
+			OwnerEpoch:        8,
+			LeaseExpiresAt:    leaseExpiry,
+		},
+	}
+	pool.runtimeStore = store
+	pool.workers[44] = &ManagedWorker{ID: 44, done: make(chan struct{})}
+	pool.healthCheckFunc = func(ctx context.Context, worker *ManagedWorker) error {
+		return errors.New("dead worker")
+	}
+
+	claimed, err := pool.claimSpecificWorker(context.Background(), 44, 7, &WorkerAssignment{
+		OrgID:          "analytics",
+		LeaseExpiresAt: leaseExpiry,
+		MaxWorkers:     3,
+	})
+	if err == nil {
+		t.Fatal("expected unhealthy claimed worker to fail liveness recheck")
+	}
+	if claimed != nil {
+		t.Fatalf("expected no claimed worker, got %#v", claimed)
+	}
+	if _, ok := pool.Worker(44); ok {
+		t.Fatal("expected unhealthy worker to be retired from the pool")
+	}
+}
+
 func TestK8sPoolReserveSharedWorkerCreatesRuntimeSpawningSlotWhenPoolIsCold(t *testing.T) {
 	pool, _ := newTestK8sPool(t, 5)
 	leaseExpiry := time.Date(2026, time.March, 20, 20, 0, 0, 0, time.UTC)

server/flightsqlingress/ingress.go

@@ -7,6 +7,7 @@ import (
 	"database/sql"
 	"encoding/base64"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"log/slog"
 	"net"
@@ -39,6 +40,15 @@ const (
 	defaultFlightSessionHeaderKey = "x-duckgres-session"
 )
 
+var ErrDurableReconnectTerminal = errors.New("durable reconnect terminal")
+
+func MarkDurableReconnectTerminal(err error) error {
+	if err == nil {
+		return nil
+	}
+	return fmt.Errorf("%w: %w", ErrDurableReconnectTerminal, err)
+}
+
 const (
 	ReapTriggerPeriodic = "periodic"
 	ReapTriggerForced   = "forced"
@@ -1735,6 +1745,9 @@ func (s *flightAuthSessionStore) reconnectByToken(ctx context.Context, token str
 	pid, executor, err := s.reconnector.ReconnectSession(ctx, *record)
 	if err != nil {
 		slog.Warn("Reconnecting durable Flight session failed.", "token", token, "error", err)
+		if errors.Is(err, ErrDurableReconnectTerminal) {
+			_ = s.durableStore.CloseSession(token, time.Now())
+		}
 		return nil, false
 	}
 

server/flightsqlingress/ingress_test.go

@@ -956,6 +956,78 @@ func TestFlightAuthSessionStoreReconnectRefreshesDurableSessionMetadata(t *testi
 	}
 }
 
+func TestFlightAuthSessionStoreReconnectFailureUpdatesDurableSessionState(t *testing.T) {
+	tests := []struct {
+		name              string
+		reconnectErr      error
+		wantState         DurableSessionState
+		wantReconnectCall int
+	}{
+		{
+			name:              "terminal stale ownership closes durable session",
+			reconnectErr:      MarkDurableReconnectTerminal(errors.New("stale owner")),
+			wantState:         DurableSessionStateClosed,
+			wantReconnectCall: 1,
+		},
+		{
+			name:              "transient reconnect failure leaves durable session active",
+			reconnectErr:      context.DeadlineExceeded,
+			wantState:         DurableSessionStateActive,
+			wantReconnectCall: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			durable := &captureDurableSessionStore{
+				records: map[string]DurableSessionRecord{
+					"durable-token": {
+						SessionToken: "durable-token",
+						Username:     "postgres",
+						OrgID:        "analytics",
+						WorkerID:     17,
+						OwnerEpoch:   4,
+						CPInstanceID: "cp-old:boot-a",
+						State:        DurableSessionStateActive,
+						ExpiresAt:    time.Now().Add(time.Hour),
+						LastSeenAt:   time.Now().Add(-time.Minute),
+					},
+				},
+			}
+			reconnectCalls := 0
+			provider := &testDurableSessionProvider{
+				durableStore: durable,
+				reconnectSessionFn: func(ctx context.Context, record DurableSessionRecord) (int32, *server.FlightExecutor, error) {
+					reconnectCalls++
+					return 0, nil, tt.reconnectErr
+				},
+			}
+			store := newFlightAuthSessionStore(provider, time.Minute, time.Hour, time.Minute, time.Hour, 0, Options{})
+
+			if session, ok := store.GetByTokenContext(context.Background(), "durable-token"); ok || session != nil {
+				t.Fatal("expected durable token reconnect to fail")
+			}
+			record, err := durable.GetSession("durable-token")
+			if err != nil {
+				t.Fatalf("GetSession: %v", err)
+			}
+			if record == nil {
+				t.Fatal("expected durable session record to remain present")
+			}
+			if record.State != tt.wantState {
+				t.Fatalf("expected durable session state %q, got %q", tt.wantState, record.State)
+			}
+
+			if session, ok := store.GetByTokenContext(context.Background(), "durable-token"); ok || session != nil {
+				t.Fatal("expected second durable token lookup to fail")
+			}
+			if reconnectCalls != tt.wantReconnectCall {
+				t.Fatalf("expected %d reconnect attempts, got %d", tt.wantReconnectCall, reconnectCalls)
+			}
+		})
+	}
+}
+
 func TestFlightAuthSessionStoreRejectsNewSessionsWhileDraining(t *testing.T) {
 	provider := &testDurableSessionProvider{
 		createSessionFn: func(ctx context.Context, username string, pid int32, memoryLimit string, threads int) (int32, *server.FlightExecutor, error) {