From ed9e5a26a88b8f3e36e70ccab0d8b9c5e14f015b Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 29 Sep 2026 18:24:25 +0000 Subject: [PATCH] Clean up canceled exec-server RPC requests immediately (#49332) ## Why Canceled RPC calls could leave entries in the pending request map until further traffic or disconnect. Cleanup should happen when the call is dropped, even if no response arrives. ## What changed Use a `PendingRequestGuard` to remove each registration when its call completes or is canceled. Protect the pending map with a synchronous mutex so cleanup can run in `Drop`, replacing opportunistic pruning and explicit error-path removals. ## Testing Add a regression test that drops an in-flight call and verifies that the pending request count immediately reaches zero without further traffic. GitOrigin-RevId: 82c0629f37a2f4bb89b8d10e167e2a09765b26eb --- codex-rs/exec-server/src/rpc.rs | 80 ++++++++++++++----- .../exec-server/src/rpc_pending_request.rs | 23 ++++++ 2 files changed, 82 insertions(+), 21 deletions(-) create mode 100644 codex-rs/exec-server/src/rpc_pending_request.rs diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 3912632ece..b1ed87a76c 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -21,7 +21,6 @@ use codex_protocol::protocol::W3cTraceContext; use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value; -use tokio::sync::Mutex; use tokio::sync::OwnedSemaphorePermit; use tokio::sync::Semaphore; use tokio::sync::SemaphorePermit; @@ -61,6 +60,10 @@ pub(crate) enum RpcCallError { PendingRequestLimitExceeded { limit: usize }, } +#[path = "rpc_pending_request.rs"] +mod pending_request; +use pending_request::PendingRequestGuard; + type PendingRequest = oneshot::Sender; type BoxFuture = Pin + Send + 'static>>; type RequestRoute = Box< @@ -322,7 +325,7 @@ where pub(crate) struct RpcClient { metrics: Option, write_tx: mpsc::Sender, - pending: Arc>>, + pending: Arc>>, inbound_request_ids: Arc>>, // Shared transport state from `JsonRpcConnection`. Calls use this to fail // immediately when the socket closes, even if no JSON-RPC error response @@ -346,7 +349,7 @@ impl RpcClient { task_handles: transport_tasks, transport, } = connection; - let pending = Arc::new(Mutex::new(HashMap::::new())); + let pending = Arc::new(StdMutex::new(HashMap::::new())); let closed = Arc::new(AtomicBool::new(false)); let (event_tx, event_rx) = mpsc::channel(128); @@ -394,7 +397,7 @@ impl RpcClient { }; closed_for_reader.store(true, Ordering::Release); - drain_pending(&pending_for_reader).await; + drain_pending(&pending_for_reader); let _ = event_tx .send(RpcClientEvent::Disconnected { reason: disconnect_reason, @@ -523,7 +526,7 @@ impl RpcClient { for task in &self.transport_tasks { task.abort(); } - drain_pending(&self.pending).await; + drain_pending(&self.pending); } // Callers keep this permit until `call_inner` returns, so an executor @@ -635,24 +638,24 @@ impl RpcClient { let request_id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); let (response_tx, response_rx) = oneshot::channel(); { - let mut pending = self.pending.lock().await; + let mut pending = self + .pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); // Registering the pending request and checking disconnect must be // atomic with the reader's drain_pending path. Otherwise a call // can sneak in after the drain and wait forever. if self.closed.load(Ordering::Acquire) || *self.disconnected_rx.borrow() { return Err(RpcCallError::Closed); } - pending.retain(|_, response_tx| !response_tx.is_closed()); pending.insert(request_id.clone(), response_tx); } - let params = match serde_json::to_value(params) { - Ok(params) => params, - Err(err) => { - self.pending.lock().await.remove(&request_id); - return Err(RpcCallError::Json(err)); - } + let _pending = PendingRequestGuard { + pending: &self.pending, + request_id: request_id.clone(), }; + let params = serde_json::to_value(params).map_err(RpcCallError::Json)?; let response = self .send_request_and_wait(method, request_id, params, response_rx, call_timeout) .await?; @@ -681,7 +684,6 @@ impl RpcClient { .await .is_err() { - self.pending.lock().await.remove(&request_id); return Err(RpcCallError::Closed); } tracing::event!( @@ -702,7 +704,6 @@ impl RpcClient { RpcCallTimeout::After(call_timeout) => match timeout(call_timeout, response_rx).await { Ok(response) => response, Err(_) => { - self.pending.lock().await.remove(&request_id); return Err(RpcCallError::TimedOut { method: method.to_string(), timeout: call_timeout, @@ -717,7 +718,10 @@ impl RpcClient { #[cfg(test)] pub(crate) async fn pending_request_count(&self) -> usize { - self.pending.lock().await.len() + self.pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len() } } @@ -835,18 +839,26 @@ where } async fn handle_server_message( - pending: &Mutex>, + pending: &StdMutex>, event_tx: &mpsc::Sender, message: JSONRPCMessage, ) -> Result<(), String> { match message { JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { - if let Some(pending) = pending.lock().await.remove(&id) { + if let Some(pending) = pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&id) + { let _ = pending.send(RpcCompletion::new(Ok(result))); } } JSONRPCMessage::Error(JSONRPCError { id, error }) => { - if let Some(pending) = pending.lock().await.remove(&id) { + if let Some(pending) = pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&id) + { let _ = pending.send(RpcCompletion::new(Err(RpcCallError::Server(error)))); } } @@ -869,9 +881,11 @@ async fn handle_server_message( Ok(()) } -async fn drain_pending(pending: &Mutex>) { +fn drain_pending(pending: &StdMutex>) { let pending = { - let mut pending = pending.lock().await; + let mut pending = pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); pending .drain() .map(|(_, pending)| pending) @@ -1172,6 +1186,30 @@ mod tests { } } + #[tokio::test] + async fn rpc_client_cancellation_removes_pending_request_without_further_traffic() { + let (client_stdin, server_reader) = tokio::io::duplex(/*max_buf_size*/ 4096); + let (_server_writer, client_stdout) = tokio::io::duplex(/*max_buf_size*/ 4096); + let connection = + JsonRpcConnection::from_stdio(client_stdout, client_stdin, "test-rpc".to_string()); + let (client, _events_rx) = RpcClient::new(connection); + let params = serde_json::json!({}); + let mut call = Box::pin(client.call::<_, serde_json::Value>("waiting", ¶ms)); + let mut lines = BufReader::new(server_reader).lines(); + + tokio::select! { + message = read_jsonrpc_line(&mut lines) => { + assert!(matches!(message, JSONRPCMessage::Request(_))); + } + result = &mut call => panic!("call completed without a response: {result:?}"), + } + assert_eq!(client.pending_request_count().await, 1); + + drop(call); + + assert_eq!(client.pending_request_count().await, 0); + } + #[tokio::test] async fn rpc_client_bounds_in_flight_calls_and_preserves_cleanup() { let (outgoing_tx, outgoing_rx) = tokio::sync::mpsc::channel(/*buffer*/ 1); diff --git a/codex-rs/exec-server/src/rpc_pending_request.rs b/codex-rs/exec-server/src/rpc_pending_request.rs new file mode 100644 index 0000000000..564c78a0f4 --- /dev/null +++ b/codex-rs/exec-server/src/rpc_pending_request.rs @@ -0,0 +1,23 @@ +//! Removes client RPC registrations when their call completes or is canceled. +//! The pending map is locked only for synchronous bookkeeping, including in Drop. + +use std::collections::HashMap; +use std::sync::Mutex; + +use codex_exec_server_protocol::RequestId; + +use super::PendingRequest; + +pub(super) struct PendingRequestGuard<'a> { + pub(super) pending: &'a Mutex>, + pub(super) request_id: RequestId, +} + +impl Drop for PendingRequestGuard<'_> { + fn drop(&mut self) { + self.pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&self.request_id); + } +}