From ddabe594e6f053f2d6d193b5768ece5f4c8c771d Mon Sep 17 00:00:00 2001 From: Alexi Christakis <167903946+alexi-openai@users.noreply.github.com> Date: Tue, 6 Oct 2026 21:11:30 +0000 Subject: [PATCH] Bound relay connection attempts and handle pongs during blocked writes (#51502) ## Why A stalled rendezvous WebSocket upgrade can prevent reconnects, while short-lived connections reset reconnect backoff too early. Blocked relay writes also prevent received pongs from clearing their heartbeat deadline. ## What changed - Bound rendezvous WebSocket connection attempts with `DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT`. - Retain exponential reconnect backoff across short-lived connections, resetting it only after a connection lasts at least 120 seconds. - Process pongs during blocked data writes while preserving the in-flight send and its independent write timeout. Prefer queued pongs over an expired heartbeat deadline. - Abort the relay writer when its owning connection task is dropped. ## Testing Add coverage for pongs arriving during blocked writes, missing and queued pongs, independent write timeouts, and writer cancellation. Extend reconnect tests to cover stalled WebSocket upgrades without re-registration and backoff reset after a stable session. GitOrigin-RevId: 0efa350c2b4ada9b89e23df8909063fbb69f33dc --- codex-rs/exec-server/src/relay.rs | 99 ++++------------ codex-rs/exec-server/src/relay_writer.rs | 107 ++++++++++++++++++ .../exec-server/src/relay_writer_tests.rs | 98 ++++++++++++++++ codex-rs/exec-server/src/remote.rs | 26 +++-- .../exec-server/src/remote/noise_tests.rs | 32 ++++-- .../src/remote/reconnect_backoff.rs | 33 +++++- .../src/remote/reconnect_backoff_tests.rs | 23 ++++ 7 files changed, 320 insertions(+), 98 deletions(-) create mode 100644 codex-rs/exec-server/src/relay_writer.rs create mode 100644 codex-rs/exec-server/src/relay_writer_tests.rs diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 6af5ee9cb9..a3fcd95a4a 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -15,6 +15,7 @@ use tokio::task::JoinSet; use tokio::time::Instant; use tokio::time::timeout; use tokio_tungstenite::tungstenite::Message; +use tokio_util::task::AbortOnDropHandle; use tracing::debug; use tracing::info; use tracing::warn; @@ -46,9 +47,10 @@ use crate::relay_proto::relay_message_frame; #[cfg(test)] use crate::server::ConnectionProcessor; use crate::telemetry::ExecutorRegistration; -use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT; use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT_REASON; -use crate::websocket_pong_watchdog::WebSocketPongWatchdog; + +#[path = "relay_writer.rs"] +mod writer; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128; @@ -526,77 +528,17 @@ where let executor_registration = ExecutorRegistration::new(environment_id.clone(), executor_registration_id.clone()) .map(Arc::new); - let (mut websocket_sink, mut websocket_stream) = stream.split(); - let (physical_outgoing_tx, mut physical_outgoing_rx) = - mpsc::channel::>(CHANNEL_CAPACITY); + let (websocket_sink, mut websocket_stream) = stream.split(); + let (physical_outgoing_tx, physical_outgoing_rx) = mpsc::channel::>(CHANNEL_CAPACITY); let (closed_stream_tx, mut closed_stream_rx) = mpsc::channel::(MAX_ACTIVE_NOISE_RELAY_STREAMS); - let (pong_tx, mut pong_rx) = mpsc::channel(1); + let (pong_tx, pong_rx) = mpsc::channel(1); // Use a separate writer so this loop never waits on the channel it drains. - let mut physical_writer_task = tokio::spawn(async move { - let mut keepalive = tokio::time::interval_at( - tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, - WEBSOCKET_KEEPALIVE_INTERVAL, - ); - keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut pong_watchdog = WebSocketPongWatchdog::new(WEBSOCKET_PONG_TIMEOUT); - let pong_deadline = tokio::time::sleep(WEBSOCKET_PONG_TIMEOUT); - tokio::pin!(pong_deadline); - loop { - let message = tokio::select! { - pong = pong_rx.recv() => { - let Some(()) = pong else { - break RendezvousDisconnectReason::LocalShutdown; - }; - pong_watchdog.received_pong(); - continue; - } - _ = &mut pong_deadline, if pong_watchdog.deadline().is_some() => { - match pong_rx.try_recv() { - Ok(()) => { - pong_watchdog.received_pong(); - continue; - } - Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { - break RendezvousDisconnectReason::PongTimeout; - } - Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { - break RendezvousDisconnectReason::LocalShutdown; - } - } - } - _ = keepalive.tick(), if pong_watchdog.deadline().is_none() => { - Message::Ping(Vec::new().into()) - } - encoded = physical_outgoing_rx.recv() => { - let Some(encoded) = encoded else { - break RendezvousDisconnectReason::LocalShutdown; - }; - Message::Binary(encoded.into()) - } - }; - let is_keepalive_ping = matches!(message, Message::Ping(_)); - let write_deadline = pong_watchdog.write_deadline(tokio::time::Instant::now()); - match tokio::time::timeout_at(write_deadline, websocket_sink.send(message)).await { - Ok(Ok(())) => { - if is_keepalive_ping { - pong_watchdog.ping_sent(tokio::time::Instant::now()); - if let Some(deadline) = pong_watchdog.deadline() { - pong_deadline.as_mut().reset(deadline); - } - } - } - Ok(Err(error)) => { - warn!("Noise multiplexed environment websocket write failed: {error}"); - break RendezvousDisconnectReason::WriteError; - } - Err(_) => { - warn!("Noise multiplexed environment websocket write timed out"); - break RendezvousDisconnectReason::WriteError; - } - } - } - }); + let mut physical_writer_task = AbortOnDropHandle::new(tokio::spawn(writer::run( + websocket_sink, + physical_outgoing_rx, + pong_rx, + ))); let mut streams: HashMap> = HashMap::new(); let mut pending_handshakes: HashMap = HashMap::new(); let mut validation_tasks: JoinSet = JoinSet::new(); @@ -1113,9 +1055,11 @@ mod tests { Ok(()) } - #[tokio::test] - async fn multiplexed_environment_sends_keepalive() -> anyhow::Result<()> { - let (client_websocket, mut server_websocket) = websocket_pair().await?; + #[tokio::test(start_paused = true)] + async fn multiplexed_environment_sends_keepalive_and_cancels_blocked_writer() + -> anyhow::Result<()> { + let (client_websocket, control, mut outbound_rx) = + ControlledWebSocket::new(/*write_ready*/ true); let runtime_paths = crate::ExecServerRuntimeOptions::new( std::env::current_exe()?, /*codex_linux_sandbox_exe*/ None, @@ -1130,10 +1074,17 @@ mod tests { AllowHarnessKeyValidator, )); - read_keepalive_ping(&mut server_websocket).await?; + assert!(matches!(outbound_rx.next().await, Some(Message::Ping(_)))); + control.set_write_blocked(); + control.send_inbound(Message::Pong(Vec::new().into()))?; + control.wait_for_blocked_write().await?; environment_task.abort(); let _ = environment_task.await; + assert_eq!( + timeout(Duration::from_millis(1), outbound_rx.next()).await?, + None + ); Ok(()) } diff --git a/codex-rs/exec-server/src/relay_writer.rs b/codex-rs/exec-server/src/relay_writer.rs new file mode 100644 index 0000000000..4355ba6cbf --- /dev/null +++ b/codex-rs/exec-server/src/relay_writer.rs @@ -0,0 +1,107 @@ +//! Keep executor writes bounded while processing pongs independently of write backpressure. + +use futures::Sink; +use futures::SinkExt; +use tokio::sync::mpsc; +use tokio::time::Instant; +use tokio_tungstenite::tungstenite::Message; +use tracing::warn; + +use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; +use crate::relay::RendezvousDisconnectReason; +use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT; +use crate::websocket_pong_watchdog::WebSocketPongWatchdog; + +pub(super) async fn run( + mut sink: S, + mut outgoing_rx: mpsc::Receiver>, + mut pong_rx: mpsc::Receiver<()>, +) -> RendezvousDisconnectReason +where + S: Sink + Unpin, + E: std::fmt::Display, +{ + let mut keepalive = tokio::time::interval_at( + Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, + WEBSOCKET_KEEPALIVE_INTERVAL, + ); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut watchdog = WebSocketPongWatchdog::new(WEBSOCKET_PONG_TIMEOUT); + loop { + let message = tokio::select! { + result = receive_pong(&mut watchdog, &mut pong_rx) => { + if let Err(reason) = result { + return reason; + } + continue; + } + _ = keepalive.tick(), if watchdog.deadline().is_none() => { + Message::Ping(Vec::new().into()) + } + encoded = outgoing_rx.recv() => { + let Some(encoded) = encoded else { + return RendezvousDisconnectReason::LocalShutdown; + }; + Message::Binary(encoded.into()) + } + }; + let is_keepalive_ping = matches!(message, Message::Ping(_)); + // A pong clears only its own deadline, never the fixed write budget. + let write = tokio::time::timeout(WEBSOCKET_PONG_TIMEOUT, sink.send(message)); + tokio::pin!(write); + loop { + tokio::select! { + // Keep the same send future alive when a pong arrives: cancelling + // and restarting a partially flushed send can duplicate a frame. + result = receive_pong(&mut watchdog, &mut pong_rx), if !is_keepalive_ping => { + if let Err(reason) = result { + return reason; + } + } + result = &mut write => { + match result { + Ok(Ok(())) => { + if is_keepalive_ping { + // Pongs received during the ping flush stay queued + // until the watchdog is armed here. + watchdog.ping_sent(Instant::now()); + } + break; + } + Ok(Err(error)) => { + warn!("Noise multiplexed environment websocket write failed: {error}"); + return RendezvousDisconnectReason::WriteError; + } + Err(_) => { + warn!("Noise multiplexed environment websocket write timed out"); + return RendezvousDisconnectReason::WriteError; + } + } + } + } + } + } +} + +async fn receive_pong( + watchdog: &mut WebSocketPongWatchdog, + pong_rx: &mut mpsc::Receiver<()>, +) -> Result<(), RendezvousDisconnectReason> { + let deadline = watchdog.deadline(); + tokio::select! { + biased; + // Honor a pong already read by the reader even if the timer is also ready. + pong = pong_rx.recv() => { + pong.ok_or(RendezvousDisconnectReason::LocalShutdown)?; + watchdog.received_pong(); + Ok(()) + } + _ = tokio::time::sleep_until(deadline.unwrap_or_else(Instant::now)), if deadline.is_some() => { + Err(RendezvousDisconnectReason::PongTimeout) + } + } +} + +#[cfg(test)] +#[path = "relay_writer_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/relay_writer_tests.rs b/codex-rs/exec-server/src/relay_writer_tests.rs new file mode 100644 index 0000000000..dbc021c8bc --- /dev/null +++ b/codex-rs/exec-server/src/relay_writer_tests.rs @@ -0,0 +1,98 @@ +//! Exercise heartbeat deadlines while a single physical write is stalled. + +use std::convert::Infallible; +use std::sync::Arc; +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio::sync::Notify; +use tokio::sync::mpsc; +use tokio_util::task::AbortOnDropHandle; + +use super::*; + +#[tokio::test(start_paused = true)] +async fn pong_during_write_clears_only_the_pong_deadline() { + for release_write in [true, false] { + let (outgoing_tx, outgoing_rx) = mpsc::channel(1); + let (pong_tx, pong_rx) = mpsc::channel(1); + let (observed_tx, mut observed_rx) = mpsc::unbounded_channel(); + let release = Arc::new(Notify::new()); + let completed = Arc::new(Notify::new()); + let sink = futures::sink::unfold( + (observed_tx, Arc::clone(&release), Arc::clone(&completed)), + |(observed_tx, release, completed), message| async move { + let is_data = matches!(message, Message::Binary(_)); + observed_tx.send(message).unwrap(); + if is_data { + release.notified().await; + completed.notify_one(); + } + Ok::<_, Infallible>((observed_tx, release, completed)) + }, + ); + let task = AbortOnDropHandle::new(tokio::spawn(run(Box::pin(sink), outgoing_rx, pong_rx))); + assert!(matches!(observed_rx.recv().await, Some(Message::Ping(_)))); + tokio::time::advance(WEBSOCKET_PONG_TIMEOUT * 3 / 4).await; + outgoing_tx.send(vec![42]).await.unwrap(); + assert_eq!( + observed_rx.recv().await, + Some(Message::Binary(vec![42].into())) + ); + pong_tx.send(()).await.unwrap(); + + // Let the old pong deadline expire while the write remains blocked. + tokio::time::advance(WEBSOCKET_PONG_TIMEOUT / 2).await; + tokio::task::yield_now().await; + assert!(!task.is_finished()); + if release_write { + release.notify_one(); + completed.notified().await; + drop(outgoing_tx); + drop(pong_tx); + assert_eq!( + task.await.unwrap(), + RendezvousDisconnectReason::LocalShutdown + ); + } else { + // Receiving a pong must not extend the independent write budget. + tokio::time::advance(WEBSOCKET_PONG_TIMEOUT).await; + assert_eq!(task.await.unwrap(), RendezvousDisconnectReason::WriteError); + } + } +} + +#[tokio::test(start_paused = true)] +async fn missing_pong_times_out_during_a_blocked_write() { + let (outgoing_tx, outgoing_rx) = mpsc::channel(1); + let (_pong_tx, pong_rx) = mpsc::channel(1); + let (observed_tx, mut observed_rx) = mpsc::unbounded_channel(); + let sink = futures::sink::unfold(observed_tx, |observed_tx, message| async move { + let is_data = matches!(message, Message::Binary(_)); + observed_tx.send(message).unwrap(); + if is_data { + std::future::pending::<()>().await; + } + Ok::<_, Infallible>(observed_tx) + }); + let task = AbortOnDropHandle::new(tokio::spawn(run(Box::pin(sink), outgoing_rx, pong_rx))); + assert!(matches!(observed_rx.recv().await, Some(Message::Ping(_)))); + tokio::time::advance(WEBSOCKET_PONG_TIMEOUT * 3 / 4).await; + outgoing_tx.send(vec![42]).await.unwrap(); + assert_eq!( + observed_rx.recv().await, + Some(Message::Binary(vec![42].into())) + ); + assert_eq!(task.await.unwrap(), RendezvousDisconnectReason::PongTimeout); +} + +#[tokio::test(start_paused = true)] +async fn queued_pong_wins_over_an_expired_deadline() { + let (pong_tx, mut pong_rx) = mpsc::channel(1); + let mut watchdog = WebSocketPongWatchdog::new(Duration::from_secs(1)); + watchdog.ping_sent(Instant::now()); + tokio::time::advance(Duration::from_secs(1)).await; + pong_tx.send(()).await.unwrap(); + assert_eq!(receive_pong(&mut watchdog, &mut pong_rx).await, Ok(())); + assert_eq!(watchdog.deadline(), None); +} diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 29715486da..d432853ed4 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -709,7 +709,7 @@ async fn run_remote_environment_connections( let identity = NoiseChannelIdentity::generate().map_err(|error| { ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}")) })?; - let mut backoff = Duration::from_secs(1); + let mut backoff = reconnect_backoff::ReconnectBackoff::new(); let mut response = client .register_environment_with_retry(&config.environment_id, &identity.public_key()) .await?; @@ -727,7 +727,7 @@ async fn run_remote_environment_connections( .await { Ok(websocket) => { - backoff = Duration::from_secs(1); + let connected_at = tokio::time::Instant::now(); let executor_registration_id = response.executor_registration_id.clone(); info!( noise_event = "rendezvous_connection", @@ -749,6 +749,7 @@ async fn run_remote_environment_connections( }, ) .await; + backoff.connection_closed(connected_at.elapsed()); info!( noise_event = "rendezvous_connection", noise_outcome = "disconnected", @@ -794,12 +795,7 @@ async fn run_remote_environment_connections( } } - sleep(reconnect_backoff::reconnect_delay( - backoff, - uuid::Uuid::new_v4().as_u64_pair().1, - )) - .await; - backoff = (backoff * 2).min(Duration::from_secs(30)); + sleep(backoff.next_delay(uuid::Uuid::new_v4().as_u64_pair().1)).await; } } @@ -820,7 +816,7 @@ async fn connect_rendezvous( http_client_factory: &HttpClientFactory, ) -> Result { let started_at = Instant::now(); - let result = async { + let result = tokio::time::timeout(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, async { let mut request = url.into_client_request()?; request.headers_mut().extend(current_rendezvous_headers()); request @@ -836,8 +832,16 @@ async fn connect_rendezvous( .connect(request, noise_relay_websocket_config()) .await .map(|(websocket, _)| websocket) - } - .await; + }) + .await + .unwrap_or_else(|_| { + Err(tokio_tungstenite::tungstenite::Error::Io( + std::io::Error::new( + std::io::ErrorKind::TimedOut, + "rendezvous websocket connection timed out", + ), + )) + }); let result_name = if result.is_ok() { "success" } else { "error" }; tracing::Span::current().record("result", result_name); telemetry.remote_rendezvous_completed(result_name, started_at.elapsed()); diff --git a/codex-rs/exec-server/src/remote/noise_tests.rs b/codex-rs/exec-server/src/remote/noise_tests.rs index 312bf3351c..0e9dfb15c2 100644 --- a/codex-rs/exec-server/src/remote/noise_tests.rs +++ b/codex-rs/exec-server/src/remote/noise_tests.rs @@ -186,22 +186,32 @@ async fn reconnect_reuses_registration_until_url_is_rejected() -> Result<()> { static_registry_auth_provider(), HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), )?; - let environment_task = tokio::spawn(run_remote_environment( - config, - ExecServerRuntimeOptions::new( - std::env::current_exe()?, - /*codex_linux_sandbox_exe*/ None, - )?, - )); + let environment_task = + tokio_util::task::AbortOnDropHandle::new(tokio::spawn(run_remote_environment( + config, + ExecServerRuntimeOptions::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?, + ))); let (first_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??; let mut first_websocket = accept_async(first_socket).await?; first_websocket.close(None).await?; - // An ordinary disconnect retries the same URL without registering again. - let (mut rejected_socket, _peer_addr) = + // An ordinary disconnect and a stalled upgrade both retry without registering again. + let (mut stalled_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??; let mut request = [0u8; 4096]; + assert!(stalled_socket.read(&mut request).await? > 0); + tokio::time::pause(); + tokio::time::advance(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT).await; + tokio::time::resume(); + + // Keep the stalled socket open until the executor starts its next attempt. + let (mut rejected_socket, _peer_addr) = + timeout(Duration::from_secs(5), listener.accept()).await??; + drop(stalled_socket); let _ = rejected_socket.read(&mut request).await?; rejected_socket .write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n") @@ -209,8 +219,8 @@ async fn reconnect_reuses_registration_until_url_is_rejected() -> Result<()> { rejected_socket.shutdown().await?; // The 4xx response discards the old registration before this attempt. - let (third_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??; - let _third_websocket = accept_async(third_socket).await?; + let (final_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??; + let _final_websocket = accept_async(final_socket).await?; registry.verify().await; environment_task.abort(); diff --git a/codex-rs/exec-server/src/remote/reconnect_backoff.rs b/codex-rs/exec-server/src/remote/reconnect_backoff.rs index 9f0638ae67..28780a7336 100644 --- a/codex-rs/exec-server/src/remote/reconnect_backoff.rs +++ b/codex-rs/exec-server/src/remote/reconnect_backoff.rs @@ -1,8 +1,37 @@ -//! Spread executor reconnects after a shared outage without exceeding the existing backoff cap. +//! Spread executor reconnects and retain backoff until a connection has been stable. use std::time::Duration; -pub(super) fn reconnect_delay(backoff: Duration, random_sample: u64) -> Duration { +const INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const MAX_BACKOFF: Duration = Duration::from_secs(30); +// Longer than the first keepalive interval plus its pong timeout (30s + 60s). +const STABLE_CONNECTION_DURATION: Duration = Duration::from_secs(120); + +pub(super) struct ReconnectBackoff { + delay: Duration, +} + +impl ReconnectBackoff { + pub(super) fn new() -> Self { + Self { + delay: INITIAL_BACKOFF, + } + } + + pub(super) fn connection_closed(&mut self, connected_for: Duration) { + if connected_for >= STABLE_CONNECTION_DURATION { + self.delay = INITIAL_BACKOFF; + } + } + + pub(super) fn next_delay(&mut self, random_sample: u64) -> Duration { + let delay = reconnect_delay(self.delay, random_sample); + self.delay = (self.delay * 2).min(MAX_BACKOFF); + delay + } +} + +fn reconnect_delay(backoff: Duration, random_sample: u64) -> Duration { let upper_ms = backoff.as_millis() as u64; let lower_ms = upper_ms / 2; Duration::from_millis(lower_ms + random_sample % (upper_ms - lower_ms + 1)) diff --git a/codex-rs/exec-server/src/remote/reconnect_backoff_tests.rs b/codex-rs/exec-server/src/remote/reconnect_backoff_tests.rs index 990db1cff6..e599dc675e 100644 --- a/codex-rs/exec-server/src/remote/reconnect_backoff_tests.rs +++ b/codex-rs/exec-server/src/remote/reconnect_backoff_tests.rs @@ -4,6 +4,8 @@ use std::time::Duration; use pretty_assertions::assert_eq; +use super::ReconnectBackoff; +use super::STABLE_CONNECTION_DURATION; use super::reconnect_delay; #[test] @@ -23,3 +25,24 @@ fn reconnect_jitter_stays_within_each_backoff_window() { ); } } + +#[test] +fn flapping_connections_retain_backoff_until_a_stable_session() { + let mut backoff = ReconnectBackoff::new(); + for seconds in [1, 2, 4, 8, 16, 30, 30] { + backoff.connection_closed(STABLE_CONNECTION_DURATION - Duration::from_millis(1)); + assert_eq!( + backoff.next_delay(/*random_sample*/ 0), + Duration::from_millis(seconds * 500) + ); + } + backoff.connection_closed(STABLE_CONNECTION_DURATION); + assert_eq!( + backoff.next_delay(/*random_sample*/ 0), + Duration::from_millis(500) + ); + assert_eq!( + backoff.next_delay(/*random_sample*/ 0), + Duration::from_secs(1) + ); +}