Treat draining control planes as live in orphan worker sweep

Follow-up to #407. The startup orphan sweep used
ListActiveControlPlaneInstanceIDs which only returned CPs in state
'active', so a CP that was mid-graceful-shutdown (state='draining',
still serving in-flight queries while waiting for them to finish)
appeared "dead" to a new peer's startup sweep. Phase 1 would then
delete the draining CP's worker pods, killing every query the
draining CP was waiting on.

The pre-existing ListStuckWorkers query already uses
`cp.state <> 'expired'` for the same "alive" notion, so this aligns
the new sweep with the existing convention.

  - Renamed ListActiveControlPlaneInstanceIDs to
    ListLiveControlPlaneInstanceIDs and changed the underlying query
    from `state = 'active'` to `state <> 'expired'`. The doc comment
    spells out why draining must count as live.
  - Renamed activeCPIDs / activeCPLabels to liveCPIDs / liveCPLabels
    in cleanupOrphanedWorkers and updated both the phase 1 pod-delete
    filter and the phase 2 owner-skip filter to use them.
  - Updated the captureRuntimeWorkerStore mock and existing test
    field names.
  - Added TestCleanupOrphanedWorkers_PreservesDrainingPeerWorkers as
    a regression test: a draining peer with both a busy worker pod
    and a stale-looking idle row must end the sweep with the pod
    intact and the row untouched.
  - Added TestListLiveControlPlaneInstanceIDsPostgres as an
    integration test against real Postgres covering active /
    draining / expired states.

Workers serving queries are in reserved/hot/activating state, which
phase 2 doesn't query at all (it only lists idle and spawning rows),
so phase 2 was already safe. The bug was strictly in phase 1's
pod-delete filter.

Author: benben

controlplane/configstore/store.go

@@ -422,16 +422,21 @@ func (cs *ConfigStore) GetControlPlaneInstance(id string) (*ControlPlaneInstance
 	return &instance, nil
 }
 
-// ListActiveControlPlaneInstanceIDs returns the IDs of control-plane instances
-// currently in the active state. Used by the K8s pool's startup orphan sweep
-// to distinguish "owned by a live peer" from "owned by a dead CP" without
-// needing N round-trips.
-func (cs *ConfigStore) ListActiveControlPlaneInstanceIDs() ([]string, error) {
+// ListLiveControlPlaneInstanceIDs returns the IDs of control-plane instances
+// that are not yet expired — i.e. either currently active OR draining (still
+// alive, waiting on in-flight queries to finish before SIGTERM completes).
+// Used by the K8s pool's startup orphan sweep to distinguish "owned by a CP
+// that is still serving traffic" from "owned by a dead CP".
+//
+// Including draining CPs is critical: a draining CP's worker pods are still
+// running queries that haven't finished yet, and treating them as orphans
+// would kill those queries mid-flight.
+func (cs *ConfigStore) ListLiveControlPlaneInstanceIDs() ([]string, error) {
 	var ids []string
 	if err := cs.db.Table(cs.runtimeTable((&ControlPlaneInstance{}).TableName())).
-		Where("state = ?", ControlPlaneInstanceStateActive).
+		Where("state <> ?", ControlPlaneInstanceStateExpired).
 		Pluck("id", &ids).Error; err != nil {
-		return nil, fmt.Errorf("list active control plane instance ids: %w", err)
+		return nil, fmt.Errorf("list live control plane instance ids: %w", err)
 	}
 	return ids, nil
 }

controlplane/k8s_pool.go

@@ -216,27 +216,32 @@ func (p *K8sWorkerPool) cleanupOrphanedWorkers() {
 		return
 	}
 
-	// Build the set of currently-active CP instance IDs (both raw and label-
+	// Build the set of currently-live CP instance IDs (both raw and label-
 	// sanitized forms) so phase 1 can compare against pod labels and phase 2
 	// can compare against owner_cp_instance_id columns.
-	activeCPIDs := map[string]bool{}
-	activeCPLabels := map[string]bool{}
+	//
+	// "Live" includes both `active` and `draining` states. A draining CP
+	// is mid-graceful-shutdown waiting on in-flight queries to finish — its
+	// worker pods are still serving traffic and must NOT be treated as
+	// orphans. Only `expired` CPs are dead enough to clean up after.
+	liveCPIDs := map[string]bool{}
+	liveCPLabels := map[string]bool{}
 	if p.runtimeStore != nil {
-		ids, err := p.runtimeStore.ListActiveControlPlaneInstanceIDs()
+		ids, err := p.runtimeStore.ListLiveControlPlaneInstanceIDs()
 		if err != nil {
-			slog.Warn("Failed to list active control-plane instances for orphan cleanup; skipping sweep to avoid wiping live peers.", "error", err)
+			slog.Warn("Failed to list live control-plane instances for orphan cleanup; skipping sweep to avoid wiping live peers.", "error", err)
 			return
 		}
 		for _, id := range ids {
-			activeCPIDs[id] = true
-			activeCPLabels[controlPlaneIDLabelValue(id)] = true
+			liveCPIDs[id] = true
+			liveCPLabels[controlPlaneIDLabelValue(id)] = true
 		}
 	}
-	// Always treat the current CP as active so the sweep is safe in
+	// Always treat the current CP as live so the sweep is safe in
 	// non-runtime-store deployments and during the brief window before this
 	// CP's heartbeat is first written.
-	activeCPIDs[p.cpInstanceID] = true
-	activeCPLabels[controlPlaneIDLabelValue(p.cpInstanceID)] = true
+	liveCPIDs[p.cpInstanceID] = true
+	liveCPLabels[controlPlaneIDLabelValue(p.cpInstanceID)] = true
 
 	// Phase 1: delete pods owned by dead CPs.
 	gracePeriod := int64(10)
@@ -245,7 +250,7 @@ func (p *K8sWorkerPool) cleanupOrphanedWorkers() {
 	var deletedNames sync.Map
 	for _, pod := range pods.Items {
 		podInstanceID := pod.Labels["duckgres/cp-instance-id"]
-		if podInstanceID == "" || activeCPLabels[podInstanceID] {
+		if podInstanceID == "" || liveCPLabels[podInstanceID] {
 			continue
 		}
 		wg.Add(1)
@@ -295,7 +300,7 @@ func (p *K8sWorkerPool) cleanupOrphanedWorkers() {
 		if livePodNames[row.PodName] {
 			continue // pod is alive in K8s; row reflects reality
 		}
-		if row.OwnerCPInstanceID != "" && activeCPIDs[row.OwnerCPInstanceID] {
+		if row.OwnerCPInstanceID != "" && liveCPIDs[row.OwnerCPInstanceID] {
 			continue // a live owner is responsible for this row's lifecycle
 		}
 		record := row

controlplane/k8s_pool_test.go

@@ -56,8 +56,8 @@ type captureRuntimeWorkerStore struct {
 	takeOverOwnerCPID     string
 	takeOverOrgID         string
 	takeOverExpectedEpoch int64
-	activeCPIDs           []string
-	activeCPIDsErr        error
+	liveCPIDs             []string
+	liveCPIDsErr          error
 	stateBeforeRows       []configstore.WorkerRecord
 	stateBeforeErr        error
 	stateBeforeCalls      int
@@ -190,14 +190,14 @@ func (s *captureRuntimeWorkerStore) TakeOverWorker(workerID int, ownerCPInstance
 	return &record, nil
 }
 
-func (s *captureRuntimeWorkerStore) ListActiveControlPlaneInstanceIDs() ([]string, error) {
+func (s *captureRuntimeWorkerStore) ListLiveControlPlaneInstanceIDs() ([]string, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	if s.activeCPIDsErr != nil {
-		return nil, s.activeCPIDsErr
+	if s.liveCPIDsErr != nil {
+		return nil, s.liveCPIDsErr
 	}
-	out := make([]string, len(s.activeCPIDs))
-	copy(out, s.activeCPIDs)
+	out := make([]string, len(s.liveCPIDs))
+	copy(out, s.liveCPIDs)
 	return out, nil
 }
 
@@ -1965,7 +1965,7 @@ func TestCleanupOrphanedWorkers_PreservesLivePeerPods(t *testing.T) {
 	deadCPID := "cp-dead:boot-789"
 
 	store := &captureRuntimeWorkerStore{
-		activeCPIDs: []string{"cp-self:boot-123", livePeerID},
+		liveCPIDs: []string{"cp-self:boot-123", livePeerID},
 	}
 	pool.runtimeStore = store
 
@@ -2037,7 +2037,7 @@ func TestCleanupOrphanedWorkers_MarksStaleIdleRowsLost(t *testing.T) {
 
 	stale := time.Now().Add(-2 * time.Hour)
 	store := &captureRuntimeWorkerStore{
-		activeCPIDs: []string{"cp-self:boot-123"},
+		liveCPIDs: []string{"cp-self:boot-123"},
 		stateBeforeRows: []configstore.WorkerRecord{
 			{
 				WorkerID:          1,
@@ -2099,7 +2099,7 @@ func TestCleanupOrphanedWorkers_PreservesInflightSpawnsByLivePeers(t *testing.T)
 
 	stale := time.Now().Add(-2 * time.Hour)
 	store := &captureRuntimeWorkerStore{
-		activeCPIDs: []string{"cp-self:boot-123", peerID},
+		liveCPIDs: []string{"cp-self:boot-123", peerID},
 		stateBeforeRows: []configstore.WorkerRecord{
 			// In-flight spawn owned by a live peer: pod isn't visible to us
 			// yet (e.g. slow image pull), but the peer is alive and will
@@ -2134,6 +2134,72 @@ func TestCleanupOrphanedWorkers_PreservesInflightSpawnsByLivePeers(t *testing.T)
 	}
 }
 
+func TestCleanupOrphanedWorkers_PreservesDrainingPeerWorkers(t *testing.T) {
+	// A draining CP is mid-graceful-shutdown, still serving in-flight queries
+	// on its worker pods. The startup sweep must treat draining CPs as live
+	// (state <> 'expired'), otherwise it would delete the worker pods and
+	// kill the queries the draining CP is waiting to finish.
+	pool, cs := newTestK8sPool(t, 5)
+	pool.cpInstanceID = "cp-self:boot-123"
+	drainingPeerID := "cp-draining-peer:boot-456"
+
+	// The mock returns whatever the test sets in liveCPIDs; the production
+	// store query is `state <> expired` so callers always get both active
+	// and draining CPs in this list.
+	store := &captureRuntimeWorkerStore{
+		liveCPIDs: []string{"cp-self:boot-123", drainingPeerID},
+		stateBeforeRows: []configstore.WorkerRecord{
+			// A warm idle row owned (still!) by the draining peer. Even
+			// though idle rows have empty owner in the production code,
+			// exercise the owner-active path here to make sure draining
+			// owners are honored.
+			{
+				WorkerID:          50,
+				PodName:           "duckgres-worker-draining-warm",
+				State:             configstore.WorkerStateIdle,
+				OwnerCPInstanceID: drainingPeerID,
+				UpdatedAt:         time.Now().Add(-2 * time.Hour),
+			},
+		},
+	}
+	pool.runtimeStore = store
+
+	// Worker pod owned by the draining peer — pretend it's serving a
+	// long-running query right now.
+	_, err := cs.CoreV1().Pods("default").Create(context.Background(), &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "duckgres-worker-draining-busy",
+			Namespace: "default",
+			Labels: map[string]string{
+				"app":                     "duckgres-worker",
+				"duckgres/cp-instance-id": controlPlaneIDLabelValue(drainingPeerID),
+			},
+		},
+	}, metav1.CreateOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	pool.cleanupOrphanedWorkers()
+
+	pods, _ := cs.CoreV1().Pods("default").List(context.Background(), metav1.ListOptions{
+		LabelSelector: "app=duckgres-worker",
+	})
+	survived := false
+	for _, p := range pods.Items {
+		if p.Name == "duckgres-worker-draining-busy" {
+			survived = true
+		}
+	}
+	if !survived {
+		t.Fatal("draining peer's worker pod was deleted — would kill in-flight queries")
+	}
+
+	if rows := store.snapshot(); len(rows) != 0 {
+		t.Fatalf("draining peer's idle row was marked lost: %#v", rows)
+	}
+}
+
 func TestCleanupOrphanedWorkers_GracePeriodProtectsFreshRows(t *testing.T) {
 	// The mock honors the cutoff itself, so passing a fresh updated_at row
 	// exercises the grace-period filter end-to-end: the sweep asks for rows
@@ -2143,7 +2209,7 @@ func TestCleanupOrphanedWorkers_GracePeriodProtectsFreshRows(t *testing.T) {
 	pool.cpInstanceID = "cp-self:boot-123"
 
 	store := &captureRuntimeWorkerStore{
-		activeCPIDs: []string{"cp-self:boot-123"},
+		liveCPIDs: []string{"cp-self:boot-123"},
 		stateBeforeRows: []configstore.WorkerRecord{
 			{
 				WorkerID:          21,

controlplane/worker_pool.go

@@ -81,7 +81,7 @@ type RuntimeWorkerStore interface {
 	CreateNeutralWarmWorkerSlot(ownerCPInstanceID, podNamePrefix string, targetWarmWorkers, maxGlobalWorkers int) (*configstore.WorkerRecord, error)
 	GetWorkerRecord(workerID int) (*configstore.WorkerRecord, error)
 	TakeOverWorker(workerID int, ownerCPInstanceID, orgID string, expectedOwnerEpoch int64) (*configstore.WorkerRecord, error)
-	ListActiveControlPlaneInstanceIDs() ([]string, error)
+	ListLiveControlPlaneInstanceIDs() ([]string, error)
 	ListWorkerRecordsByStatesBefore(states []configstore.WorkerState, updatedBefore time.Time) ([]configstore.WorkerRecord, error)
 }
 

tests/configstore/runtime_store_postgres_test.go

@@ -277,6 +277,70 @@ func TestExpireControlPlaneInstancesPostgres(t *testing.T) {
 	}
 }
 
+func TestListLiveControlPlaneInstanceIDsPostgres(t *testing.T) {
+	store := newIsolatedConfigStore(t)
+	now := time.Date(2026, time.April, 9, 12, 0, 0, 0, time.UTC)
+
+	// Active CP — should appear.
+	if err := store.UpsertControlPlaneInstance(&configstore.ControlPlaneInstance{
+		ID:              "cp-active:boot-a",
+		PodName:         "duckgres-active",
+		PodUID:          "pod-active",
+		BootID:          "boot-a",
+		State:           configstore.ControlPlaneInstanceStateActive,
+		StartedAt:       now.Add(-time.Hour),
+		LastHeartbeatAt: now,
+	}); err != nil {
+		t.Fatalf("UpsertControlPlaneInstance(active): %v", err)
+	}
+	// Draining CP — must also appear (still serving in-flight queries).
+	drainingAt := now.Add(-time.Minute)
+	if err := store.UpsertControlPlaneInstance(&configstore.ControlPlaneInstance{
+		ID:              "cp-draining:boot-b",
+		PodName:         "duckgres-draining",
+		PodUID:          "pod-draining",
+		BootID:          "boot-b",
+		State:           configstore.ControlPlaneInstanceStateDraining,
+		StartedAt:       now.Add(-time.Hour),
+		LastHeartbeatAt: now.Add(-30 * time.Second),
+		DrainingAt:      &drainingAt,
+	}); err != nil {
+		t.Fatalf("UpsertControlPlaneInstance(draining): %v", err)
+	}
+	// Expired CP — must NOT appear.
+	expiredAt := now.Add(-2 * time.Minute)
+	if err := store.UpsertControlPlaneInstance(&configstore.ControlPlaneInstance{
+		ID:              "cp-expired:boot-c",
+		PodName:         "duckgres-expired",
+		PodUID:          "pod-expired",
+		BootID:          "boot-c",
+		State:           configstore.ControlPlaneInstanceStateExpired,
+		StartedAt:       now.Add(-time.Hour),
+		LastHeartbeatAt: now.Add(-5 * time.Minute),
+		ExpiredAt:       &expiredAt,
+	}); err != nil {
+		t.Fatalf("UpsertControlPlaneInstance(expired): %v", err)
+	}
+
+	ids, err := store.ListLiveControlPlaneInstanceIDs()
+	if err != nil {
+		t.Fatalf("ListLiveControlPlaneInstanceIDs: %v", err)
+	}
+	got := map[string]bool{}
+	for _, id := range ids {
+		got[id] = true
+	}
+	if !got["cp-active:boot-a"] {
+		t.Error("expected active CP to be listed as live")
+	}
+	if !got["cp-draining:boot-b"] {
+		t.Error("expected draining CP to be listed as live (still serving in-flight queries)")
+	}
+	if got["cp-expired:boot-c"] {
+		t.Error("expected expired CP to NOT be listed as live")
+	}
+}
+
 func TestExpireDrainingControlPlaneInstancesPostgres(t *testing.T) {
 	store := newIsolatedConfigStore(t)