[Rust Frontend] Batch auto-abort requests by engine (#44591)

Signed-off-by: Hugh Ryan <197298026+HueCodes@users.noreply.github.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>

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<ClientInner>,
     mut abort_rx: mpsc::UnboundedReceiver<AbortRequest>,
 ) {
-    // 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<AbortRequest> = 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<EngineId, Vec<String>> = 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"
+                );
+            }
         }
     }
 }

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<String> = 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();