clusterd: Non-blocking storage reconnection via channel adapter

Replace the blocking storage reconnection path with an async channel
adapter (spawn_storage_channel_adapter) that mirrors compute's pattern:

1. Receives new client connections from ClusterClient asynchronously
2. Buffers commands until InitializationComplete, collecting them
3. Delivers the reconciliation batch to the worker through a channel
4. Forwards ongoing commands and responses through persistent channels

The worker never blocks waiting for a new storage client — it checks
for reconciliation batches with try_recv each iteration.

Also extract StorageState::reconcile_with_commands from Worker::reconcile
so the reconciliation logic can be called without the Worker wrapper
(and without the blocking drain loop).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: antiguru

src/clusterd/src/unified.rs

@@ -321,6 +321,21 @@ fn run_unified_worker<A: Allocate + 'static>(
 
     let compute_metrics = config.compute_metrics.for_worker(worker_id);
 
+    // --- Set up storage channel adapter ---
+    // The adapter handles reconnections: it buffers commands until
+    // InitializationComplete, then delivers the batch for reconciliation.
+    // Ongoing commands and responses flow through persistent channels.
+    let (storage_reconcile_tx, storage_reconcile_rx) = std::sync::mpsc::channel();
+    let (storage_cmd_tx, storage_cmd_rx) = std::sync::mpsc::channel();
+    let (storage_resp_tx, storage_resp_rx) = mpsc::unbounded_channel();
+    spawn_storage_channel_adapter(
+        storage_client_rx,
+        storage_reconcile_tx,
+        storage_cmd_tx,
+        storage_resp_rx,
+        thread::current(),
+    );
+
     // --- Set up storage state ---
     let (internal_cmd_tx, internal_cmd_rx) =
         internal_control::setup_command_sequencer(timely_worker);
@@ -400,83 +415,42 @@ fn run_unified_worker<A: Allocate + 'static>(
         .expect_err("change to first nonce");
     compute_worker.set_nonce(nonce);
 
-    // Wait for first storage client connection.
-    let mut storage_client_rx = storage_client_rx;
-    let Some((_nonce, mut storage_cmd_rx, storage_resp_tx)) = storage_client_rx.blocking_recv()
-    else {
-        return;
-    };
-
-    // Main unified loop: alternates between compute and storage steps.
-    // On compute nonce change, reconcile compute and continue.
-    // On storage channel disconnect, wait for a new storage client.
+    // Main unified loop.
+    // Compute reconnection: handled by spawn_channel_adapter + nonce mechanism.
+    // Storage reconnection: handled by spawn_storage_channel_adapter which
+    // delivers reconciliation batches through storage_reconcile_rx.
     loop {
-        let result = run_unified_client(
+        let Err(NonceChange(nonce)) = run_unified_loop(
             &mut compute_worker,
             &mut storage_state,
-            &mut storage_cmd_rx,
+            &storage_reconcile_rx,
+            &storage_cmd_rx,
             &storage_resp_tx,
         );
-        match result {
-            UnifiedLoopExit::ComputeNonceChange(nonce) => {
-                compute_worker.set_nonce(nonce);
-            }
-            UnifiedLoopExit::StorageDisconnect => {
-                // Wait for a new storage client connection.
-                let Some((_nonce, new_cmd_rx, new_resp_tx)) = storage_client_rx.blocking_recv()
-                else {
-                    return;
-                };
-                storage_cmd_rx = new_cmd_rx;
-                // We need to restart the loop with the new storage channels.
-                // The run_unified_client function will pick up the new channels.
-                let _ = new_resp_tx; // Will be used in next iteration
-                // TODO: Handle storage reconnection properly
-                continue;
-            }
-        }
+        compute_worker.set_nonce(nonce);
     }
 }
 
-enum UnifiedLoopExit {
-    ComputeNonceChange(Uuid),
-    StorageDisconnect,
-}
-
-fn run_unified_client<A: Allocate + 'static>(
+fn run_unified_loop<A: Allocate + 'static>(
     compute: &mut ComputeWorker<'_, A>,
     storage_state: &mut StorageState,
-    storage_cmd_rx: &mut mpsc::UnboundedReceiver<StorageCommand>,
+    storage_reconcile_rx: &std::sync::mpsc::Receiver<Vec<StorageCommand>>,
+    storage_cmd_rx: &std::sync::mpsc::Receiver<StorageCommand>,
     storage_resp_tx: &mpsc::UnboundedSender<StorageResponse>,
-) -> UnifiedLoopExit {
+) -> Result<std::convert::Infallible, NonceChange> {
     // Reconcile compute.
-    if let Err(NonceChange(nonce)) = compute.reconcile() {
-        return UnifiedLoopExit::ComputeNonceChange(nonce);
-    }
-
-    // Reconcile storage (drain until InitializationComplete).
-    {
-        let mut commands = vec![];
-        loop {
-            match storage_cmd_rx.blocking_recv() {
-                Some(StorageCommand::InitializationComplete) => break,
-                Some(command) => commands.push(command),
-                None => return UnifiedLoopExit::StorageDisconnect,
-            }
-        }
-        // Apply reconciled commands.
-        // TODO: Full storage reconciliation logic (matching Worker::reconcile).
-        // For now, apply commands directly.
-        for command in commands {
-            storage_state.handle_storage_command(command);
-        }
-    }
+    compute.reconcile()?;
 
     let mut last_compute_maintenance = Instant::now();
     let mut last_storage_maintenance = Instant::now();
     let mut last_stats_time = Instant::now();
 
     loop {
+        // --- Check for storage reconnection ---
+        if let Ok(commands) = storage_reconcile_rx.try_recv() {
+            storage_state.reconcile_with_commands(commands);
+        }
+
         // --- Compute maintenance ---
         let compute_maintenance_interval = compute
             .compute_state
@@ -521,8 +495,9 @@ fn run_unified_client<A: Allocate + 'static>(
             .statistics_collection_interval;
         let stats_remaining = stats_interval.saturating_sub(last_stats_time.elapsed());
 
+        // The storage channel adapter unparks the worker thread when it sends
+        // commands, so we don't need to check storage_cmd_rx here.
         let can_park = compute.command_rx.is_empty()
-            && storage_cmd_rx.is_empty()
             && storage_state.async_worker.is_empty();
 
         let sleep = min_durations(&[compute_sleep, storage_sleep, Some(stats_remaining)]);
@@ -537,9 +512,7 @@ fn run_unified_client<A: Allocate + 'static>(
         timer.observe_duration();
 
         // --- Compute step ---
-        if let Err(NonceChange(nonce)) = compute.handle_pending_commands() {
-            return UnifiedLoopExit::ComputeNonceChange(nonce);
-        }
+        compute.handle_pending_commands()?;
         if let Some(mut cs) = compute.activate_compute() {
             cs.process_peeks();
             cs.process_subscribes();
@@ -565,20 +538,8 @@ fn run_unified_client<A: Allocate + 'static>(
         }
 
         // Drain storage commands
-        let mut storage_disconnected = false;
-        loop {
-            match storage_cmd_rx.try_recv() {
-                Ok(cmd) => storage_state.handle_storage_command(cmd),
-                Err(mpsc::error::TryRecvError::Empty) => break,
-                Err(mpsc::error::TryRecvError::Disconnected) => {
-                    storage_disconnected = true;
-                    break;
-                }
-            }
-        }
-
-        if storage_disconnected {
-            return UnifiedLoopExit::StorageDisconnect;
+        while let Ok(cmd) = storage_cmd_rx.try_recv() {
+            storage_state.handle_storage_command(cmd);
         }
 
         // Drain async worker responses
@@ -593,6 +554,72 @@ fn run_unified_client<A: Allocate + 'static>(
     }
 }
 
+/// Spawn a task that bridges [`ClusterClient`] connections to persistent worker channels.
+///
+/// For each new storage client connection:
+/// 1. Drains commands until `InitializationComplete`, collecting them into a batch
+/// 2. Sends the batch through `reconcile_tx` for the worker to reconcile in one go
+/// 3. Forwards ongoing commands through `command_tx`
+/// 4. Forwards worker responses from `response_rx` back to the active client
+///
+/// This ensures the worker thread never blocks waiting for a new client connection.
+fn spawn_storage_channel_adapter(
+    mut client_rx: mpsc::UnboundedReceiver<(
+        Uuid,
+        mpsc::UnboundedReceiver<StorageCommand>,
+        mpsc::UnboundedSender<StorageResponse>,
+    )>,
+    reconcile_tx: std::sync::mpsc::Sender<Vec<StorageCommand>>,
+    command_tx: std::sync::mpsc::Sender<StorageCommand>,
+    mut response_rx: mpsc::UnboundedReceiver<StorageResponse>,
+    worker_thread: thread::Thread,
+) {
+    mz_ore::task::spawn(
+        || "storage-channel-adapter",
+        async move {
+            while let Some((_nonce, mut cmd_rx, resp_tx)) = client_rx.recv().await {
+                // Phase 1: Buffer commands until InitializationComplete.
+                let mut commands = vec![];
+                let init_complete = loop {
+                    match cmd_rx.recv().await {
+                        Some(StorageCommand::InitializationComplete) => break true,
+                        Some(cmd) => commands.push(cmd),
+                        None => break false,
+                    }
+                };
+                if !init_complete {
+                    continue;
+                }
+
+                // Deliver reconciliation batch to the worker.
+                if reconcile_tx.send(commands).is_err() {
+                    break;
+                }
+                worker_thread.unpark();
+
+                // Phase 2: Forward ongoing commands and responses.
+                loop {
+                    tokio::select! {
+                        cmd = cmd_rx.recv() => match cmd {
+                            Some(cmd) => {
+                                if command_tx.send(cmd).is_err() {
+                                    return;
+                                }
+                                worker_thread.unpark();
+                            }
+                            None => break, // Client disconnected, wait for next.
+                        },
+                        resp = response_rx.recv() => match resp {
+                            Some(resp) => { let _ = resp_tx.send(resp); }
+                            None => return, // Worker gone.
+                        },
+                    }
+                }
+            }
+        },
+    );
+}
+
 /// Returns the minimum of a set of optional durations.
 fn min_durations(durations: &[Option<Duration>]) -> Option<Duration> {
     durations.iter().copied().flatten().min()