Skip to content

Commit 92907e2

Browse files
committed
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>
1 parent 837fdfd commit 92907e2

6 files changed

Lines changed: 320 additions & 25 deletions

File tree

‎crates/openshell-sdk/README.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,11 @@ endpoint. Authorization is stripped by default. Select `BearerPassthrough` only
6262
when the sandbox application validates its own bearer credential. The returned
6363
`SandboxRef::service_urls` map contains each routed URL under the same name.
6464

65+
`exec` sends the sandbox's owner replica in the `x-openshell-replica` header so
66+
an edge proxy with replica routing can skip the peer relay. The owner is cached
67+
for 30 seconds, and a call is retried once without the header if the owner's
68+
route has no ready endpoints.
69+
6570
Curated calls without a workspace argument explicitly select the `default`
6671
workspace. Cross-workspace listing uses the separate `*_all_workspaces`
6772
methods and requires Platform Admin access.

‎crates/openshell-sdk/src/client.rs‎

Lines changed: 68 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use crate::error::{Result, SdkError};
1414
use crate::pagination::{Page, Pager};
1515
use crate::raw::AuthedGrpcClient;
1616
use crate::refresh::{RefreshedToken, TokenSource};
17+
use crate::routing::{OwnerCache, owner_replica, retry_unrouted, routed_to, with_owner_hint};
1718
use crate::types::{
1819
DeleteOptions, DeletionResult, ExecOptions, ExecResult, Health, ListOptions, SandboxPhase,
1920
SandboxRef, SandboxSpec, SandboxTemplateCreateSpec, SandboxTemplateListOptions,
@@ -26,6 +27,7 @@ use std::collections::HashMap;
2627
use std::future::Future;
2728
use std::sync::Arc;
2829
use std::time::{Duration, Instant};
30+
use tonic::metadata::AsciiMetadataValue;
2931
use tonic::transport::Channel;
3032

3133
/// Async client for a single `OpenShell` gateway.
@@ -43,6 +45,7 @@ pub struct OpenShellClient {
4345
/// Live bearer slot the interceptor reads; refreshed tokens are written
4446
/// here so rotation reaches in-flight requests. `None` for non-OIDC auth.
4547
bearer_slot: Option<BearerSlot>,
48+
owners: OwnerCache,
4649
}
4750

4851
impl OpenShellClient {
@@ -60,6 +63,7 @@ impl OpenShellClient {
6063
interceptor,
6164
token_source,
6265
bearer_slot,
66+
owners: OwnerCache::default(),
6367
})
6468
}
6569

@@ -75,6 +79,7 @@ impl OpenShellClient {
7579
interceptor,
7680
token_source: None,
7781
bearer_slot,
82+
owners: OwnerCache::default(),
7883
}
7984
}
8085

@@ -625,16 +630,7 @@ impl OpenShellClient {
625630
no_login_shell: opts.no_login_shell,
626631
};
627632

628-
// Open the stream under the same OIDC-aware auth policy as unary RPCs
629-
// (proactive refresh, then one reactive retry on `Unauthenticated`).
630-
// Mid-stream rotation is out of scope; streaming retry is tracked
631-
// separately.
632-
let mut stream = self
633-
.unary(|mut grpc| {
634-
let request = request.clone();
635-
async move { grpc.exec_sandbox(request).await }
636-
})
637-
.await?;
633+
let mut stream = self.open_exec("default", request).await?;
638634

639635
let mut stdout = Vec::new();
640636
let mut stderr = Vec::new();
@@ -786,19 +782,27 @@ impl OpenShellClient {
786782
/// response, force a refresh and retry exactly once. No-op auth behaves
787783
/// as a plain single call.
788784
async fn unary<T, F, Fut>(&self, call: F) -> Result<T>
785+
where
786+
F: Fn(AuthedGrpcClient) -> Fut,
787+
Fut: Future<Output = std::result::Result<tonic::Response<T>, tonic::Status>>,
788+
{
789+
self.unary_response(call)
790+
.await
791+
.map(tonic::Response::into_inner)
792+
}
793+
794+
/// [`OpenShellClient::unary`] that keeps the response metadata.
795+
async fn unary_response<T, F, Fut>(&self, call: F) -> Result<tonic::Response<T>>
789796
where
790797
F: Fn(AuthedGrpcClient) -> Fut,
791798
Fut: Future<Output = std::result::Result<tonic::Response<T>, tonic::Status>>,
792799
{
793800
self.ensure_fresh().await?;
794801
match call(self.raw_grpc()).await {
795-
Ok(resp) => Ok(resp.into_inner()),
802+
Ok(resp) => Ok(resp),
796803
Err(status) if status.code() == tonic::Code::Unauthenticated => {
797804
if self.refresh_on_unauthorized().await? {
798-
call(self.raw_grpc())
799-
.await
800-
.map(tonic::Response::into_inner)
801-
.map_err(map_status)
805+
call(self.raw_grpc()).await.map_err(map_status)
802806
} else {
803807
Err(map_status(status))
804808
}
@@ -807,6 +811,53 @@ impl OpenShellClient {
807811
}
808812
}
809813

814+
/// Best effort: a failed lookup leaves the call unrouted.
815+
async fn owner_of(&self, workspace: &str, name: &str) -> Option<AsciiMetadataValue> {
816+
if let Some(cached) = self.owners.get(workspace, name) {
817+
return cached.0;
818+
}
819+
let response = self
820+
.unary_response(|mut grpc| {
821+
let request = with_owner_hint(proto::GetSandboxRequest {
822+
name: name.to_string(),
823+
workspace_scope: Some(proto::workspace_selector(workspace)),
824+
});
825+
async move { grpc.get_sandbox(request).await }
826+
})
827+
.await
828+
.ok()?;
829+
let owner = owner_replica(&response);
830+
self.owners.insert(workspace, name, owner.clone());
831+
owner
832+
}
833+
834+
/// Open `ExecSandbox` on the owner replica under the unary auth policy.
835+
async fn open_exec(
836+
&self,
837+
workspace: &str,
838+
request: proto::ExecSandboxRequest,
839+
) -> Result<tonic::Streaming<proto::ExecSandboxEvent>> {
840+
let owner = self.owner_of(workspace, &request.sandbox).await;
841+
let owners = &self.owners;
842+
self.unary(|mut grpc| {
843+
let request = request.clone();
844+
let owner = owner.clone();
845+
async move {
846+
match grpc
847+
.exec_sandbox(routed_to(request.clone(), owner.as_ref()))
848+
.await
849+
{
850+
Err(status) if retry_unrouted(owner.as_ref(), &status) => {
851+
owners.forget(workspace, &request.sandbox);
852+
grpc.exec_sandbox(request).await
853+
}
854+
response => response,
855+
}
856+
}
857+
})
858+
.await
859+
}
860+
810861
/// Proactive refresh: if a token source is wired and the token is within
811862
/// the refresh skew of expiry, mint a new one and store it in the live
812863
/// bearer slot. Tokens with no advertised expiry are left untouched.
@@ -1204,13 +1255,7 @@ impl WorkspaceScopedClient {
12041255
no_login_shell: opts.no_login_shell,
12051256
};
12061257

1207-
let mut stream = self
1208-
.client
1209-
.unary(|mut grpc| {
1210-
let request = request.clone();
1211-
async move { grpc.exec_sandbox(request).await }
1212-
})
1213-
.await?;
1258+
let mut stream = self.client.open_exec(&self.workspace, request).await?;
12141259

12151260
let mut stdout = Vec::new();
12161261
let mut stderr = Vec::new();
@@ -1547,6 +1592,7 @@ mod tests {
15471592
interceptor,
15481593
token_source: Some(source),
15491594
bearer_slot,
1595+
owners: OwnerCache::default(),
15501596
}
15511597
}
15521598

‎crates/openshell-sdk/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ pub mod pagination;
4040
pub mod provider_readiness;
4141
pub mod raw;
4242
pub mod refresh;
43+
mod routing;
4344
pub mod transport;
4445
pub mod types;
4546

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
//! Routes long-lived sandbox calls to the replica that owns the sandbox.
5+
6+
use openshell_core::replica_routing::{
7+
OWNER_REPLICA_HEADER, OWNER_REQUEST_HEADER, ROUTE_REPLICA_HEADER,
8+
};
9+
use std::collections::HashMap;
10+
use std::sync::{Arc, Mutex, PoisonError};
11+
use std::time::{Duration, Instant};
12+
use tonic::metadata::AsciiMetadataValue;
13+
14+
/// A stale owner is safe: a live non-owner replica relays the call.
15+
const OWNER_TTL: Duration = Duration::from_secs(30);
16+
const MAX_OWNERS: usize = 1024;
17+
18+
type Entries = HashMap<(String, String), (Instant, Option<AsciiMetadataValue>)>;
19+
20+
/// Owner the gateway last reported; `None` means it reported none.
21+
pub struct Cached(pub Option<AsciiMetadataValue>);
22+
23+
/// Reported owners keyed by workspace and sandbox name.
24+
#[derive(Clone, Default)]
25+
pub struct OwnerCache(Arc<Mutex<Entries>>);
26+
27+
impl OwnerCache {
28+
pub fn get(&self, workspace: &str, name: &str) -> Option<Cached> {
29+
let entries = self.0.lock().unwrap_or_else(PoisonError::into_inner);
30+
let (stored, owner) = entries.get(&key(workspace, name))?;
31+
(stored.elapsed() < OWNER_TTL).then(|| Cached(owner.clone()))
32+
}
33+
34+
pub fn insert(&self, workspace: &str, name: &str, owner: Option<AsciiMetadataValue>) {
35+
let mut entries = self.0.lock().unwrap_or_else(PoisonError::into_inner);
36+
if entries.len() >= MAX_OWNERS {
37+
entries.clear();
38+
}
39+
entries.insert(key(workspace, name), (Instant::now(), owner));
40+
}
41+
42+
pub fn forget(&self, workspace: &str, name: &str) {
43+
self.0
44+
.lock()
45+
.unwrap_or_else(PoisonError::into_inner)
46+
.remove(&key(workspace, name));
47+
}
48+
}
49+
50+
fn key(workspace: &str, name: &str) -> (String, String) {
51+
(workspace.to_string(), name.to_string())
52+
}
53+
54+
pub fn with_owner_hint<T>(message: T) -> tonic::Request<T> {
55+
let mut request = tonic::Request::new(message);
56+
request
57+
.metadata_mut()
58+
.insert(OWNER_REQUEST_HEADER, AsciiMetadataValue::from_static("1"));
59+
request
60+
}
61+
62+
pub fn owner_replica<T>(response: &tonic::Response<T>) -> Option<AsciiMetadataValue> {
63+
response.metadata().get(OWNER_REPLICA_HEADER).cloned()
64+
}
65+
66+
pub fn routed_to<T>(message: T, replica: Option<&AsciiMetadataValue>) -> tonic::Request<T> {
67+
let mut request = tonic::Request::new(message);
68+
if let Some(replica) = replica {
69+
request
70+
.metadata_mut()
71+
.insert(ROUTE_REPLICA_HEADER, replica.clone());
72+
}
73+
request
74+
}
75+
76+
/// Envoy Gateway answers a replica route with no ready endpoints with
77+
/// `grpc-status: 14` and no message; gateway errors always carry one.
78+
pub fn retry_unrouted(replica: Option<&AsciiMetadataValue>, status: &tonic::Status) -> bool {
79+
let message = status.message();
80+
replica.is_some()
81+
&& status.code() == tonic::Code::Unavailable
82+
&& (message.is_empty()
83+
|| message.contains("no healthy upstream")
84+
|| message.contains("connection failure"))
85+
}
86+
87+
#[cfg(test)]
88+
mod tests {
89+
use super::*;
90+
91+
#[test]
92+
fn only_unrouted_unavailable_is_retried() {
93+
let replica = AsciiMetadataValue::from_static("openshell-1");
94+
let unavailable = |message| tonic::Status::unavailable(message);
95+
assert!(retry_unrouted(Some(&replica), &unavailable("")));
96+
assert!(retry_unrouted(
97+
Some(&replica),
98+
&unavailable("no healthy upstream")
99+
));
100+
assert!(!retry_unrouted(None, &unavailable("")));
101+
assert!(!retry_unrouted(
102+
Some(&replica),
103+
&unavailable("supervisor not connected")
104+
));
105+
assert!(!retry_unrouted(
106+
Some(&replica),
107+
&tonic::Status::not_found("")
108+
));
109+
}
110+
111+
#[test]
112+
fn cache_distinguishes_unknown_from_no_owner() {
113+
let cache = OwnerCache::default();
114+
assert!(cache.get("default", "box").is_none());
115+
cache.insert("default", "box", None);
116+
assert!(cache.get("default", "box").is_some_and(|c| c.0.is_none()));
117+
assert!(cache.get("team", "box").is_none());
118+
cache.forget("default", "box");
119+
assert!(cache.get("default", "box").is_none());
120+
}
121+
}

0 commit comments

Comments
 (0)