From 182b64ba13ee2a14329766bd819abfb9a7ebe441 Mon Sep 17 00:00:00 2001 From: Drew Newberry Date: Wed, 30 Sep 2026 16:44:55 -0700 Subject: [PATCH 1/4] fix(runtime): recover SSH relays and bound startup diagnostics Signed-off-by: Drew Newberry --- crates/openshell-cli/src/ssh.rs | 51 +++++--- .../tests/ssh_proxy_shutdown_integration.rs | 117 ++++++++++++++++++ crates/openshell-driver-docker/src/lib.rs | 11 ++ crates/openshell-driver-docker/src/tests.rs | 26 ++++ .../src/supervisor_session.rs | 41 ++++-- 5 files changed, 222 insertions(+), 24 deletions(-) create mode 100644 crates/openshell-cli/tests/ssh_proxy_shutdown_integration.rs diff --git a/crates/openshell-cli/src/ssh.rs b/crates/openshell-cli/src/ssh.rs index 8d818062ca..a7dc992e07 100644 --- a/crates/openshell-cli/src/ssh.rs +++ b/crates/openshell-cli/src/ssh.rs @@ -25,7 +25,7 @@ use std::os::unix::process::CommandExt; use std::path::{Path, PathBuf}; use std::process::{Command, ExitStatus, Stdio}; use std::time::{Duration, Instant}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::io::AsyncWriteExt; use tokio::net::TcpStream; use tokio::process::{Child, Command as TokioCommand}; use tokio_stream::wrappers::ReceiverStream; @@ -1926,20 +1926,37 @@ pub async fn sandbox_ssh_proxy( .into_diagnostic()? .into_inner(); - let stdin = tokio::io::stdin(); let stdout = tokio::io::stdout(); - let to_remote = tokio::spawn(async move { - let mut stdin = stdin; - let mut buf = vec![0u8; 64 * 1024]; - while let Ok(n) = stdin.read(&mut buf).await { - if n == 0 { - break; + // Tokio stdin uses an uncancellable read on the runtime's blocking pool. + // If the relay closes while SSH still holds the pipe open, runtime shutdown + // would wait forever for that read. A dedicated thread can be left behind + // when this ProxyCommand process exits without holding the runtime alive. + let (stdin_tx, mut stdin_rx) = tokio::sync::mpsc::channel(8); + std::thread::Builder::new() + .name("ssh-proxy-stdin".into()) + .spawn(move || { + use std::io::Read as _; + let mut stdin = std::io::stdin().lock(); + let mut buf = vec![0u8; 64 * 1024]; + loop { + match stdin.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => { + if stdin_tx.blocking_send(buf[..n].to_vec()).is_err() { + break; + } + } + } } + }) + .into_diagnostic()?; + let to_remote = tokio::spawn(async move { + while let Some(data) = stdin_rx.recv().await { if tx .send(TcpForwardFrame { payload: Some(openshell_core::proto::tcp_forward_frame::Payload::Data( - buf[..n].to_vec(), + data, )), }) .await @@ -1952,8 +1969,10 @@ pub async fn sandbox_ssh_proxy( let from_remote = tokio::spawn(async move { let mut stdout = stdout; loop { - let Ok(Some(frame)) = response.message().await else { - break; + let frame = match response.message().await { + Ok(Some(frame)) => frame, + Ok(None) => return Ok::<_, Report>(()), + Err(error) => return Err(error).into_diagnostic(), }; let Some(openshell_core::proto::tcp_forward_frame::Payload::Data(data)) = frame.payload else { @@ -1962,16 +1981,14 @@ pub async fn sandbox_ssh_proxy( if data.is_empty() { continue; } - if stdout.write_all(&data).await.is_err() { - break; - } - let _ = stdout.flush().await; + stdout.write_all(&data).await.into_diagnostic()?; + stdout.flush().await.into_diagnostic()?; } }); - let _ = from_remote.await; + let result = from_remote.await; to_remote.abort(); - Ok(()) + result.into_diagnostic()? } fn grpc_server_from_ssh_gateway_url(gateway_url: &str) -> Result { diff --git a/crates/openshell-cli/tests/ssh_proxy_shutdown_integration.rs b/crates/openshell-cli/tests/ssh_proxy_shutdown_integration.rs new file mode 100644 index 0000000000..76f9b32f5e --- /dev/null +++ b/crates/openshell-cli/tests/ssh_proxy_shutdown_integration.rs @@ -0,0 +1,117 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use bytes::Bytes; +use http_body_util::{BodyExt, StreamBody}; +use hyper::body::Frame; +use hyper::service::service_fn; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use openshell_core::proto::{TcpForwardFrame, tcp_forward_frame}; +use prost::Message; +use std::convert::Infallible; +use std::process::Stdio; +use std::time::Duration; +use tokio::io::AsyncWriteExt; + +async fn proxy_exits_with_stdin_open(relay_status: tonic::Status) { + let expected_code = relay_status.code(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let service = service_fn(move |request: hyper::Request| { + let status = relay_status.clone(); + async move { + assert_eq!(request.uri().path(), "/openshell.v1.OpenShell/ForwardTcp"); + let body = StreamBody::new(futures::stream::once(async move { + // Wait until stdin actually traverses the proxy. After this + // chunk its reader blocks again, with the parent pipe open. + let mut inbound = request.into_body(); + let mut pending = Vec::new(); + 'input: while let Some(frame) = inbound.frame().await { + if let Ok(data) = frame.unwrap().into_data() { + pending.extend_from_slice(&data); + while pending.len() >= 5 { + let length = + u32::from_be_bytes(pending[1..5].try_into().unwrap()) as usize; + if pending.len() < 5 + length { + break; + } + let frame = + TcpForwardFrame::decode(&pending[5..5 + length]).unwrap(); + pending.drain(..5 + length); + if matches!(frame.payload, Some(tcp_forward_frame::Payload::Data(data)) if data == b"probe") + { + break 'input; + } + } + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + let trailers = status.into_http::<()>().into_parts().0.headers; + Ok::, Infallible>(Frame::trailers(trailers)) + })); + let mut response = hyper::Response::new(body); + response + .headers_mut() + .insert("content-type", "application/grpc".parse().unwrap()); + Ok::<_, Infallible>(response) + } + }); + let _ = hyper::server::conn::http2::Builder::new(TokioExecutor::new()) + .serve_connection(TokioIo::new(socket), service) + .await; + }); + let config = tempfile::tempdir().unwrap(); + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_openshell")) + .args([ + "ssh-proxy", + "--gateway", + &format!("http://{address}/proxy/connect"), + "--sandbox", + "repro", + "--token", + "test-token", + ]) + .env("XDG_CONFIG_HOME", config.path()) + .env("OPENSHELL_TELEMETRY_ENABLED", "false") + .stdin(Stdio::piped()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .unwrap(); + child + .stdin + .as_mut() + .unwrap() + .write_all(b"probe") + .await + .unwrap(); + // Keep child.stdin alive throughout wait: closing it would conceal the bug. + let status = tokio::time::timeout(Duration::from_secs(5), child.wait()) + .await + .expect("SSH proxy must exit even while the parent holds stdin open") + .unwrap(); + let mut stderr = String::new(); + tokio::io::AsyncReadExt::read_to_string(child.stderr.as_mut().unwrap(), &mut stderr) + .await + .unwrap(); + if expected_code == tonic::Code::Ok { + assert!(status.success(), "{stderr}"); + } else { + assert!(!status.success(), "relay failure must propagate to SSH"); + assert!(stderr.contains("Deadline expired"), "{stderr}"); + } + server.abort(); +} + +#[tokio::test] +async fn proxy_exits_after_relay_error_with_stdin_open() { + proxy_exits_with_stdin_open(tonic::Status::deadline_exceeded("relay open timed out")).await; +} + +#[tokio::test] +async fn proxy_exits_after_clean_relay_close_with_stdin_open() { + proxy_exits_with_stdin_open(tonic::Status::ok("")).await; +} diff --git a/crates/openshell-driver-docker/src/lib.rs b/crates/openshell-driver-docker/src/lib.rs index d4a1f2f71d..1269c8ebf3 100644 --- a/crates/openshell-driver-docker/src/lib.rs +++ b/crates/openshell-driver-docker/src/lib.rs @@ -5376,6 +5376,7 @@ async fn wait_for_docker_supervisor_ready( _ if state.running == Some(false) => { let log_tail = docker_container_log_tail(docker, supervisor_id).await; let sandbox_log_tail = docker_container_log_tail(docker, sandbox_id).await; + warn!(sandbox_id, supervisor_id, supervisor_logs = %log_tail, sandbox_logs = %sandbox_log_tail, "Docker supervisor exited before becoming ready"); return Err(Status::unavailable(format!( "Docker supervisor exited before becoming ready{}{}", format_log_tail(&log_tail), @@ -5392,8 +5393,18 @@ fn format_log_tail(log_tail: &str) -> String { } fn format_named_log_tail(label: &str, log_tail: &str) -> String { + // gRPC status messages travel in HTTP/2 headers. Two 16 KiB container + // tails exceed the client's 16 KiB header budget and hide the real error + // behind PROTOCOL_ERROR. Allow for up to 3x percent-encoding expansion. + const MAX_STATUS_LOG_TAIL_BYTES: usize = 1024; if log_tail.is_empty() { String::new() + } else if log_tail.len() > MAX_STATUS_LOG_TAIL_BYTES { + let mut start = log_tail.len() - MAX_STATUS_LOG_TAIL_BYTES; + while !log_tail.is_char_boundary(start) { + start += 1; + } + format!("; {label}: [truncated] {}", &log_tail[start..]) } else { format!("; {label}: {log_tail}") } diff --git a/crates/openshell-driver-docker/src/tests.rs b/crates/openshell-driver-docker/src/tests.rs index bc747312c3..750e8d720e 100644 --- a/crates/openshell-driver-docker/src/tests.rs +++ b/crates/openshell-driver-docker/src/tests.rs @@ -24,6 +24,32 @@ use std::io::Read as _; use std::sync::Arc; use tempfile::TempDir; +#[test] +fn startup_error_log_tails_fit_grpc_header_budget() { + // Multibyte text exercises both the UTF-8 cut and worst-case gRPC message + // percent encoding. Preserve the final diagnostic from each container. + let logs = format!("{}\nstartup timed out", "🦀".repeat(8192)); + let message = format!( + "Docker supervisor exited before becoming ready{}{}", + format_log_tail(&logs), + format_named_log_tail("sandbox log tail", &logs), + ); + assert_eq!(message.matches("[truncated]").count(), 2); + assert_eq!(message.matches("startup timed out").count(), 2); + let response = Status::unavailable(message).into_http::<()>(); + let header_bytes: usize = response + .headers() + .iter() + .map(|(name, value)| name.as_str().len() + value.as_bytes().len() + 32) + .sum(); + assert!( + header_bytes < 16 * 1024, + "status headers: {header_bytes} bytes" + ); + assert_eq!(format_log_tail("small error"), "; log tail: small error"); + assert!(format_log_tail("").is_empty()); +} + fn test_launch_authentication() -> Vec { serde_json::to_vec(&SandboxLaunchAuthentication { supervisor: SupervisorAuthBundle { diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 710ecc1406..9b2f0f988a 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -933,7 +933,7 @@ impl SupervisorSessionRegistry { payload: Some(gateway_message::Payload::RelayOpen(relay_open)), }; if tx.send(msg).await.is_err() { - warn!(sandbox_id = %sandbox_id, channel_id = %channel_id, "supervisor session: failed to replay pending relay to superseding session"); + warn!(sandbox_id = %sandbox_id, channel_id = %channel_id, "supervisor session: failed to replay pending relay to connected session"); break; } } @@ -1960,12 +1960,13 @@ async fn establish_supervisor_session( } state.telemetry.sandbox_session_connected(&sandbox_id); - if superseded { - state - .supervisor_sessions - .replay_pending_relays(&sandbox_id, &tx) - .await; - } + // A disconnected session may already have removed its registration while + // an unclaimed RelayOpen remains pending. Replay on every accepted session, + // including reconnects that did not supersede a live registration. + state + .supervisor_sessions + .replay_pending_relays(&sandbox_id, &tx) + .await; // Step 4: Spawn the session loop that reads inbound messages. let state_clone = Arc::clone(&state); @@ -2920,6 +2921,32 @@ mod tests { .expect("shutdown signal should arrive at superseded session"); } + #[tokio::test] + async fn replay_pending_relays_after_disconnected_session_was_removed() { + let registry = SupervisorSessionRegistry::new(); + let (tx_old, mut rx_old) = mpsc::channel(4); + registry.register("sbx".into(), "old".into(), tx_old, make_shutdown()); + let (channel_id, relay_rx) = registry + .open_relay("sbx", Duration::from_secs(1)) + .await + .unwrap(); + rx_old.recv().await.unwrap(); + registry.remove_if_current("sbx", "old"); + + let (tx_new, mut rx_new) = mpsc::channel(4); + assert!(!registry.register("sbx".into(), "new".into(), tx_new.clone(), make_shutdown())); + registry.replay_pending_relays("sbx", &tx_new).await; + let replayed = rx_new.recv().await.unwrap(); + let Some(gateway_message::Payload::RelayOpen(open)) = replayed.payload else { + panic!("expected replayed RelayOpen"); + }; + assert_eq!(open.channel_id, channel_id); + let _claimed = registry + .claim_relay(&channel_id, Some(&sandbox_principal("sbx"))) + .unwrap(); + assert!(relay_rx.await.unwrap().is_ok()); + } + #[tokio::test] async fn replay_pending_relays_reissues_open_to_superseding_session() { let registry = SupervisorSessionRegistry::new(); From e8dc61f9d63ccb68d619350c4973df50680f982f Mon Sep 17 00:00:00 2001 From: Drew Newberry Date: Wed, 30 Sep 2026 17:54:43 -0700 Subject: [PATCH 2/4] fix(server): deliver pending relays once per supervisor session Signed-off-by: Drew Newberry --- .../src/supervisor_session.rs | 235 +++++++++++++----- 1 file changed, 178 insertions(+), 57 deletions(-) diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 9b2f0f988a..95991241a5 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -336,6 +336,8 @@ struct PendingRelay { sandbox_id: String, relay_open: RelayOpen, created_at: Instant, + /// Last session whose outbound queue received this RelayOpen. + delivered_session_id: Option, } #[derive(Debug)] @@ -474,13 +476,19 @@ impl SupervisorSessionRegistry { &self, sandbox_id: &str, timeout: Duration, - ) -> Result, Status> { + ) -> Result<(String, mpsc::Sender), Status> { let deadline = Instant::now() + timeout; let mut backoff = SESSION_WAIT_INITIAL_BACKOFF; loop { - if let Some(tx) = self.lookup_session(sandbox_id) { - return Ok(tx); + let session = self + .sessions + .lock() + .unwrap() + .get(sandbox_id) + .map(|session| (session.session_id.clone(), session.tx.clone())); + if let Some(session) = session { + return Ok(session); } if Instant::now() + backoff > deadline { return Err(Status::unavailable("supervisor session not connected")); @@ -490,6 +498,7 @@ impl SupervisorSessionRegistry { } } + #[cfg(test)] fn lookup_session(&self, sandbox_id: &str) -> Option> { self.sessions .lock() @@ -791,10 +800,6 @@ impl SupervisorSessionRegistry { if relay_open.channel_id.is_empty() { return Err(Status::invalid_argument("relay channel_id is required")); } - let tx = self - .wait_for_session(sandbox_id, session_wait_timeout) - .await?; - let channel_id = relay_open.channel_id.clone(); // Register the pending relay before sending RelayOpen to avoid a race. @@ -802,44 +807,62 @@ impl SupervisorSessionRegistry { // so two concurrent calls can't both observe "under the cap" and then // both insert past it. let (relay_tx, relay_rx) = oneshot::channel(); - { - let mut pending = self.pending_relays.lock().unwrap(); - if pending.len() >= MAX_PENDING_RELAYS { - return Err(Status::resource_exhausted(format!( - "gateway relay capacity reached ({MAX_PENDING_RELAYS} in flight)" - ))); - } - let per_sandbox = pending - .values() - .filter(|p| p.sandbox_id == sandbox_id) - .count(); - if per_sandbox >= MAX_PENDING_RELAYS_PER_SANDBOX { - return Err(Status::resource_exhausted(format!( - "per-sandbox relay limit reached ({MAX_PENDING_RELAYS_PER_SANDBOX} in flight for {sandbox_id})" - ))); + let mut relay_tx = Some(relay_tx); + loop { + let (session_id, tx) = self + .wait_for_session(sandbox_id, session_wait_timeout) + .await?; + // Reserve capacity before taking synchronous locks. The session + // may change while waiting; validate it again before insertion. + let permit = tx + .reserve() + .await + .map_err(|_| Status::unavailable("supervisor session disconnected"))?; + let sent = { + let sessions = self.sessions.lock().unwrap(); + if !sessions + .get(sandbox_id) + .is_some_and(|session| session.session_id == session_id) + { + false + } else { + let mut pending = self.pending_relays.lock().unwrap(); + if pending.len() >= MAX_PENDING_RELAYS { + return Err(Status::resource_exhausted(format!( + "gateway relay capacity reached ({MAX_PENDING_RELAYS} in flight)" + ))); + } + let per_sandbox = pending + .values() + .filter(|p| p.sandbox_id == sandbox_id) + .count(); + if per_sandbox >= MAX_PENDING_RELAYS_PER_SANDBOX { + return Err(Status::resource_exhausted(format!( + "per-sandbox relay limit reached ({MAX_PENDING_RELAYS_PER_SANDBOX} in flight for {sandbox_id})" + ))); + } + pending.insert( + channel_id.clone(), + PendingRelay { + sender: relay_tx.take().unwrap(), + sandbox_id: sandbox_id.to_string(), + relay_open: relay_open.clone(), + created_at: Instant::now(), + delivered_session_id: Some(session_id), + }, + ); + // Insertion, delivery selection, and enqueueing are atomic with + // respect to registration and replay. No await holds these locks. + permit.send(GatewayMessage { + payload: Some(gateway_message::Payload::RelayOpen(relay_open.clone())), + }); + true + } + }; + if sent { + return Ok((channel_id, relay_rx)); } - pending.insert( - channel_id.clone(), - PendingRelay { - sender: relay_tx, - sandbox_id: sandbox_id.to_string(), - relay_open: relay_open.clone(), - created_at: Instant::now(), - }, - ); - } - - let msg = GatewayMessage { - payload: Some(gateway_message::Payload::RelayOpen(relay_open)), - }; - - if tx.send(msg).await.is_err() { - // Session dropped between our lookup and send. - self.pending_relays.lock().unwrap().remove(&channel_id); - return Err(Status::unavailable("supervisor session disconnected")); } - - Ok((channel_id, relay_rx)) } pub fn fail_pending_relay(&self, channel_id: &str, error: String) -> bool { @@ -918,23 +941,43 @@ impl SupervisorSessionRegistry { self.remove(sandbox_id); } - pub async fn replay_pending_relays(&self, sandbox_id: &str, tx: &mpsc::Sender) { + pub async fn replay_pending_relays( + &self, + sandbox_id: &str, + session_id: &str, + tx: &mpsc::Sender, + ) { for channel_id in self.pending_channel_ids(sandbox_id) { - let relay_open = { + let needs_replay = { let pending = self.pending_relays.lock().unwrap(); - pending - .get(&channel_id) - .map(|pending| pending.relay_open.clone()) + pending.get(&channel_id).is_some_and(|pending| { + pending.delivered_session_id.as_deref() != Some(session_id) + }) }; - let Some(relay_open) = relay_open else { + if !needs_replay { continue; - }; - let msg = GatewayMessage { - payload: Some(gateway_message::Payload::RelayOpen(relay_open)), - }; - if tx.send(msg).await.is_err() { + } + let Ok(permit) = tx.reserve().await else { warn!(sandbox_id = %sandbox_id, channel_id = %channel_id, "supervisor session: failed to replay pending relay to connected session"); break; + }; + let sessions = self.sessions.lock().unwrap(); + if !sessions + .get(sandbox_id) + .is_some_and(|session| session.session_id == session_id) + { + break; + } + let mut pending = self.pending_relays.lock().unwrap(); + if let Some(pending) = pending.get_mut(&channel_id) + && pending.delivered_session_id.as_deref() != Some(session_id) + { + pending.delivered_session_id = Some(session_id.to_string()); + permit.send(GatewayMessage { + payload: Some(gateway_message::Payload::RelayOpen( + pending.relay_open.clone(), + )), + }); } } } @@ -1965,7 +2008,7 @@ async fn establish_supervisor_session( // including reconnects that did not supersede a live registration. state .supervisor_sessions - .replay_pending_relays(&sandbox_id, &tx) + .replay_pending_relays(&sandbox_id, &session_id, &tx) .await; // Step 4: Spawn the session loop that reads inbound messages. @@ -2497,6 +2540,7 @@ mod tests { service_id: String::new(), }, created_at, + delivered_session_id: None, } } @@ -2921,6 +2965,83 @@ mod tests { .expect("shutdown signal should arrive at superseded session"); } + #[tokio::test] + async fn replay_pending_relays_does_not_duplicate_forward_opened_after_registration() { + let registry = SupervisorSessionRegistry::new(); + let (tx_old, mut rx_old) = mpsc::channel(4); + registry.register("sbx".into(), "old".into(), tx_old, make_shutdown()); + let (old_channel, _old_relay_rx) = registry + .open_relay("sbx", Duration::from_secs(1)) + .await + .unwrap(); + rx_old.recv().await.unwrap(); + registry.remove_if_current("sbx", "old"); + + // Session establishment is paused between registration and replay. + let (tx_new, mut rx_new) = mpsc::channel(4); + registry.register("sbx".into(), "new".into(), tx_new.clone(), make_shutdown()); + let (new_channel, new_relay_rx) = registry + .open_relay("sbx", Duration::from_secs(1)) + .await + .unwrap(); + registry.replay_pending_relays("sbx", "new", &tx_new).await; + registry.replay_pending_relays("sbx", "new", &tx_new).await; + + for expected in [new_channel.clone(), old_channel] { + let Some(gateway_message::Payload::RelayOpen(open)) = + rx_new.recv().await.unwrap().payload + else { + panic!("expected RelayOpen"); + }; + assert_eq!(open.channel_id, expected); + } + assert!(matches!( + rx_new.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + let principal = sandbox_principal("sbx"); + let _claimed = registry + .claim_relay(&new_channel, Some(&principal)) + .unwrap(); + assert!(new_relay_rx.await.unwrap().is_ok()); + assert_eq!( + registry + .claim_relay(&new_channel, Some(&principal)) + .unwrap_err() + .code(), + tonic::Code::NotFound + ); + } + + #[tokio::test] + async fn open_relay_rechecks_session_after_waiting_for_queue_capacity() { + let registry = SupervisorSessionRegistry::new(); + let (tx_old, mut rx_old) = mpsc::channel(1); + registry.register("sbx".into(), "old".into(), tx_old.clone(), make_shutdown()); + tx_old.send(GatewayMessage::default()).await.unwrap(); + let open = registry.open_relay("sbx", Duration::from_secs(1)); + tokio::pin!(open); + assert!(futures_util::poll!(&mut open).is_pending()); + + let (tx_new, mut rx_new) = mpsc::channel(4); + registry.register("sbx".into(), "new".into(), tx_new.clone(), make_shutdown()); + registry.replay_pending_relays("sbx", "new", &tx_new).await; + rx_old.recv().await.unwrap(); + let (channel_id, _relay_rx) = open.await.unwrap(); + registry.replay_pending_relays("sbx", "new", &tx_new).await; + let Some(gateway_message::Payload::RelayOpen(message)) = + rx_new.recv().await.unwrap().payload + else { + panic!("expected RelayOpen on replacement session"); + }; + assert_eq!(message.channel_id, channel_id); + assert!(rx_old.try_recv().is_err()); + assert!(matches!( + rx_new.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + #[tokio::test] async fn replay_pending_relays_after_disconnected_session_was_removed() { let registry = SupervisorSessionRegistry::new(); @@ -2935,7 +3056,7 @@ mod tests { let (tx_new, mut rx_new) = mpsc::channel(4); assert!(!registry.register("sbx".into(), "new".into(), tx_new.clone(), make_shutdown())); - registry.replay_pending_relays("sbx", &tx_new).await; + registry.replay_pending_relays("sbx", "new", &tx_new).await; let replayed = rx_new.recv().await.unwrap(); let Some(gateway_message::Payload::RelayOpen(open)) = replayed.payload else { panic!("expected replayed RelayOpen"); @@ -2983,7 +3104,7 @@ mod tests { assert!(superseded); registry - .replay_pending_relays("sbx", ®istry.lookup_session("sbx").unwrap()) + .replay_pending_relays("sbx", "s-new", ®istry.lookup_session("sbx").unwrap()) .await; let replayed = rx_new From b3429877501efaa4efddfd71e0e80172e7f8457e Mon Sep 17 00:00:00 2001 From: Drew Newberry Date: Wed, 30 Sep 2026 18:14:20 -0700 Subject: [PATCH 3/4] fix(server): satisfy relay delivery clippy diagnostics Signed-off-by: Drew Newberry --- crates/openshell-server/src/supervisor_session.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 95991241a5..3c979dd422 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -336,7 +336,7 @@ struct PendingRelay { sandbox_id: String, relay_open: RelayOpen, created_at: Instant, - /// Last session whose outbound queue received this RelayOpen. + /// Last session whose outbound queue received this `RelayOpen`. delivered_session_id: Option, } @@ -820,12 +820,10 @@ impl SupervisorSessionRegistry { .map_err(|_| Status::unavailable("supervisor session disconnected"))?; let sent = { let sessions = self.sessions.lock().unwrap(); - if !sessions + if sessions .get(sandbox_id) .is_some_and(|session| session.session_id == session_id) { - false - } else { let mut pending = self.pending_relays.lock().unwrap(); if pending.len() >= MAX_PENDING_RELAYS { return Err(Status::resource_exhausted(format!( @@ -857,6 +855,8 @@ impl SupervisorSessionRegistry { payload: Some(gateway_message::Payload::RelayOpen(relay_open.clone())), }); true + } else { + false } }; if sent { @@ -962,9 +962,9 @@ impl SupervisorSessionRegistry { break; }; let sessions = self.sessions.lock().unwrap(); - if !sessions + if sessions .get(sandbox_id) - .is_some_and(|session| session.session_id == session_id) + .is_none_or(|session| session.session_id != session_id) { break; } From 32db1b210a49294bb696a0468667bede29856f33 Mon Sep 17 00:00:00 2001 From: Drew Newberry Date: Wed, 30 Sep 2026 18:21:30 -0700 Subject: [PATCH 4/4] fix(server): bound relay setup with one absolute deadline Signed-off-by: Drew Newberry --- crates/openshell-server/Cargo.toml | 1 + .../src/supervisor_session.rs | 393 ++++++++++++------ 2 files changed, 262 insertions(+), 132 deletions(-) diff --git a/crates/openshell-server/Cargo.toml b/crates/openshell-server/Cargo.toml index 54c44a7bed..fc143a9d8c 100644 --- a/crates/openshell-server/Cargo.toml +++ b/crates/openshell-server/Cargo.toml @@ -129,6 +129,7 @@ tonic-prost-build = { workspace = true } protoc-bin-vendored = { workspace = true } [dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } # Tests import the example profiles from providers/ the way an operator # would; the feature is test-only and never reaches a release binary. openshell-providers = { path = "../openshell-providers", features = ["example-profiles"] } diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 3c979dd422..42a624bce6 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -468,16 +468,15 @@ impl SupervisorSessionRegistry { None } - /// Look up the sender for a supervisor session, waiting up to `timeout` - /// for it to appear if absent. + /// Look up a supervisor session without extending the setup deadline. /// /// Uses exponential backoff (100ms → 2s) while polling the sessions map. async fn wait_for_session( &self, sandbox_id: &str, - timeout: Duration, + deadline: tokio::time::Instant, + wait_for_missing_session: bool, ) -> Result<(String, mpsc::Sender), Status> { - let deadline = Instant::now() + timeout; let mut backoff = SESSION_WAIT_INITIAL_BACKOFF; loop { @@ -490,7 +489,7 @@ impl SupervisorSessionRegistry { if let Some(session) = session { return Ok(session); } - if Instant::now() + backoff > deadline { + if !wait_for_missing_session || tokio::time::Instant::now() + backoff > deadline { return Err(Status::unavailable("supervisor session not connected")); } tokio::time::sleep(backoff).await; @@ -737,7 +736,8 @@ impl SupervisorSessionRegistry { /// blip, gateway restart, supervisor restart) and the supervisor is /// in its reconnect backoff loop /// - /// Callers pick the timeout based on how much patience the caller needs. + /// The timeout bounds session lookup, reconnect retries, and outbound queue + /// capacity waits together. Callers pick it based on their patience. /// A first `sandbox connect` right after `sandbox create` may need to /// wait for the supervisor's initial TLS + gRPC handshake (tens of /// seconds on a slow cluster), while mid-lifetime calls typically just @@ -796,73 +796,107 @@ impl SupervisorSessionRegistry { oneshot::Receiver>, ), Status, + > { + self.open_relay_with_message_until( + sandbox_id, + relay_open, + tokio::time::Instant::now() + session_wait_timeout, + true, + ) + .await + } + + /// Local routing skips waiting for a missing session, but queue capacity + /// and reconnect retries still share the routing caller's setup deadline. + async fn open_relay_with_message_until( + &self, + sandbox_id: &str, + relay_open: RelayOpen, + deadline: tokio::time::Instant, + wait_for_missing_session: bool, + ) -> Result< + ( + String, + oneshot::Receiver>, + ), + Status, > { if relay_open.channel_id.is_empty() { return Err(Status::invalid_argument("relay channel_id is required")); } - let channel_id = relay_open.channel_id.clone(); - - // Register the pending relay before sending RelayOpen to avoid a race. - // Both caps are checked and the insert happens under a single lock hold - // so two concurrent calls can't both observe "under the cap" and then - // both insert past it. - let (relay_tx, relay_rx) = oneshot::channel(); - let mut relay_tx = Some(relay_tx); - loop { - let (session_id, tx) = self - .wait_for_session(sandbox_id, session_wait_timeout) - .await?; - // Reserve capacity before taking synchronous locks. The session - // may change while waiting; validate it again before insertion. - let permit = tx - .reserve() - .await - .map_err(|_| Status::unavailable("supervisor session disconnected"))?; - let sent = { - let sessions = self.sessions.lock().unwrap(); - if sessions - .get(sandbox_id) - .is_some_and(|session| session.session_id == session_id) - { - let mut pending = self.pending_relays.lock().unwrap(); - if pending.len() >= MAX_PENDING_RELAYS { - return Err(Status::resource_exhausted(format!( - "gateway relay capacity reached ({MAX_PENDING_RELAYS} in flight)" - ))); - } - let per_sandbox = pending - .values() - .filter(|p| p.sandbox_id == sandbox_id) - .count(); - if per_sandbox >= MAX_PENDING_RELAYS_PER_SANDBOX { - return Err(Status::resource_exhausted(format!( - "per-sandbox relay limit reached ({MAX_PENDING_RELAYS_PER_SANDBOX} in flight for {sandbox_id})" - ))); + tokio::time::timeout_at(deadline, async { + let channel_id = relay_open.channel_id.clone(); + + // Register the pending relay before sending RelayOpen to avoid a race. + // Both caps are checked and the insert happens under a single lock hold + // so two concurrent calls can't both observe "under the cap" and then + // both insert past it. + let (relay_tx, relay_rx) = oneshot::channel(); + let mut relay_tx = Some(relay_tx); + loop { + if tokio::time::Instant::now() >= deadline { + return Err(Status::deadline_exceeded("supervisor relay setup timed out")); + } + let (session_id, tx) = self + .wait_for_session(sandbox_id, deadline, wait_for_missing_session) + .await?; + // Reserve capacity before taking synchronous locks. The session + // may change while waiting; validate it again before insertion. + let permit = tx + .reserve() + .await + .map_err(|_| Status::unavailable("supervisor session disconnected"))?; + if tokio::time::Instant::now() >= deadline { + return Err(Status::deadline_exceeded("supervisor relay setup timed out")); + } + let sent = { + let sessions = self.sessions.lock().unwrap(); + if sessions + .get(sandbox_id) + .is_some_and(|session| session.session_id == session_id) + { + let mut pending = self.pending_relays.lock().unwrap(); + if pending.len() >= MAX_PENDING_RELAYS { + return Err(Status::resource_exhausted(format!( + "gateway relay capacity reached ({MAX_PENDING_RELAYS} in flight)" + ))); + } + let per_sandbox = pending + .values() + .filter(|p| p.sandbox_id == sandbox_id) + .count(); + if per_sandbox >= MAX_PENDING_RELAYS_PER_SANDBOX { + return Err(Status::resource_exhausted(format!( + "per-sandbox relay limit reached ({MAX_PENDING_RELAYS_PER_SANDBOX} in flight for {sandbox_id})" + ))); + } + pending.insert( + channel_id.clone(), + PendingRelay { + sender: relay_tx.take().unwrap(), + sandbox_id: sandbox_id.to_string(), + relay_open: relay_open.clone(), + created_at: Instant::now(), + delivered_session_id: Some(session_id), + }, + ); + // Insertion, delivery selection, and enqueueing are atomic with + // respect to registration and replay. No await holds these locks. + permit.send(GatewayMessage { + payload: Some(gateway_message::Payload::RelayOpen(relay_open.clone())), + }); + true + } else { + false } - pending.insert( - channel_id.clone(), - PendingRelay { - sender: relay_tx.take().unwrap(), - sandbox_id: sandbox_id.to_string(), - relay_open: relay_open.clone(), - created_at: Instant::now(), - delivered_session_id: Some(session_id), - }, - ); - // Insertion, delivery selection, and enqueueing are atomic with - // respect to registration and replay. No await holds these locks. - permit.send(GatewayMessage { - payload: Some(gateway_message::Payload::RelayOpen(relay_open.clone())), - }); - true - } else { - false + }; + if sent { + return Ok((channel_id, relay_rx)); } - }; - if sent { - return Ok((channel_id, relay_rx)); } - } + }) + .await + .map_err(|_| Status::deadline_exceeded("supervisor relay setup timed out"))? } pub fn fail_pending_relay(&self, channel_id: &str, error: String) -> bool { @@ -1433,85 +1467,89 @@ pub async fn open_routed_relay_with_message( ), Status, > { - let deadline = Instant::now() + session_wait_timeout; - let mut backoff = SESSION_WAIT_INITIAL_BACKOFF; - let owner_index = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL); - loop { - if state.supervisor_sessions.has_session(sandbox_id) { - match state - .supervisor_sessions - .open_relay_with_message(sandbox_id, relay_open.clone(), Duration::ZERO) - .await - { - Ok(relay) => return Ok(relay), - Err(status) if status.code() == tonic::Code::Unavailable => { - // The session can migrate after `has_session` but before - // RelayOpen reaches its sender. Fall through and reread the - // persisted owner instead of surfacing a handoff race. - warn!( - sandbox_id, - error = %status, - "local supervisor relay disappeared during open; resolving owner again" - ); + let deadline = tokio::time::Instant::now() + session_wait_timeout; + tokio::time::timeout_at(deadline, async { + let mut backoff = SESSION_WAIT_INITIAL_BACKOFF; + let owner_index = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL); + loop { + if state.supervisor_sessions.has_session(sandbox_id) { + match state + .supervisor_sessions + .open_relay_with_message_until(sandbox_id, relay_open.clone(), deadline, false) + .await + { + Ok(relay) => return Ok(relay), + Err(status) if status.code() == tonic::Code::Unavailable => { + // The session can migrate after `has_session` but before + // RelayOpen reaches its sender. Fall through and reread the + // persisted owner instead of surfacing a handoff race. + warn!( + sandbox_id, + error = %status, + "local supervisor relay disappeared during open; resolving owner again" + ); + } + Err(status) => return Err(status), } - Err(status) => return Err(status), } - } - if let Some(owner) = resolve_owner(state, &owner_index, sandbox_id).await? - && owner_is_fresh(&owner) - { - if owner.owner_replica_id == state.replica_id { - warn!( - sandbox_id, - owner_replica_id = %owner.owner_replica_id, - "supervisor owner record points at this replica but no local session is registered; retrying" - ); - state.peer_routes.evict_owner(sandbox_id); - if Instant::now() + backoff > deadline { - return Err(Status::unavailable("supervisor session not connected")); - } - tokio::time::sleep(backoff).await; - backoff = (backoff * 2).min(SESSION_WAIT_MAX_BACKOFF); - continue; - } - if owner_endpoint_is_local_only(&owner.owner_peer_endpoint) { - return Err(Status::failed_precondition(format!( - "sandbox is owned by gateway replica {} which advertises no peer endpoint; \ - set OPENSHELL_PEER_ENDPOINT on every replica to route across replicas", - owner.owner_replica_id - ))); - } - match open_peer_relay( - state, - owner.owner_peer_endpoint.clone(), - sandbox_id, - relay_open.clone(), - ) - .await + if let Some(owner) = resolve_owner(state, &owner_index, sandbox_id).await? + && owner_is_fresh(&owner) { - Ok(relay) => return Ok(relay), - Err(status) => { + if owner.owner_replica_id == state.replica_id { warn!( sandbox_id, owner_replica_id = %owner.owner_replica_id, - owner_peer_endpoint = %owner.owner_peer_endpoint, - error = %status, - "gateway peer owner relay open failed; retrying until session wait timeout" + "supervisor owner record points at this replica but no local session is registered; retrying" ); - // The record may name a replaced pod, so retry against a - // fresh read rather than the cached endpoint. state.peer_routes.evict_owner(sandbox_id); + if tokio::time::Instant::now() + backoff > deadline { + return Err(Status::unavailable("supervisor session not connected")); + } + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(SESSION_WAIT_MAX_BACKOFF); + continue; + } + if owner_endpoint_is_local_only(&owner.owner_peer_endpoint) { + return Err(Status::failed_precondition(format!( + "sandbox is owned by gateway replica {} which advertises no peer endpoint; \ + set OPENSHELL_PEER_ENDPOINT on every replica to route across replicas", + owner.owner_replica_id + ))); + } + match open_peer_relay( + state, + owner.owner_peer_endpoint.clone(), + sandbox_id, + relay_open.clone(), + ) + .await + { + Ok(relay) => return Ok(relay), + Err(status) => { + warn!( + sandbox_id, + owner_replica_id = %owner.owner_replica_id, + owner_peer_endpoint = %owner.owner_peer_endpoint, + error = %status, + "gateway peer owner relay open failed; retrying until session wait timeout" + ); + // The record may name a replaced pod, so retry against a + // fresh read rather than the cached endpoint. + state.peer_routes.evict_owner(sandbox_id); + } } } - } - if Instant::now() + backoff > deadline { - return Err(Status::unavailable("supervisor session not connected")); + if tokio::time::Instant::now() + backoff > deadline { + return Err(Status::unavailable("supervisor session not connected")); + } + tokio::time::sleep(backoff).await; + backoff = (backoff * 2).min(SESSION_WAIT_MAX_BACKOFF); } - tokio::time::sleep(backoff).await; - backoff = (backoff * 2).min(SESSION_WAIT_MAX_BACKOFF); - } + }) + .await + .map_err(|_| Status::deadline_exceeded("supervisor relay setup timed out"))? } /// Reads the owning replica, reusing a recent result when one is cached. @@ -2965,6 +3003,97 @@ mod tests { .expect("shutdown signal should arrive at superseded session"); } + #[tokio::test] + async fn local_relay_attempt_skips_missing_session_without_spending_setup_budget() { + let registry = SupervisorSessionRegistry::new(); + let started = tokio::time::Instant::now(); + let attempt = registry.open_relay_with_message_until( + "missing", + RelayOpen { + channel_id: "test-channel".into(), + ..Default::default() + }, + started + Duration::from_secs(15), + false, + ); + tokio::pin!(attempt); + let std::task::Poll::Ready(Err(error)) = futures_util::poll!(&mut attempt) else { + panic!("local routing must immediately fall back when its session disappears"); + }; + assert_eq!(error.code(), tonic::Code::Unavailable); + assert!(registry.pending_relays.lock().unwrap().is_empty()); + } + + #[tokio::test(start_paused = true)] + async fn relay_setup_deadline_bounds_queue_wait_for_nonwaiting_local_attempt() { + let registry = SupervisorSessionRegistry::new(); + let (tx, mut rx) = mpsc::channel(1); + registry.register("sbx".into(), "session".into(), tx.clone(), make_shutdown()); + tx.send(GatewayMessage::default()).await.unwrap(); + let deadline = tokio::time::Instant::now() + Duration::from_secs(15); + let attempt = registry.open_relay_with_message_until( + "sbx", + RelayOpen { + channel_id: "test-channel".into(), + ..Default::default() + }, + deadline, + false, + ); + tokio::pin!(attempt); + // Skipping a missing-session wait must not imply a zero queue budget. + assert!(futures_util::poll!(&mut attempt).is_pending()); + tokio::time::advance(Duration::from_secs(15)).await; + assert_eq!( + attempt.await.unwrap_err().code(), + tonic::Code::DeadlineExceeded + ); + assert!(registry.pending_relays.lock().unwrap().is_empty()); + rx.recv().await.unwrap(); + assert!(matches!( + rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + } + + #[tokio::test(start_paused = true)] + async fn relay_setup_deadline_is_not_reset_by_reconnects() { + let registry = SupervisorSessionRegistry::new(); + let (tx, mut rx) = mpsc::channel(1); + registry.register( + "sbx".into(), + "session-0".into(), + tx.clone(), + make_shutdown(), + ); + tx.send(GatewayMessage::default()).await.unwrap(); + let started = tokio::time::Instant::now(); + let attempt = registry.open_relay("sbx", Duration::from_secs(9)); + tokio::pin!(attempt); + assert!(futures_util::poll!(&mut attempt).is_pending()); + for generation in 1..=2 { + tokio::time::advance(Duration::from_secs(3)).await; + let (new_tx, new_rx) = mpsc::channel(1); + new_tx.send(GatewayMessage::default()).await.unwrap(); + registry.register( + "sbx".into(), + format!("session-{generation}"), + new_tx, + make_shutdown(), + ); + rx.recv().await.unwrap(); + assert!(futures_util::poll!(&mut attempt).is_pending()); + rx = new_rx; + } + tokio::time::advance(Duration::from_secs(3)).await; + assert_eq!( + attempt.await.unwrap_err().code(), + tonic::Code::DeadlineExceeded + ); + assert_eq!(started.elapsed(), Duration::from_secs(9)); + assert!(registry.pending_relays.lock().unwrap().is_empty()); + } + #[tokio::test] async fn replay_pending_relays_does_not_duplicate_forward_opened_after_registration() { let registry = SupervisorSessionRegistry::new();