Skip to content

Commit 40afc33

Browse files
committed
feat: route long-lived sandbox connections to the owning gateway replica
Signed-off-by: divesh <dgude@nvidia.com>
1 parent 0f9cdb8 commit 40afc33

7 files changed

Lines changed: 387 additions & 74 deletions

File tree

‎crates/openshell-cli/src/run.rs‎

Lines changed: 77 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,9 @@ pub use crate::commands::provider_readiness::{ProviderWaitOptions, sandbox_provi
3737

3838
use crate::color::Colorize;
3939
use crate::policy_update::build_policy_update_plan;
40-
use crate::tls::{TlsOptions, grpc_client};
40+
use crate::tls::{
41+
TlsOptions, grpc_client, owner_replica, retry_unrouted, routed_to, with_owner_hint,
42+
};
4143
use futures::StreamExt;
4244
use miette::{IntoDiagnostic, Result, WrapErr, miette};
4345
use openshell_bootstrap::{
@@ -71,6 +73,7 @@ use std::io::{ErrorKind, IsTerminal, Read, Write};
7173
use std::path::{Path, PathBuf};
7274
use std::process::Command;
7375
use std::time::{Duration, Instant};
76+
use tonic::metadata::AsciiMetadataValue;
7477
use tonic::{Code, Status};
7578

7679
const PROVISIONAL_CONTAINER_EXIT_RECONCILIATION_TIMEOUT: Duration = Duration::from_secs(5);
@@ -1990,15 +1993,17 @@ pub async fn sandbox_exec_grpc(
19901993
let mut client = grpc_client(server, tls).await?;
19911994

19921995
// Resolve sandbox name to id.
1993-
let sandbox = client
1994-
.get_sandbox(GetSandboxRequest {
1996+
let response = client
1997+
.get_sandbox(with_owner_hint(GetSandboxRequest {
19951998
name: name.to_string(),
19961999
workspace_scope: Some(openshell_core::proto::workspace_selector(
19972000
workspace.to_string(),
19982001
)),
1999-
})
2002+
}))
20002003
.await
2001-
.into_diagnostic()?
2004+
.into_diagnostic()?;
2005+
let replica = owner_replica(&response);
2006+
let sandbox = response
20022007
.into_inner()
20032008
.sandbox
20042009
.ok_or_else(|| miette::miette!("sandbox not found"))?;
@@ -2086,16 +2091,26 @@ pub async fn sandbox_exec_grpc(
20862091
stdin_is_terminal,
20872092
std::mem::take(&mut request.stdin),
20882093
stdin_rest,
2094+
replica.as_ref(),
20892095
)
20902096
.await;
20912097
}
20922098

20932099
// Make the streaming gRPC call.
2094-
let mut stream = client
2095-
.exec_sandbox(request)
2096-
.await
2097-
.into_diagnostic()?
2098-
.into_inner();
2100+
let retry = replica.is_some().then(|| request.clone());
2101+
let mut stream = match (
2102+
client
2103+
.exec_sandbox(routed_to(request, replica.as_ref()))
2104+
.await,
2105+
retry,
2106+
) {
2107+
(Err(status), Some(request)) if retry_unrouted(replica.as_ref(), &status) => {
2108+
client.exec_sandbox(request).await
2109+
}
2110+
(response, _) => response,
2111+
}
2112+
.into_diagnostic()?
2113+
.into_inner();
20992114

21002115
// Stream output to terminal in real-time.
21012116
let mut exit_code = 0i32;
@@ -2148,7 +2163,7 @@ pub async fn service_forward_tcp(
21482163
let (bind_addr, bind_port) = parse_tcp_forward_spec(local, target_port)?;
21492164
let mut client = grpc_client(server, tls).await?;
21502165

2151-
fetch_ready_sandbox_for_forward(&mut client, name, workspace).await?;
2166+
let mut replica = fetch_ready_sandbox_for_forward(&mut client, name, workspace).await?;
21522167

21532168
let listener = tokio::net::TcpListener::bind((bind_addr.as_str(), bind_port))
21542169
.await
@@ -2179,7 +2194,7 @@ pub async fn service_forward_tcp(
21792194
}
21802195

21812196
_ = health_check.tick() => {
2182-
fetch_ready_sandbox_for_forward(&mut client, name, workspace).await?;
2197+
replica = fetch_ready_sandbox_for_forward(&mut client, name, workspace).await?;
21832198
}
21842199

21852200
accepted = listener.accept() => {
@@ -2193,6 +2208,7 @@ pub async fn service_forward_tcp(
21932208
let target_host = target_host.to_string();
21942209
let service_id = format!("service-forward:{name}:{target_host}:{target_port}");
21952210
let fatal_tx = fatal_tx.clone();
2211+
let replica = replica.clone();
21962212
tokio::spawn(async move {
21972213
let token = match create_forward_session_token(
21982214
&mut client,
@@ -2217,6 +2233,7 @@ pub async fn service_forward_tcp(
22172233
target_port,
22182234
service_id,
22192235
token.clone(),
2236+
replica.as_ref(),
22202237
)
22212238
.await
22222239
{
@@ -2251,18 +2268,19 @@ async fn create_forward_session_token(
22512268
Ok(response.into_inner().token)
22522269
}
22532270

2271+
/// Confirm the sandbox is still ready, returning the replica that owns it.
22542272
async fn fetch_ready_sandbox_for_forward(
22552273
client: &mut crate::tls::GrpcClient,
22562274
name: &str,
22572275
workspace: &str,
2258-
) -> Result<Sandbox> {
2276+
) -> Result<Option<AsciiMetadataValue>> {
22592277
let response = match client
2260-
.get_sandbox(GetSandboxRequest {
2278+
.get_sandbox(with_owner_hint(GetSandboxRequest {
22612279
name: name.to_string(),
22622280
workspace_scope: Some(openshell_core::proto::workspace_selector(
22632281
workspace.to_string(),
22642282
)),
2265-
})
2283+
}))
22662284
.await
22672285
{
22682286
Ok(response) => response,
@@ -2274,6 +2292,7 @@ async fn fetch_ready_sandbox_for_forward(
22742292
Err(status) => return Err(status).into_diagnostic(),
22752293
};
22762294

2295+
let replica = owner_replica(&response);
22772296
let sandbox = response
22782297
.into_inner()
22792298
.sandbox
@@ -2287,7 +2306,7 @@ async fn fetch_ready_sandbox_for_forward(
22872306
));
22882307
}
22892308

2290-
Ok(sandbox)
2309+
Ok(replica)
22912310
}
22922311

22932312
#[derive(Debug)]
@@ -2353,11 +2372,9 @@ async fn forward_one_tcp_connection(
23532372
target_port: u16,
23542373
service_id: String,
23552374
authorization_token: String,
2375+
replica: Option<&AsciiMetadataValue>,
23562376
) -> std::result::Result<(), ForwardTcpConnectionError> {
2357-
use tokio_stream::wrappers::ReceiverStream;
2358-
2359-
let (tx, rx) = tokio::sync::mpsc::channel::<TcpForwardFrame>(16);
2360-
tx.send(TcpForwardFrame {
2377+
let init = TcpForwardFrame {
23612378
payload: Some(openshell_core::proto::tcp_forward_frame::Payload::Init(
23622379
TcpForwardInit {
23632380
sandbox: sandbox_name,
@@ -2370,11 +2387,10 @@ async fn forward_one_tcp_connection(
23702387
authorization_token,
23712388
},
23722389
)),
2373-
})
2374-
.await
2375-
.map_err(|_| ForwardTcpConnectionError::transient("failed to initialize forward stream"))?;
2390+
};
2391+
let (tx, response) = crate::ssh::open_forward_tcp(client, init, replica).await;
23762392

2377-
let response = match client.forward_tcp(ReceiverStream::new(rx)).await {
2393+
let response = match response {
23782394
Ok(response) => response.into_inner(),
23792395
Err(status) => {
23802396
let err = ForwardTcpConnectionError::from_status(status);
@@ -2493,6 +2509,7 @@ async fn sandbox_exec_streaming_grpc(
24932509
stdin_is_terminal: bool,
24942510
stdin_prefix: Vec<u8>,
24952511
stdin_rest: Option<tokio::sync::mpsc::Receiver<std::io::Result<Vec<u8>>>>,
2512+
replica: Option<&AsciiMetadataValue>,
24962513
) -> Result<i32> {
24972514
#[cfg(unix)]
24982515
use openshell_core::proto::ExecSandboxWindowResize;
@@ -2505,36 +2522,44 @@ async fn sandbox_exec_streaming_grpc(
25052522
(0, 0)
25062523
};
25072524

2508-
let (input_tx, input_rx) = tokio::sync::mpsc::channel::<ExecSandboxInput>(64);
2509-
2510-
// Send the start message with exec metadata.
2511-
input_tx
2512-
.send(ExecSandboxInput {
2513-
payload: Some(exec_sandbox_input::Payload::Start(ExecSandboxRequest {
2514-
request_id: String::new(),
2515-
sandbox: sandbox.object_name().to_string(),
2516-
workspace_scope: Some(openshell_core::proto::workspace_selector(
2517-
(sandbox.object_workspace()).to_string(),
2518-
)),
2519-
command: command.to_vec(),
2520-
workdir: workdir.unwrap_or_default().to_string(),
2521-
environment: environment.clone(),
2522-
no_login_shell,
2523-
execution_timeout: proto_execution_timeout(timeout_seconds)?,
2524-
stdin: Vec::new(),
2525-
tty,
2526-
cols,
2527-
rows,
2528-
})),
2529-
})
2530-
.await
2531-
.into_diagnostic()?;
2525+
// The start message carries the exec metadata.
2526+
let start = ExecSandboxInput {
2527+
payload: Some(exec_sandbox_input::Payload::Start(ExecSandboxRequest {
2528+
request_id: String::new(),
2529+
sandbox: sandbox.object_name().to_string(),
2530+
workspace_scope: Some(openshell_core::proto::workspace_selector(
2531+
(sandbox.object_workspace()).to_string(),
2532+
)),
2533+
command: command.to_vec(),
2534+
workdir: workdir.unwrap_or_default().to_string(),
2535+
environment: environment.clone(),
2536+
no_login_shell,
2537+
execution_timeout: proto_execution_timeout(timeout_seconds)?,
2538+
stdin: Vec::new(),
2539+
tty,
2540+
cols,
2541+
rows,
2542+
})),
2543+
};
25322544

2533-
let mut stream = client
2534-
.exec_sandbox_interactive(ReceiverStream::new(input_rx))
2545+
let (mut input_tx, input_rx) = tokio::sync::mpsc::channel::<ExecSandboxInput>(64);
2546+
input_tx.send(start.clone()).await.into_diagnostic()?;
2547+
let mut stream = match client
2548+
.exec_sandbox_interactive(routed_to(ReceiverStream::new(input_rx), replica))
25352549
.await
2536-
.into_diagnostic()?
2537-
.into_inner();
2550+
{
2551+
Err(status) if retry_unrouted(replica, &status) => {
2552+
let (tx, rx) = tokio::sync::mpsc::channel::<ExecSandboxInput>(64);
2553+
tx.send(start).await.into_diagnostic()?;
2554+
input_tx = tx;
2555+
client
2556+
.exec_sandbox_interactive(ReceiverStream::new(rx))
2557+
.await
2558+
}
2559+
response => response,
2560+
}
2561+
.into_diagnostic()?
2562+
.into_inner();
25382563

25392564
// Raw mode is only appropriate for an interactive terminal, not a pipe.
25402565
let raw_guard = if tty && stdin_is_terminal {

0 commit comments

Comments
 (0)