diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index e7b6e42b6a..63689ea3db 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -5,7 +5,7 @@ use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; -use futures::future::{join_all, try_join_all}; +use futures::future::join_all; use itertools::Itertools; use serde::Serialize; use serde_json::Value as JsonValue; @@ -450,6 +450,11 @@ impl EngineCoreClient { self.config.transport_mode.data_parallel_size() } + #[cfg(test)] + pub(crate) fn pending_utility_call_count(&self) -> usize { + self.inner.pending_utility_call_count() + } + /// Return the engine-side indices connected to this client. pub fn engine_indices(&self) -> Vec { self.engines @@ -596,6 +601,16 @@ impl EngineCoreClient { let (engine_id, rx) = self.inner.register_request(request_id.clone(), lora_name, data_parallel_rank)?; + // Construct the output stream first before actually sending the request to the engine. + // This ensures that cancelling the future will properly clean up the request registration + // via `Drop` of the stream. + let stream = EngineCoreOutputStream::new( + request_id.clone(), + engine_id.engine_index().unwrap_or(0), + self.abort_tx.clone(), + rx, + ); + let result: Result<()> = async { if let Some(coordinator) = self.coordinator.as_ref() { let snapshot = coordinator.snapshot(); @@ -622,12 +637,7 @@ impl EngineCoreClient { return Err(error); } - Ok(EngineCoreOutputStream::new( - request_id, - engine_id.engine_index().unwrap_or(0), - self.abort_tx.clone(), - rx, - )) + Ok(stream) } /// Abort currently in-flight requests by request ID. @@ -651,8 +661,9 @@ impl EngineCoreClient { } /// Call a typed utility method on all connected engines, returning one - /// decoded result per connected engine if all calls succeed or an error - /// if any call fails. + /// decoded result per connected engine if all calls succeed. The client + /// waits for every engine outcome before returning an error so callers can + /// safely compensate partially applied mutations. /// /// Callers should pass utility arguments using Rust tuple semantics so the /// encoded payload matches Python's `(client_index, call_id, @@ -669,62 +680,47 @@ impl EngineCoreClient { "sending utility request" ); - // Phase 1: allocate one call id per engine and build the per-engine - // request payloads up-front. Any failure here (registry closed, encode - // error) must roll back the call ids already allocated so they do not - // leak in the utility registry until shutdown. - let mut pending_calls = Vec::with_capacity(self.engines.len()); - let mut prepared_sends = Vec::with_capacity(self.engines.len()); + /// Removes utility waiters if a call future is cancelled before completion. + struct UtilityCallGuard<'a> { + inner: &'a ClientInner, + call_ids: Vec, + } + + impl Drop for UtilityCallGuard<'_> { + fn drop(&mut self) { + self.inner.unregister_utility_calls(self.call_ids.drain(..)); + } + } + + let mut call_guard = UtilityCallGuard { + inner: self.inner.as_ref(), + call_ids: Vec::with_capacity(self.engines.len()), + }; + + let mut prepared_calls = Vec::with_capacity(self.engines.len()); for engine in &self.engines { - let (call_id, rx) = match self.inner.allocate_and_register_utility_call() { - Ok(pair) => pair, - Err(err) => { - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id)); - return Err(err); - } - }; - let request = match EngineCoreUtilityRequest::new( - self.config.client_index, - call_id, - method, - &args, - ) { - Ok(request) => request, - Err(err) => { - self.inner.unregister_utility_calls( - pending_calls.iter().map(|(id, _)| *id).chain(std::iter::once(call_id)), - ); - return Err(err); - } - }; - pending_calls.push((call_id, rx)); - prepared_sends.push((&engine.engine_id, request)); + let (call_id, rx) = self.inner.allocate_and_register_utility_call()?; + call_guard.call_ids.push(call_id); + let request = + EngineCoreUtilityRequest::new(self.config.client_index, call_id, method, &args)?; + prepared_calls.push((&engine.engine_id, call_id, rx, request)); } - // Phase 2: dispatch every utility request concurrently. `try_join_all` - // fails fast on the first transport error and drops the remaining send - // futures; any engines that already received the request will reply, - // but those replies are simply dropped because we roll back the call - // ids below. - let send_futures = prepared_sends.iter().map(|(engine_id, request)| { - self.inner.send_to_engine(engine_id, EngineCoreRequestType::Utility, request) - }); - if let Err(err) = try_join_all(send_futures).await { - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id)); - return Err(err); - } - - // Phase 3: wait for all engines to respond and preserve the per-engine - // result list. - let futures = pending_calls.into_iter().map(|(call_id, rx)| async move { - rx.await - .map_err(|_| Error::UtilityCallClosed { - method: method.to_string(), - call_id, - })?? - .into_typed_result(method) - }); - try_join_all(futures).await + let outcomes = join_all(prepared_calls.into_iter().map( + |(engine_id, call_id, rx, request)| async move { + self.inner + .send_to_engine(engine_id, EngineCoreRequestType::Utility, &request) + .await?; + rx.await + .map_err(|_| Error::UtilityCallClosed { + method: method.to_string(), + call_id, + })?? + .into_typed_result(method) + }, + )) + .await; + outcomes.into_iter().collect() } /// Call a utility method on all connected engines and return the shared diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 9c539bf180..117473ce47 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -118,6 +118,11 @@ impl ClientInner { self.utility_reg.lock().unregister_many(call_ids); } + #[cfg(test)] + pub fn pending_utility_call_count(&self) -> usize { + self.utility_reg.lock().len() + } + /// Undo a request registration when `add_request()` fails. pub fn rollback_request(&self, request_id: &str) { let _ = self.request_reg.lock().remove(request_id); diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 59f77bf71c..724ae5b64c 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -441,6 +441,11 @@ impl UtilityRegistry { self.utility_calls.contains_key(&call_id) } + #[cfg(test)] + pub fn len(&self) -> usize { + self.utility_calls.len() + } + pub fn is_closed(&self) -> bool { self.closed } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 2f988e2bfe..df9eec90b3 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -1431,23 +1431,20 @@ async fn is_sleeping_wrapper_sends_typed_request_and_returns_typed_response() { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn call_utility_failure_message_surfaces_as_error() { +async fn call_utility_waits_for_all_engines_before_returning_error() { init_tracing(); let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let engine_id = b"engine-utility-fail".to_vec(); - - let (shutdown_tx, engine_task) = spawn_mock_engine_task( + let (failure_sent_tx, failure_sent_rx) = oneshot::channel(); + let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task( handshake_address.clone(), - engine_id.clone(), - |dealer, push| { + EngineId::from_engine_index(0).into_frame().to_vec(), + move |dealer, push| { Box::pin(async move { let utility = recv_engine_message(dealer).await; - assert_eq!(utility[0].as_ref(), &[0x03]); let payload = decode_value(&utility[1]); let call_id = payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id"); - send_outputs( push, UtilityCallOutput { @@ -1462,35 +1459,134 @@ async fn call_utility_failure_message_surfaces_as_error() { .into(), ) .await; + let _ = failure_sent_tx.send(()); + }) + }, + ); + let (second_received_tx, second_received_rx) = oneshot::channel(); + let (release_second_tx, release_second_rx) = oneshot::channel(); + let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task( + handshake_address.clone(), + EngineId::from_engine_index(1).into_frame().to_vec(), + move |dealer, push| { + Box::pin(async move { + let utility = recv_engine_message(dealer).await; + let payload = decode_value(&utility[1]); + let call_id = + payload.as_array().and_then(|array| array[1].as_u64()).expect("call_id"); + let _ = second_received_tx.send(()); + let _ = release_second_rx.await; + send_outputs( + push, + UtilityCallOutput { + engine_index: 1, + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(utility_result_value(true)), + }, + } + .into(), + ) + .await; }) }, ); - let client = connect_client_with_ipc( - handshake_test_config( - handshake_address, - 1, - "test-model", - Duration::from_secs(2), - 0, - None, - ), - &ipc, - ) - .await; + let client = std::sync::Arc::new( + connect_client_with_ipc( + handshake_test_config( + handshake_address, + 2, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await, + ); + let call_client = client.clone(); + let call = + tokio::spawn(async move { call_client.call_utility::("test_mutation", ()).await }); - let error = client.call_utility::("is_sleeping", ()).await.unwrap_err(); + failure_sent_rx.await.unwrap(); + second_received_rx.await.unwrap(); + timeout(Duration::from_secs(2), async { + while client.pending_utility_call_count() != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("wait for first engine failure"); + assert!(!call.is_finished()); + + let _ = release_second_tx.send(()); + let error = call.await.unwrap().unwrap_err(); assert!(matches!( error, Error::UtilityCallFailed { method, message, .. - } if method == "is_sleeping" && message == "boom" + } if method == "test_mutation" && message == "boom" )); - let _ = shutdown_tx.send(()); - engine_task.await.unwrap(); + let _ = shutdown_tx_0.send(()); + let _ = shutdown_tx_1.send(()); + engine_task_0.await.unwrap(); + engine_task_1.await.unwrap(); + let client = std::sync::Arc::try_unwrap(client) + .unwrap_or_else(|_| panic!("utility task retained client after completion")); + client.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_utility_call_unregisters_waiter() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let (received_tx, received_rx) = oneshot::channel(); + let (_shutdown, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + vec![0x00, 0x00], + |dealer, _push| { + Box::pin(async move { + let _utility = recv_engine_message(dealer).await; + let _ = received_tx.send(()); + std::future::pending::<()>().await; + }) + }, + ); + let client = std::sync::Arc::new( + connect_client_with_ipc( + handshake_test_config( + handshake_address, + 1, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await, + ); + let call_client = client.clone(); + let call = + tokio::spawn(async move { call_client.call_utility::("add_lora", ()).await }); + + received_rx.await.unwrap(); + assert_eq!(client.pending_utility_call_count(), 1); + call.abort(); + assert!(call.await.unwrap_err().is_cancelled()); + assert_eq!(client.pending_utility_call_count(), 0); + + engine_task.abort(); + let client = std::sync::Arc::try_unwrap(client) + .unwrap_or_else(|_| panic!("utility task retained client after cancellation")); client.shutdown().await.unwrap(); }