[Rust Frontend] Start current wave for a stale DP FirstRequest (#46833)
Signed-off-by: Blas Rodriguez Irizar <rodrigblas@gmail.com>
Blas Rodriguez Irizar committed
Jun 30, 2026 at 06:13 UTC
e45c8a9f4bad29472991ef01b2fa754a2f565723
2 files changed
+84
-9
rust/src/engine-core-client/src/coordinator/handle.rs
+21
@@ -20,6 +20,27 @@ pub(crate) struct CoordinatorStateSnapshot {
20
pub engines_running: bool,
21
}
22
23
+impl CoordinatorStateSnapshot {
24
+ /// Resume the engines for a `FirstRequest` and return the wave to broadcast
25
+ /// and the engine to exclude from the wakeup.
26
+ ///
27
+ /// The request may have been stamped with a `request_wave` older than
28
+ /// `current_wave` if a `WaveComplete` advanced it after the command was
29
+ /// enqueued. Such a request still needs serving, so the current wave is
30
+ /// broadcast to every engine (`exclude = None`); the wave is never rewound.
31
+ /// A non-stale request excludes the engine that already received it. Mirrors
32
+ /// the Python coordinator's front-end path.
33
+ pub(crate) fn start_wave_for_first_request(
34
+ &mut self,
35
+ request_wave: u32,
36
+ target_engine_index: u32,
37
+ ) -> (u32, Option<u32>) {
38
+ self.engines_running = true;
39
+ let exclude = (request_wave >= self.current_wave).then_some(target_engine_index);
40
+ (self.current_wave, exclude)
41
+ }
42
+}
43
+
44
/// Shared in-process coordinator state.
45
pub(crate) type CoordinatorState = Mutex<CoordinatorStateSnapshot>;
46
rust/src/engine-core-client/src/coordinator/inproc.rs
+63
-9
@@ -27,9 +27,10 @@ use crate::protocol::{
27
struct StartDpWaveMessage {
28
/// DP wave number that all engines should start processing.
29
wave: u32,
30
- /// Engine index that already received the triggering request and should not
31
- /// receive an extra wakeup notification.
32
- exclude_engine_index: u32,
30
+ /// Engine index that already received the triggering request and so does not
31
+ /// need an extra wakeup. `None` wakes every engine (used when the triggering
32
+ /// request was for a stale wave).
33
+ exclude_engine_index: Option<u32>,
34
}
35
36
/// Background half of the in-process coordinator.
@@ -57,7 +58,11 @@ impl InProcCoordinatorRunner {
58
}
59
60
/// Broadcast Python-compatible `START_DP_WAVE` to all connected engines.
60
- async fn broadcast_start_wave(&mut self, wave: u32, exclude_engine_index: u32) -> Result<()> {
61
+ async fn broadcast_start_wave(
62
+ &mut self,
63
+ wave: u32,
64
+ exclude_engine_index: Option<u32>,
65
+ ) -> Result<()> {
66
let payload = encode_msgpack(&StartDpWaveMessage {
67
wave,
68
exclude_engine_index,
@@ -86,13 +91,17 @@ impl InProcCoordinatorRunner {
91
engine_id: target_engine_id.to_vec(),
92
}
93
})?;
89
- self.state.lock().current_wave = wave;
94
+ let (current_wave, exclude) = {
95
+ let mut state = self.state.lock();
96
+ state.start_wave_for_first_request(wave, target_engine_index)
97
+ };
98
debug!(
91
- wave,
92
- exclude_engine_index = target_engine_index,
99
+ current_wave,
100
+ request_wave = wave,
101
+ ?exclude,
102
"starting DP wave after first request while engines were paused"
103
);
95
- self.broadcast_start_wave(wave, target_engine_index).await?;
104
+ self.broadcast_start_wave(current_wave, exclude).await?;
105
}
106
}
107
Ok(())
@@ -150,7 +159,7 @@ impl InProcCoordinatorRunner {
159
exclude_engine_index = engine_index,
160
"starting DP wave after stale-wave notification from engine"
161
);
153
- self.broadcast_start_wave(wave, engine_index).await?;
162
+ self.broadcast_start_wave(wave, Some(engine_index)).await?;
163
}
164
}
165
},
@@ -202,3 +211,48 @@ impl InProcCoordinatorRunner {
211
inner.close_registries(Arc::new(error));
212
}
213
}
214
+
215
+#[cfg(test)]
216
+mod tests {
217
+ use crate::coordinator::handle::CoordinatorStateSnapshot;
218
+
219
+ /// A `FirstRequest` for the current wave starts that wave and excludes the
220
+ /// engine that already received the triggering request.
221
+ #[test]
222
+ fn first_request_for_current_wave_excludes_target() {
223
+ let mut state = CoordinatorStateSnapshot {
224
+ current_wave: 3,
225
+ engines_running: false,
226
+ };
227
+
228
+ let (wave, exclude) = state.start_wave_for_first_request(3, 2);
229
+
230
+ assert_eq!(wave, 3);
231
+ assert_eq!(exclude, Some(2));
232
+ assert!(state.engines_running);
233
+ assert_eq!(state.current_wave, 3);
234
+ }
235
+
236
+ /// A `FirstRequest` whose wave was superseded by a racing `WaveComplete`
237
+ /// (`request_wave < current_wave`) must still start the request's wave: it
238
+ /// broadcasts the current wave and wakes every engine (`exclude = None`)
239
+ /// rather than rewinding the wave or dropping the request.
240
+ #[test]
241
+ fn stale_first_request_starts_current_wave_for_all_engines() {
242
+ let mut state = CoordinatorStateSnapshot {
243
+ current_wave: 4,
244
+ engines_running: false,
245
+ };
246
+
247
+ // Request stamped with wave 3 while the coordinator already advanced to 4.
248
+ let (wave, exclude) = state.start_wave_for_first_request(3, 2);
249
+
250
+ assert_eq!(
251
+ wave, 4,
252
+ "must broadcast the current wave, not the stale one"
253
+ );
254
+ assert_eq!(exclude, None, "a stale request must wake every engine");
255
+ assert!(state.engines_running);
256
+ assert_eq!(state.current_wave, 4, "wave must not be rewound");
257
+ }
258
+}