feat(sdk): route Rust SDK exec to the sandbox owner replica (#4790)

* feat(sdk): route Rust SDK exec to the sandbox owner replica

The Rust SDK now does what the CLI already does for long-lived sandbox
calls: it asks the gateway which replica owns the sandbox's supervisor
session, sends that owner in the x-openshell-replica header on
ExecSandbox, and retries once without the header when the edge proxy
reports that the owner's route has no ready endpoints.

The owner is cached per client for 30 seconds, keyed by workspace and
sandbox name, so repeated execs do not add a GetSandbox round trip. A
failed owner lookup sends the call unrouted and leaves any error to the
exec itself.

Part of #4223

Signed-off-by: divesh <dgude@nvidia.com>

* refactor(sdk): share replica routing helpers and address review feedback

Move owner_replica, with_owner_hint, routed_to and retry_unrouted into
openshell-core so the CLI and SDK share one copy. Curated get_sandbox now
caches the owner, owner lookup returns auth and availability errors, and a
routed Unavailable drops the cached owner.

Signed-off-by: divesh <dgude@nvidia.com>

---------

Signed-off-by: divesh <dgude@nvidia.com>
This commit is contained in:
Divesh
2026-10-28 21:49:05 +00:00
committed by GitHub
parent def7da876c
commit c5850d7573
8 changed files with 436 additions and 168 deletions
+4 -122
View File
@@ -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<GrpcClient> {
Ok(client.raw_grpc())
}
/// Gateway replica that owns the sandbox named in `response`, when it says.
pub fn owner_replica<T>(response: &tonic::Response<T>) -> Option<AsciiMetadataValue> {
response
.metadata()
.get(openshell_core::replica_routing::OWNER_REPLICA_HEADER)
.cloned()
}
pub fn with_owner_hint<T>(message: T) -> tonic::Request<T> {
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<T>(message: T, replica: Option<&AsciiMetadataValue>) -> tonic::Request<T> {
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();
@@ -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<T>(response: &tonic::Response<T>) -> Option<AsciiMetadataValue> {
response.metadata().get(OWNER_REPLICA_HEADER).cloned()
}
pub fn with_owner_hint<T>(message: T) -> tonic::Request<T> {
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<T>(message: T, replica: Option<&AsciiMetadataValue>) -> tonic::Request<T> {
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)
));
}
}
}
+5
View File
@@ -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.
+96 -43
View File
@@ -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<BearerSlot>,
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<SandboxRef> {
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<T, F, Fut>(&self, call: F) -> Result<T>
where
F: Fn(AuthedGrpcClient) -> Fut,
Fut: Future<Output = std::result::Result<tonic::Response<T>, tonic::Status>>,
{
self.unary_response(call)
.await
.map(tonic::Response::into_inner)
}
/// [`OpenShellClient::unary`] that keeps the response metadata.
async fn unary_response<T, F, Fut>(&self, call: F) -> Result<tonic::Response<T>>
where
F: Fn(AuthedGrpcClient) -> Fut,
Fut: Future<Output = std::result::Result<tonic::Response<T>, 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<Option<AsciiMetadataValue>> {
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<SandboxRef> {
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<tonic::Streaming<proto::ExecSandboxEvent>> {
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<SandboxRef> {
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(),
}
}
+1
View File
@@ -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;
+65
View File
@@ -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<AsciiMetadataValue>)>;
/// Owner the gateway last reported; `None` means it reported none.
pub struct Cached(pub Option<AsciiMetadataValue>);
/// Reported owners keyed by workspace and sandbox name.
#[derive(Clone, Default)]
pub struct OwnerCache(Arc<Mutex<Entries>>);
impl OwnerCache {
pub fn get(&self, workspace: &str, name: &str) -> Option<Cached> {
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<AsciiMetadataValue>) {
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());
}
}
+151 -2
View File
@@ -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<Option<proto::ListSandboxesRequest>>,
list_requests: Mutex<Vec<proto::ListSandboxesRequest>>,
last_exec_request: Mutex<Option<proto::ExecSandboxRequest>>,
owner_replica: Option<&'static str>,
owner_lookups: AtomicU32,
/// `x-openshell-replica` per `ExecSandbox` call; empty when unrouted.
exec_routes: Mutex<Vec<String>>,
routed_exec_error: Option<Status>,
last_workspace_request: Mutex<Option<String>>,
get_calls: AtomicU32,
phase_sequence: Vec<proto::SandboxPhase>,
@@ -492,6 +498,12 @@ impl OpenShell for TestOpenShell {
&self,
request: tonic::Request<proto::GetSandboxRequest>,
) -> Result<Response<proto::SandboxResponse>, 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<proto::ExecSandboxRequest>,
) -> Result<Response<Self::ExecSandboxStream>, 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<i32, Option<tonic::Code>>;
/// Run `count` execs on `my-box`; return the mock and each exit or error code.
async fn run_execs(state: MockState, count: usize) -> (Arc<MockState>, Vec<ExecOutcome>) {
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<String> {
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<AtomicU32>,
+1 -1
View File
@@ -257,7 +257,7 @@ helm upgrade <release-name> 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.