diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 9a66ad84cc1..44eae7e3e54 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -1,5 +1,4 @@ use std::collections::BTreeMap; -use std::slice; use std::sync::Arc; use arc_swap::ArcSwapOption; @@ -253,33 +252,47 @@ pub(crate) async fn run_abort_loop( inner: Arc, mut abort_rx: mpsc::UnboundedReceiver, ) { - // TODO: receive and abort requests in batch - while let Some(AbortRequest { request_id, cause }) = abort_rx.recv().await { - let Some(engine_id) = inner.take_auto_abort_target(&request_id) else { - debug!(request_id, "skip auto-abort for inactive request"); - continue; - }; + // Coalesce bursts of auto-aborts into a single Abort message per engine. + // A dropped-stream storm (e.g. many clients disconnecting at once under + // high concurrency) would otherwise issue one engine round-trip per + // request. `recv_many` returns as soon as at least one item is ready, so a + // lone abort is still forwarded promptly. + const MAX_DRAIN: usize = 1024; + let mut batch: Vec = Vec::new(); - match cause { - AbortCause::DroppedStream => { - info!(request_id, "auto-aborting request due to dropped stream") - } - AbortCause::StopStringMatched => { - debug!( - request_id, - "auto-aborting request due to stop string matched" - ) + while abort_rx.recv_many(&mut batch, MAX_DRAIN).await > 0 { + let mut by_engine: BTreeMap> = BTreeMap::new(); + + for AbortRequest { request_id, cause } in batch.drain(..) { + let Some(engine_id) = inner.take_auto_abort_target(&request_id) else { + debug!(request_id, "skip auto-abort for inactive request"); + continue; + }; + + match cause { + AbortCause::DroppedStream => { + info!(request_id, "auto-aborting request due to dropped stream") + } + AbortCause::StopStringMatched => { + debug!( + request_id, + "auto-aborting request due to stop string matched" + ) + } } + + by_engine.entry(engine_id).or_default().push(request_id); } - if let Err(error) = inner.do_abort_requests(&engine_id, slice::from_ref(&request_id)).await - { - warn!( - request_id, - ?engine_id, - error = %error.as_report(), - "failed to auto-abort dropped request stream" - ); + for (engine_id, request_ids) in by_engine { + if let Err(error) = inner.do_abort_requests(&engine_id, &request_ids).await { + warn!( + ?engine_id, + ?request_ids, + error = %error.as_report(), + "failed to auto-abort request streams" + ); + } } } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 9a92ffe447e..32530d6d385 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -1225,6 +1225,86 @@ async fn dropping_a_live_stream_triggers_abort() { client.shutdown().await.unwrap(); } +#[tokio::test] +async fn dropping_multiple_live_streams_aborts_all_in_a_burst() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-burst".to_vec(); + let request_ids = ["req-1", "req-2", "req-3"]; + + let (shutdown_tx, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + |dealer, push| { + Box::pin(async move { + for _ in 0..3 { + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + } + send_outputs( + push, + EngineCoreOutputs { + outputs: vec![ + request_output("req-1", vec![99], None), + request_output("req-2", vec![99], None), + request_output("req-3", vec![99], None), + ], + ..Default::default() + }, + ) + .await; + + let abort = + timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap(); + assert_eq!(abort[0].as_ref(), &[0x01]); + let ids: Vec = rmp_serde::from_slice(&abort[1]).unwrap(); + assert_eq!( + ids, + vec![ + "req-1".to_string(), + "req-2".to_string(), + "req-3".to_string() + ] + ); + assert!( + timeout(Duration::from_millis(100), recv_engine_message(dealer)).await.is_err() + ); + }) + }, + ); + + let client = connect_client_with_ipc( + handshake_test_config( + handshake_address, + 1, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await; + + // Open every request first so all three adds reach the engine before it + // emits outputs, then drain the first token from each stream. + let mut streams = Vec::new(); + for id in request_ids { + streams.push(client.call(sample_request_with_id(id)).await.unwrap()); + } + for stream in streams.iter_mut() { + let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); + assert_eq!(first.new_token_ids, vec![99]); + } + // Drop the whole burst back-to-back so the abort worker can batch them. + drop(streams); + + let _ = shutdown_tx.send(()); + engine_task.await.unwrap(); + client.shutdown().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn dispatcher_failure_propagates_to_streams_and_future_calls() { init_tracing();