Merge pull request #32017 from teskje/compute-future-epoch

compute: handle future epochs in channel adapter

Author: teskje

src/compute/src/server.rs

@@ -747,28 +747,50 @@ fn spawn_channel_adapter(
             // for previous clients.
             let mut epoch = 0;
 
+            // It's possible that we receive responses with epochs from the future: Worker 0 might
+            // have increased its epoch before us and broadcasted it to our Timely cluster. When we
+            // receive a response with a future epoch, we need to wait with forwarding it until we
+            // have increased our own epoch sufficiently (by observing new client connections). We
+            // need to stash the response in the meantime.
+            let mut stashed_response = None;
+
             while let Ok((command_rx, response_tx)) = client_rx.recv() {
                 epoch += 1;
 
-                // Serve this connection until we see any of the channels disconnect.
-                loop {
+                // Wait for a new response while forwarding received commands.
+                let serve_rx_channels = || loop {
                     crossbeam_channel::select! {
                         recv(command_rx) -> msg => match msg {
                             Ok(cmd) => command_tx.send((cmd, epoch)),
-                            Err(_) => break,
+                            Err(_) => return Err(()),
                         },
                         recv(response_rx) -> msg => {
-                            let (resp, resp_epoch) = msg.expect("worker connected");
-
-                            if resp_epoch < epoch {
-                                continue; // response for a previous connection
-                            } else if resp_epoch > epoch {
-                                panic!("epoch from the future: {resp_epoch} > {epoch}");
-                            }
+                            return Ok(msg.expect("worker connected"));
+                        }
+                    }
+                };
 
-                            if response_tx.send(resp).is_err() {
-                                break;
-                            }
+                // Serve this connection until we see any of the channels disconnect.
+                loop {
+                    let (resp, resp_epoch) = match stashed_response.take() {
+                        Some(stashed) => stashed,
+                        None => match serve_rx_channels() {
+                            Ok(response) => response,
+                            Err(()) => break,
+                        },
+                    };
+
+                    if resp_epoch < epoch {
+                        // Response for a previous connection; discard it.
+                        continue;
+                    } else if resp_epoch > epoch {
+                        // Response for a future connection; stash it and reconnect.
+                        stashed_response = Some((resp, resp_epoch));
+                        break;
+                    } else {
+                        // Response for the current connection; forward it.
+                        if response_tx.send(resp).is_err() {
+                            break;
                         }
                     }
                 }