From e45c8a9f4bad29472991ef01b2fa754a2f565723 Mon Sep 17 00:00:00 2001 From: Blas Rodriguez Irizar Date: Tue, 30 Jun 2026 06:13:09 +0100 Subject: [PATCH] [Rust Frontend] Start current wave for a stale DP FirstRequest (#46833) Signed-off-by: Blas Rodriguez Irizar --- .../src/coordinator/handle.rs | 21 ++++++ .../src/coordinator/inproc.rs | 72 ++++++++++++++++--- 2 files changed, 84 insertions(+), 9 deletions(-) diff --git a/rust/src/engine-core-client/src/coordinator/handle.rs b/rust/src/engine-core-client/src/coordinator/handle.rs index dca6f70de1f..f063b624d3b 100644 --- a/rust/src/engine-core-client/src/coordinator/handle.rs +++ b/rust/src/engine-core-client/src/coordinator/handle.rs @@ -20,6 +20,27 @@ pub(crate) struct CoordinatorStateSnapshot { pub engines_running: bool, } +impl CoordinatorStateSnapshot { + /// Resume the engines for a `FirstRequest` and return the wave to broadcast + /// and the engine to exclude from the wakeup. + /// + /// The request may have been stamped with a `request_wave` older than + /// `current_wave` if a `WaveComplete` advanced it after the command was + /// enqueued. Such a request still needs serving, so the current wave is + /// broadcast to every engine (`exclude = None`); the wave is never rewound. + /// A non-stale request excludes the engine that already received it. Mirrors + /// the Python coordinator's front-end path. + pub(crate) fn start_wave_for_first_request( + &mut self, + request_wave: u32, + target_engine_index: u32, + ) -> (u32, Option) { + self.engines_running = true; + let exclude = (request_wave >= self.current_wave).then_some(target_engine_index); + (self.current_wave, exclude) + } +} + /// Shared in-process coordinator state. pub(crate) type CoordinatorState = Mutex; diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs index 26ab0c5a0e7..54c9f810d03 100644 --- a/rust/src/engine-core-client/src/coordinator/inproc.rs +++ b/rust/src/engine-core-client/src/coordinator/inproc.rs @@ -27,9 +27,10 @@ use crate::protocol::{ struct StartDpWaveMessage { /// DP wave number that all engines should start processing. wave: u32, - /// Engine index that already received the triggering request and should not - /// receive an extra wakeup notification. - exclude_engine_index: u32, + /// Engine index that already received the triggering request and so does not + /// need an extra wakeup. `None` wakes every engine (used when the triggering + /// request was for a stale wave). + exclude_engine_index: Option, } /// Background half of the in-process coordinator. @@ -57,7 +58,11 @@ impl InProcCoordinatorRunner { } /// Broadcast Python-compatible `START_DP_WAVE` to all connected engines. - async fn broadcast_start_wave(&mut self, wave: u32, exclude_engine_index: u32) -> Result<()> { + async fn broadcast_start_wave( + &mut self, + wave: u32, + exclude_engine_index: Option, + ) -> Result<()> { let payload = encode_msgpack(&StartDpWaveMessage { wave, exclude_engine_index, @@ -86,13 +91,17 @@ impl InProcCoordinatorRunner { engine_id: target_engine_id.to_vec(), } })?; - self.state.lock().current_wave = wave; + let (current_wave, exclude) = { + let mut state = self.state.lock(); + state.start_wave_for_first_request(wave, target_engine_index) + }; debug!( - wave, - exclude_engine_index = target_engine_index, + current_wave, + request_wave = wave, + ?exclude, "starting DP wave after first request while engines were paused" ); - self.broadcast_start_wave(wave, target_engine_index).await?; + self.broadcast_start_wave(current_wave, exclude).await?; } } Ok(()) @@ -150,7 +159,7 @@ impl InProcCoordinatorRunner { exclude_engine_index = engine_index, "starting DP wave after stale-wave notification from engine" ); - self.broadcast_start_wave(wave, engine_index).await?; + self.broadcast_start_wave(wave, Some(engine_index)).await?; } } }, @@ -202,3 +211,48 @@ impl InProcCoordinatorRunner { inner.close_registries(Arc::new(error)); } } + +#[cfg(test)] +mod tests { + use crate::coordinator::handle::CoordinatorStateSnapshot; + + /// A `FirstRequest` for the current wave starts that wave and excludes the + /// engine that already received the triggering request. + #[test] + fn first_request_for_current_wave_excludes_target() { + let mut state = CoordinatorStateSnapshot { + current_wave: 3, + engines_running: false, + }; + + let (wave, exclude) = state.start_wave_for_first_request(3, 2); + + assert_eq!(wave, 3); + assert_eq!(exclude, Some(2)); + assert!(state.engines_running); + assert_eq!(state.current_wave, 3); + } + + /// A `FirstRequest` whose wave was superseded by a racing `WaveComplete` + /// (`request_wave < current_wave`) must still start the request's wave: it + /// broadcasts the current wave and wakes every engine (`exclude = None`) + /// rather than rewinding the wave or dropping the request. + #[test] + fn stale_first_request_starts_current_wave_for_all_engines() { + let mut state = CoordinatorStateSnapshot { + current_wave: 4, + engines_running: false, + }; + + // Request stamped with wave 3 while the coordinator already advanced to 4. + let (wave, exclude) = state.start_wave_for_first_request(3, 2); + + assert_eq!( + wave, 4, + "must broadcast the current wave, not the stale one" + ); + assert_eq!(exclude, None, "a stale request must wake every engine"); + assert!(state.engines_running); + assert_eq!(state.current_wave, 4, "wave must not be rewound"); + } +}