diff --git a/crates/openshell-cli/src/tls.rs b/crates/openshell-cli/src/tls.rs index 6dceb395d..d34e92de2 100644 --- a/crates/openshell-cli/src/tls.rs +++ b/crates/openshell-cli/src/tls.rs @@ -2,6 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use miette::{IntoDiagnostic, Result, WrapErr}; +pub use openshell_core::replica_routing::{ + owner_replica, retry_unrouted, routed_to, with_owner_hint, +}; use openshell_sdk::{AuthConfig, ClientConfig, ClientIdentity, OpenShellClient}; use rustls::{ RootCertStore, @@ -10,7 +13,6 @@ use rustls::{ }; use std::io::Cursor; use std::path::PathBuf; -use tonic::metadata::AsciiMetadataValue; use tonic::transport::Channel; /// Concrete gRPC client type used by all commands. @@ -318,132 +320,12 @@ pub async fn grpc_client(server: &str, tls: &TlsOptions) -> Result { Ok(client.raw_grpc()) } -/// Gateway replica that owns the sandbox named in `response`, when it says. -pub fn owner_replica(response: &tonic::Response) -> Option { - response - .metadata() - .get(openshell_core::replica_routing::OWNER_REPLICA_HEADER) - .cloned() -} - -pub fn with_owner_hint(message: T) -> tonic::Request { - let mut request = tonic::Request::new(message); - request.metadata_mut().insert( - openshell_core::replica_routing::OWNER_REQUEST_HEADER, - AsciiMetadataValue::from_static("1"), - ); - request -} - -/// Wrap `message` so the edge proxy can route it straight to `replica`. -pub fn routed_to(message: T, replica: Option<&AsciiMetadataValue>) -> tonic::Request { - let mut request = tonic::Request::new(message); - if let Some(replica) = replica { - request.metadata_mut().insert( - openshell_core::replica_routing::ROUTE_REPLICA_HEADER, - replica.clone(), - ); - } - request -} - -/// A routed call is retried once without routing only when no gateway -/// received it. -/// -/// Envoy Gateway answers a replica route whose Service has no ready endpoints -/// with `grpc-status: 14` and no `grpc-message`, while gateway errors and -/// mid-request resets always carry a message. -pub fn retry_unrouted(replica: Option<&AsciiMetadataValue>, status: &tonic::Status) -> bool { - let message = status.message(); - replica.is_some() - && status.code() == tonic::Code::Unavailable - && (message.is_empty() - || message.contains("no healthy upstream") - || message.contains("connection failure")) -} - #[cfg(test)] mod tests { - use super::{ - TlsOptions, owner_replica, retry_unrouted, routed_to, sdk_client_config, with_owner_hint, - }; - use openshell_core::replica_routing::{ - OWNER_REPLICA_HEADER, OWNER_REQUEST_HEADER, ROUTE_REPLICA_HEADER, - }; + use super::{TlsOptions, sdk_client_config}; use openshell_sdk::AuthConfig; use std::io::Write; - #[test] - fn owner_hint_is_requested_explicitly() { - assert!( - with_owner_hint(()) - .metadata() - .contains_key(OWNER_REQUEST_HEADER) - ); - } - - #[test] - fn owner_replica_is_carried_to_the_routed_request() { - let mut response = tonic::Response::new(()); - response - .metadata_mut() - .insert(OWNER_REPLICA_HEADER, "openshell-1".parse().unwrap()); - let replica = owner_replica(&response); - let request = routed_to((), replica.as_ref()); - assert_eq!( - request.metadata().get(ROUTE_REPLICA_HEADER).unwrap(), - "openshell-1" - ); - } - - #[test] - fn requests_without_a_known_owner_are_not_routed() { - let replica = owner_replica(&tonic::Response::new(())); - assert!(replica.is_none()); - let request = routed_to((), replica.as_ref()); - assert!(request.metadata().get(ROUTE_REPLICA_HEADER).is_none()); - } - - #[test] - fn only_unavailable_routed_calls_are_retried() { - let replica = "openshell-1".parse().unwrap(); - let unavailable = tonic::Status::unavailable("no healthy upstream"); - assert!(retry_unrouted(Some(&replica), &unavailable)); - assert!(!retry_unrouted(None, &unavailable)); - assert!(!retry_unrouted( - Some(&replica), - &tonic::Status::permission_denied("denied") - )); - } - - #[test] - fn calls_are_retried_only_when_no_gateway_received_them() { - let replica = "openshell-1".parse().unwrap(); - for never_arrived in [ - "", - "no healthy upstream", - "upstream connect error or disconnect/reset before headers. reset reason: connection failure, transport failure reason: delayed connect error: Connection refused", - ] { - assert!(retry_unrouted( - Some(&replica), - &tonic::Status::unavailable(never_arrived) - )); - } - for may_have_started in [ - "supervisor session not connected", - "upstream connect error or disconnect/reset before headers. reset reason: connection termination", - ] { - assert!(!retry_unrouted( - Some(&replica), - &tonic::Status::unavailable(may_have_started) - )); - } - assert!(!retry_unrouted( - None, - &tonic::Status::unavailable("no healthy upstream") - )); - } - #[test] fn sdk_config_carries_cli_mtls_materials() { let mut ca = tempfile::NamedTempFile::new().unwrap(); diff --git a/crates/openshell-core/src/replica_routing.rs b/crates/openshell-core/src/replica_routing.rs index 9b389182f..537f1b3f8 100644 --- a/crates/openshell-core/src/replica_routing.rs +++ b/crates/openshell-core/src/replica_routing.rs @@ -4,6 +4,8 @@ //! Metadata that lets an edge proxy send a sandbox's long-lived connections //! straight to the gateway replica holding its supervisor session. +use tonic::metadata::AsciiMetadataValue; + /// Response metadata naming the replica that owns a sandbox's supervisor session. pub const OWNER_REPLICA_HEADER: &str = "x-openshell-owner"; @@ -12,3 +14,114 @@ pub const OWNER_REQUEST_HEADER: &str = "x-openshell-want-owner"; /// Request metadata asking the edge proxy to route to a specific replica. pub const ROUTE_REPLICA_HEADER: &str = "x-openshell-replica"; + +/// Gateway replica that owns the sandbox named in `response`, when it says. +pub fn owner_replica(response: &tonic::Response) -> Option { + response.metadata().get(OWNER_REPLICA_HEADER).cloned() +} + +pub fn with_owner_hint(message: T) -> tonic::Request { + let mut request = tonic::Request::new(message); + request + .metadata_mut() + .insert(OWNER_REQUEST_HEADER, AsciiMetadataValue::from_static("1")); + request +} + +/// Wrap `message` so the edge proxy can route it straight to `replica`. +pub fn routed_to(message: T, replica: Option<&AsciiMetadataValue>) -> tonic::Request { + let mut request = tonic::Request::new(message); + if let Some(replica) = replica { + request + .metadata_mut() + .insert(ROUTE_REPLICA_HEADER, replica.clone()); + } + request +} + +/// A routed call is retried once without routing only when no gateway +/// received it. +/// +/// Envoy Gateway answers a replica route whose Service has no ready endpoints +/// with `grpc-status: 14` and no `grpc-message`, while gateway errors and +/// mid-request resets always carry a message. +pub fn retry_unrouted(replica: Option<&AsciiMetadataValue>, status: &tonic::Status) -> bool { + let message = status.message(); + replica.is_some() + && status.code() == tonic::Code::Unavailable + && (message.is_empty() + || message.contains("no healthy upstream") + || message.contains("connection failure")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn owner_hint_is_requested_explicitly() { + assert!( + with_owner_hint(()) + .metadata() + .contains_key(OWNER_REQUEST_HEADER) + ); + } + + #[test] + fn owner_replica_is_carried_to_the_routed_request() { + let mut response = tonic::Response::new(()); + response + .metadata_mut() + .insert(OWNER_REPLICA_HEADER, "openshell-1".parse().unwrap()); + let replica = owner_replica(&response); + let request = routed_to((), replica.as_ref()); + assert_eq!( + request.metadata().get(ROUTE_REPLICA_HEADER).unwrap(), + "openshell-1" + ); + } + + #[test] + fn requests_without_a_known_owner_are_not_routed() { + let replica = owner_replica(&tonic::Response::new(())); + assert!(replica.is_none()); + let request = routed_to((), replica.as_ref()); + assert!(request.metadata().get(ROUTE_REPLICA_HEADER).is_none()); + } + + #[test] + fn only_unavailable_routed_calls_are_retried() { + let replica = "openshell-1".parse().unwrap(); + let unavailable = tonic::Status::unavailable("no healthy upstream"); + assert!(retry_unrouted(Some(&replica), &unavailable)); + assert!(!retry_unrouted(None, &unavailable)); + assert!(!retry_unrouted( + Some(&replica), + &tonic::Status::permission_denied("denied") + )); + } + + #[test] + fn calls_are_retried_only_when_no_gateway_received_them() { + let replica = "openshell-1".parse().unwrap(); + for never_arrived in [ + "", + "no healthy upstream", + "upstream connect error or disconnect/reset before headers. reset reason: connection failure, transport failure reason: delayed connect error: Connection refused", + ] { + assert!(retry_unrouted( + Some(&replica), + &tonic::Status::unavailable(never_arrived) + )); + } + for may_have_started in [ + "supervisor session not connected", + "upstream connect error or disconnect/reset before headers. reset reason: connection termination", + ] { + assert!(!retry_unrouted( + Some(&replica), + &tonic::Status::unavailable(may_have_started) + )); + } + } +} diff --git a/crates/openshell-sdk/README.md b/crates/openshell-sdk/README.md index 1e2e1a189..4df4cde8a 100644 --- a/crates/openshell-sdk/README.md +++ b/crates/openshell-sdk/README.md @@ -61,6 +61,11 @@ endpoint. Authorization is stripped by default. Select `BearerPassthrough` only when the sandbox application validates its own bearer credential. The returned `SandboxRef::service_urls` map contains each routed URL under the same name. +`exec` sends the sandbox's owner replica in the `x-openshell-replica` header so +an edge proxy with replica routing can skip the peer relay. The owner is cached +for 30 seconds, and a call is retried once without the header if the owner's +route has no ready endpoints. + Curated calls without a workspace argument explicitly select the `default` workspace. Cross-workspace listing uses the separate `*_all_workspaces` methods and requires Platform Admin access. diff --git a/crates/openshell-sdk/src/client.rs b/crates/openshell-sdk/src/client.rs index d0904245b..59f0dacb7 100644 --- a/crates/openshell-sdk/src/client.rs +++ b/crates/openshell-sdk/src/client.rs @@ -14,6 +14,7 @@ use crate::error::{Result, SdkError}; use crate::pagination::{Page, Pager}; use crate::raw::AuthedGrpcClient; use crate::refresh::{RefreshedToken, TokenSource}; +use crate::routing::OwnerCache; use crate::types::{ DeleteOptions, DeletionResult, ExecOptions, ExecResult, Health, ListOptions, ResumeOptions, SandboxPhase, SandboxRef, SandboxSpec, SandboxTemplateCreateSpec, SandboxTemplateListOptions, @@ -22,10 +23,12 @@ use crate::types::{ use crate::{WatchEvent, WatchOptions, transport}; use futures::{Stream, StreamExt}; use openshell_core::proto; +use openshell_core::replica_routing::{owner_replica, retry_unrouted, routed_to, with_owner_hint}; use std::collections::HashMap; use std::future::Future; use std::sync::Arc; use std::time::{Duration, Instant}; +use tonic::metadata::AsciiMetadataValue; use tonic::transport::Channel; /// Async client for a single `OpenShell` gateway. @@ -43,6 +46,7 @@ pub struct OpenShellClient { /// Live bearer slot the interceptor reads; refreshed tokens are written /// here so rotation reaches in-flight requests. `None` for non-OIDC auth. bearer_slot: Option, + owners: OwnerCache, } impl OpenShellClient { @@ -60,6 +64,7 @@ impl OpenShellClient { interceptor, token_source, bearer_slot, + owners: OwnerCache::default(), }) } @@ -75,6 +80,7 @@ impl OpenShellClient { interceptor, token_source: None, bearer_slot, + owners: OwnerCache::default(), } } @@ -299,16 +305,7 @@ impl OpenShellClient { /// Fetch a sandbox by name. pub async fn get_sandbox(&self, name: &str) -> Result { - let response = self - .unary(|mut grpc| { - let request = proto::GetSandboxRequest { - name: name.to_string(), - workspace_scope: Some(proto::workspace_selector("default")), - }; - async move { grpc.get_sandbox(request).await } - }) - .await?; - sandbox_from_response(response.sandbox) + self.get_sandbox_in("default", name).await } /// List sandboxes. @@ -661,16 +658,7 @@ impl OpenShellClient { no_login_shell: opts.no_login_shell, }; - // Open the stream under the same OIDC-aware auth policy as unary RPCs - // (proactive refresh, then one reactive retry on `Unauthenticated`). - // Mid-stream rotation is out of scope; streaming retry is tracked - // separately. - let mut stream = self - .unary(|mut grpc| { - let request = request.clone(); - async move { grpc.exec_sandbox(request).await } - }) - .await?; + let mut stream = self.open_exec("default", request).await?; let mut stdout = Vec::new(); let mut stderr = Vec::new(); @@ -822,19 +810,27 @@ impl OpenShellClient { /// response, force a refresh and retry exactly once. No-op auth behaves /// as a plain single call. async fn unary(&self, call: F) -> Result + where + F: Fn(AuthedGrpcClient) -> Fut, + Fut: Future, tonic::Status>>, + { + self.unary_response(call) + .await + .map(tonic::Response::into_inner) + } + + /// [`OpenShellClient::unary`] that keeps the response metadata. + async fn unary_response(&self, call: F) -> Result> where F: Fn(AuthedGrpcClient) -> Fut, Fut: Future, tonic::Status>>, { self.ensure_fresh().await?; match call(self.raw_grpc()).await { - Ok(resp) => Ok(resp.into_inner()), + Ok(resp) => Ok(resp), Err(status) if status.code() == tonic::Code::Unauthenticated => { if self.refresh_on_unauthorized().await? { - call(self.raw_grpc()) - .await - .map(tonic::Response::into_inner) - .map_err(map_status) + call(self.raw_grpc()).await.map_err(map_status) } else { Err(map_status(status)) } @@ -843,6 +839,78 @@ impl OpenShellClient { } } + /// Errors that exec would hit anyway are returned; any other failed lookup + /// leaves the call unrouted. + async fn owner_of(&self, workspace: &str, name: &str) -> Result> { + if let Some(cached) = self.owners.get(workspace, name) { + return Ok(cached.0); + } + match self.get_sandbox_in(workspace, name).await { + Ok(_) => Ok(self.owners.get(workspace, name).and_then(|cached| cached.0)), + Err(err) + if matches!( + err.grpc_status().map(tonic::Status::code), + Some(tonic::Code::Unauthenticated | tonic::Code::Unavailable) + ) => + { + Err(err) + } + Err(err) => { + tracing::debug!(error = %err, "sandbox owner lookup failed; exec is unrouted"); + Ok(None) + } + } + } + + /// `GetSandbox` that also caches the owner replica the gateway reports. + async fn get_sandbox_in(&self, workspace: &str, name: &str) -> Result { + let response = self + .unary_response(|mut grpc| { + let request = with_owner_hint(proto::GetSandboxRequest { + name: name.to_string(), + workspace_scope: Some(proto::workspace_selector(workspace)), + }); + async move { grpc.get_sandbox(request).await } + }) + .await?; + self.owners + .insert(workspace, name, owner_replica(&response)); + sandbox_from_response(response.into_inner().sandbox) + } + + /// Open `ExecSandbox` on the owner replica under the unary auth policy. + async fn open_exec( + &self, + workspace: &str, + request: proto::ExecSandboxRequest, + ) -> Result> { + let owner = self.owner_of(workspace, &request.sandbox).await?; + let owners = &self.owners; + self.unary(|mut grpc| { + let request = request.clone(); + let owner = owner.clone(); + async move { + match grpc + .exec_sandbox(routed_to(request.clone(), owner.as_ref())) + .await + { + Err(status) if retry_unrouted(owner.as_ref(), &status) => { + owners.forget(workspace, &request.sandbox); + grpc.exec_sandbox(request).await + } + Err(status) => { + if owner.is_some() && status.code() == tonic::Code::Unavailable { + owners.forget(workspace, &request.sandbox); + } + Err(status) + } + response => response, + } + } + }) + .await + } + /// Proactive refresh: if a token source is wired and the token is within /// the refresh skew of expiry, mint a new one and store it in the live /// bearer slot. Tokens with no advertised expiry are left untouched. @@ -1050,17 +1118,7 @@ impl WorkspaceScopedClient { /// Fetch a sandbox by name in this workspace. pub async fn get_sandbox(&self, name: &str) -> Result { - let response = self - .client - .unary(|mut grpc| { - let request = proto::GetSandboxRequest { - name: name.to_string(), - workspace_scope: Some(proto::workspace_selector(&self.workspace)), - }; - async move { grpc.get_sandbox(request).await } - }) - .await?; - sandbox_from_response(response.sandbox) + self.client.get_sandbox_in(&self.workspace, name).await } /// List sandboxes in this workspace. @@ -1279,13 +1337,7 @@ impl WorkspaceScopedClient { no_login_shell: opts.no_login_shell, }; - let mut stream = self - .client - .unary(|mut grpc| { - let request = request.clone(); - async move { grpc.exec_sandbox(request).await } - }) - .await?; + let mut stream = self.client.open_exec(&self.workspace, request).await?; let mut stdout = Vec::new(); let mut stderr = Vec::new(); @@ -1622,6 +1674,7 @@ mod tests { interceptor, token_source: Some(source), bearer_slot, + owners: OwnerCache::default(), } } diff --git a/crates/openshell-sdk/src/lib.rs b/crates/openshell-sdk/src/lib.rs index 83793bc54..bede00910 100644 --- a/crates/openshell-sdk/src/lib.rs +++ b/crates/openshell-sdk/src/lib.rs @@ -40,6 +40,7 @@ pub mod pagination; pub mod provider_readiness; pub mod raw; pub mod refresh; +mod routing; pub mod transport; pub mod types; diff --git a/crates/openshell-sdk/src/routing.rs b/crates/openshell-sdk/src/routing.rs new file mode 100644 index 000000000..95594c399 --- /dev/null +++ b/crates/openshell-sdk/src/routing.rs @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Owner replicas the gateway reported for sandboxes. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, PoisonError}; +use std::time::{Duration, Instant}; +use tonic::metadata::AsciiMetadataValue; + +/// A stale owner is safe: a live non-owner replica relays the call. +const OWNER_TTL: Duration = Duration::from_secs(30); +const MAX_OWNERS: usize = 1024; + +type Entries = HashMap<(String, String), (Instant, Option)>; + +/// Owner the gateway last reported; `None` means it reported none. +pub struct Cached(pub Option); + +/// Reported owners keyed by workspace and sandbox name. +#[derive(Clone, Default)] +pub struct OwnerCache(Arc>); + +impl OwnerCache { + pub fn get(&self, workspace: &str, name: &str) -> Option { + let entries = self.0.lock().unwrap_or_else(PoisonError::into_inner); + let (stored, owner) = entries.get(&key(workspace, name))?; + (stored.elapsed() < OWNER_TTL).then(|| Cached(owner.clone())) + } + + pub fn insert(&self, workspace: &str, name: &str, owner: Option) { + let mut entries = self.0.lock().unwrap_or_else(PoisonError::into_inner); + if entries.len() >= MAX_OWNERS { + entries.clear(); + } + entries.insert(key(workspace, name), (Instant::now(), owner)); + } + + pub fn forget(&self, workspace: &str, name: &str) { + self.0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&key(workspace, name)); + } +} + +fn key(workspace: &str, name: &str) -> (String, String) { + (workspace.to_string(), name.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cache_distinguishes_unknown_from_no_owner() { + let cache = OwnerCache::default(); + assert!(cache.get("default", "box").is_none()); + cache.insert("default", "box", None); + assert!(cache.get("default", "box").is_some_and(|c| c.0.is_none())); + assert!(cache.get("team", "box").is_none()); + cache.forget("default", "box"); + assert!(cache.get("default", "box").is_none()); + } +} diff --git a/crates/openshell-sdk/tests/client_mock.rs b/crates/openshell-sdk/tests/client_mock.rs index ce7ac3b36..b730c41d1 100644 --- a/crates/openshell-sdk/tests/client_mock.rs +++ b/crates/openshell-sdk/tests/client_mock.rs @@ -10,6 +10,7 @@ use openshell_core::proto; use openshell_core::proto::open_shell_server::{OpenShell, OpenShellServer}; +use openshell_core::replica_routing; use openshell_sdk::{ AuthConfig, ClientConfig, ExecOptions, ListOptions, OpenShellClient, Refresh, RefreshError, RefreshedToken, SandboxPhase, SandboxSpec, SandboxTemplateCreateSpec, @@ -62,6 +63,11 @@ struct MockState { last_list_request: Mutex>, list_requests: Mutex>, last_exec_request: Mutex>, + owner_replica: Option<&'static str>, + owner_lookups: AtomicU32, + /// `x-openshell-replica` per `ExecSandbox` call; empty when unrouted. + exec_routes: Mutex>, + routed_exec_error: Option, last_workspace_request: Mutex>, get_calls: AtomicU32, phase_sequence: Vec, @@ -492,6 +498,12 @@ impl OpenShell for TestOpenShell { &self, request: tonic::Request, ) -> Result, Status> { + let wants_owner = request + .metadata() + .contains_key(replica_routing::OWNER_REQUEST_HEADER); + if wants_owner { + self.state.owner_lookups.fetch_add(1, Ordering::SeqCst); + } let req = request.into_inner(); let name = req.name.clone(); *self.state.last_get_name.lock().await = Some(name.clone()); @@ -523,10 +535,17 @@ impl OpenShell for TestOpenShell { if let Some(id) = &self.state.get_sandbox_id { sandbox.metadata.as_mut().unwrap().id.clone_from(id); } - Ok(Response::new(proto::SandboxResponse { + let mut response = Response::new(proto::SandboxResponse { sandbox: Some(sandbox), service_urls: HashMap::new(), - })) + }); + if wants_owner && let Some(owner) = self.state.owner_replica { + response.metadata_mut().insert( + replica_routing::OWNER_REPLICA_HEADER, + owner.parse().unwrap(), + ); + } + Ok(response) } async fn list_sandboxes( @@ -657,6 +676,15 @@ impl OpenShell for TestOpenShell { &self, request: tonic::Request, ) -> Result, Status> { + let route = request + .metadata() + .get(replica_routing::ROUTE_REPLICA_HEADER) + .map_or_else(String::new, |value| value.to_str().unwrap().to_string()); + let routed = !route.is_empty(); + self.state.exec_routes.lock().await.push(route); + if routed && let Some(error) = &self.state.routed_exec_error { + return Err(error.clone()); + } *self.state.last_exec_request.lock().await = Some(request.into_inner()); let (tx, rx) = tokio::sync::mpsc::channel(8); tokio::spawn(async move { @@ -1833,6 +1861,127 @@ async fn exec_buffers_stdout_stderr_and_exit() { ); } +type ExecOutcome = Result>; + +/// Run `count` execs on `my-box`; return the mock and each exit or error code. +async fn run_execs(state: MockState, count: usize) -> (Arc, Vec) { + let state = Arc::new(state); + let client = connect(&start_mock(state.clone()).await).await; + let mut outcomes = Vec::new(); + for _ in 0..count { + let outcome = client + .exec("my-box", &["echo".to_string()], ExecOptions::default()) + .await; + outcomes.push( + outcome + .map(|result| result.exit_code) + .map_err(|err| err.grpc_status().map(Status::code)), + ); + } + (state, outcomes) +} + +async fn routes(state: &MockState) -> Vec { + state.exec_routes.lock().await.clone() +} + +#[tokio::test] +async fn exec_routes_to_the_reported_owner_and_reuses_it() { + let (state, outcomes) = run_execs( + MockState { + owner_replica: Some("openshell-1"), + ..Default::default() + }, + 2, + ) + .await; + assert_eq!(outcomes, [Ok(7); 2]); + assert_eq!(routes(&state).await, ["openshell-1"; 2]); + assert_eq!(state.owner_lookups.load(Ordering::SeqCst), 1); + + let (state, outcomes) = run_execs(MockState::default(), 2).await; + assert_eq!(outcomes, [Ok(7); 2]); + assert_eq!(routes(&state).await, ["", ""]); + assert_eq!(state.owner_lookups.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn exec_retries_unrouted_only_when_the_owner_route_has_no_endpoints() { + let (state, outcomes) = run_execs( + MockState { + owner_replica: Some("openshell-1"), + routed_exec_error: Some(Status::unavailable("")), + ..Default::default() + }, + 2, + ) + .await; + assert_eq!(outcomes, [Ok(7); 2]); + assert_eq!(routes(&state).await, ["openshell-1", "", "openshell-1", ""]); + assert_eq!(state.owner_lookups.load(Ordering::SeqCst), 2); + + let (state, outcomes) = run_execs( + MockState { + owner_replica: Some("openshell-1"), + routed_exec_error: Some(Status::unavailable("supervisor not connected")), + ..Default::default() + }, + 2, + ) + .await; + assert_eq!(outcomes, [Err(Some(tonic::Code::Unavailable)); 2]); + assert_eq!(routes(&state).await, ["openshell-1"; 2]); + assert_eq!(state.owner_lookups.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn exec_runs_unrouted_only_when_the_owner_lookup_fails_for_another_reason() { + let (state, outcomes) = run_execs( + MockState { + owner_replica: Some("openshell-1"), + get_error: Some(Status::internal("lookup failed")), + ..Default::default() + }, + 1, + ) + .await; + assert_eq!(outcomes, [Ok(7)]); + assert_eq!(routes(&state).await, [""]); + + let (state, outcomes) = run_execs( + MockState { + get_error: Some(Status::unavailable("gateway down")), + ..Default::default() + }, + 1, + ) + .await; + assert_eq!(outcomes, [Err(Some(tonic::Code::Unavailable))]); + assert!(routes(&state).await.is_empty()); +} + +#[tokio::test] +async fn workspace_get_sandbox_caches_the_owner_for_exec() { + let state = Arc::new(MockState { + owner_replica: Some("openshell-1"), + ..Default::default() + }); + let client = connect(&start_mock(state.clone()).await).await; + let team = client.workspace("team"); + team.get_sandbox("my-box").await.unwrap(); + let result = team + .exec("my-box", &["echo".to_string()], ExecOptions::default()) + .await + .unwrap(); + assert_eq!(result.exit_code, 7); + assert_eq!( + state.last_get_workspace.lock().await.as_deref(), + Some("team") + ); + assert_eq!(state.get_calls.load(Ordering::SeqCst), 1); + assert_eq!(routes(&state).await, ["openshell-1"]); +} + /// Refresher that hands out a fixed "fresh-token" and counts invocations. struct OneShotRefresher { calls: Arc, diff --git a/docs/kubernetes/ingress.mdx b/docs/kubernetes/ingress.mdx index b5ff06342..faf08482a 100644 --- a/docs/kubernetes/ingress.mdx +++ b/docs/kubernetes/ingress.mdx @@ -257,7 +257,7 @@ helm upgrade oci://ghcr.io/nvidia/openshell/helm-chart \ --set grpcRoute.replicaRouting.enabled=true ``` -The chart renders one Service per replica and a GRPCRoute rule that matches the `x-openshell-replica` request header. The CLI sets that header from the owner the gateway reports and retries without it if the named replica is unavailable. Replica routing requires `workload.kind=statefulset`, `workload.allowMultiReplicaStatefulSet=true`, and at most 15 replicas. +The chart renders one Service per replica and a GRPCRoute rule that matches the `x-openshell-replica` request header. The CLI and the Rust SDK set that header from the owner the gateway reports and retry without it if the named replica is unavailable. Replica routing requires `workload.kind=statefulset`, `workload.allowMultiReplicaStatefulSet=true`, and at most 15 replicas. If the release uses `workload.kind=deployment`, this upgrade replaces the Deployment with a StatefulSet, so every gateway pod restarts at once. The StatefulSet also requests a 1Gi PersistentVolumeClaim per replica, which stays `Pending` on clusters without a default StorageClass.