Skip to content

Commit 7085db9

Browse files
FrostGoddrew
authored andcommitted
fix(server): address HA gateway review findings (#3449)
- Let a gateway own supervisor sessions without a peer endpoint. Requiring one whenever the store is PostgreSQL broke every single-instance PostgreSQL deployment, because no sandbox supervisor could connect. A cross-replica request to an owner that advertises no endpoint now fails immediately naming the cause, instead of retrying until the wait timeout. - Close a supervisor session on heartbeat only when another replica owns it, or after renewals fail for the ownership TTL. A database error no longer drops every session heartbeating during an outage. - Clamp owner record ages at zero so a skewed or corrupt stored timestamp cannot produce a negative age. - Bound the cross-object advisory lock with a lock timeout, so a stuck holder fails instead of blocking every mutation in the fleet. - Refuse to start when a peer endpoint is configured on a multi-replica backend but peer authentication is unavailable, and warn when a multi-replica backend has no peer endpoint at all. - Reject a plaintext peer endpoint when the gateway serves TLS. - Skip the sandbox watch poller on single-replica backends, where the local update bus already sees every write. - Rate-limit the peer owner cache sweep so an insert no longer scans the whole map under the lock. - Retry GET and HEAD on a pooled upstream the sandbox closed, instead of returning 502, and drop an emptied endpoint from the pool right away. - Document the gateway peer environment variables and the post-rollout ownership skew operators should expect. Signed-off-by: divesh <dgude@nvidia.com>
1 parent f36ae35 commit 7085db9

9 files changed

Lines changed: 410 additions & 53 deletions

File tree

‎architecture/gateway.md‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,11 @@ deadline. Each retry re-reads the owner record, so a supervisor reconnect or
379379
heartbeat can surface a new owner; if no fresh reachable owner appears before
380380
the deadline, the client operation fails rather than electing an owner itself.
381381

382+
Nothing redistributes established sessions, so after a rolling restart the last
383+
surviving replica holds most sessions and a new replica serves none until
384+
sandboxes reconnect. That skew decays only as sandboxes churn. Client traffic
385+
stays correct throughout because a non-owner relays to the owner.
386+
382387
File upload and download use tar-over-SSH through the same relay path. A gateway
383388
pod termination drops the active SSH proxy byte stream, so the CLI retries the
384389
whole sync operation with a fresh SSH session instead of attempting mid-stream
@@ -393,10 +398,12 @@ also trust the chart CA, present the chart-generated client certificate for
393398
mTLS, and verify the stable gateway Service DNS name even when connecting to a
394399
Deployment pod IP.
395400

396-
`WatchSandbox` uses the local update bus for same-replica writes. One shared
397-
poller per gateway observes resource-version changes made by other replicas and
398-
feeds that bus for all local watchers, avoiding a database poll per client
399-
stream.
401+
`WatchSandbox` uses the local update bus for same-replica writes. On
402+
multi-replica backends one shared poller per gateway observes resource-version
403+
changes made by other replicas and feeds that bus for all local watchers,
404+
avoiding a database poll per client stream. SQLite deployments do not run the
405+
poller because they are single-replica and the local bus already sees every
406+
write.
400407

401408
Mutations whose invariants span sandbox, provider-profile, policy, or provider
402409
records take a process-local mutex and a shared PostgreSQL advisory lock. The

‎crates/openshell-server/src/compute/mod.rs‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3378,9 +3378,9 @@ impl ComputeRuntime {
33783378
else {
33793379
return Ok(None);
33803380
};
3381-
let age_ms = openshell_core::time::now_ms() - owner.updated_at_ms;
3382-
let ttl_ms = i64::try_from(OWNER_TTL.as_millis()).unwrap_or(i64::MAX);
3383-
Ok((age_ms < ttl_ms).then_some(owner.supervisor_instance_id))
3381+
Ok(owner
3382+
.is_fresh(OWNER_TTL)
3383+
.then_some(owner.supervisor_instance_id))
33843384
}
33853385

33863386
async fn set_supervisor_session_state(
@@ -4028,8 +4028,9 @@ impl ComputeRuntime {
40284028
}
40294029

40304030
let sandbox = decode_sandbox_record(&current_record)?;
4031-
let age_ms =
4032-
openshell_core::time::now_ms().saturating_sub(current_record.created_at_ms);
4031+
let age_ms = openshell_core::time::now_ms()
4032+
.saturating_sub(current_record.created_at_ms)
4033+
.max(0);
40334034
if age_ms < grace_ms {
40344035
return Ok(());
40354036
}

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

Lines changed: 94 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,27 @@ fn derive_peer_endpoint(config: &Config) -> Option<String> {
492492
))
493493
}
494494

495+
/// Reject a plaintext peer endpoint on a gateway that serves TLS.
496+
///
497+
/// Peer relay traffic carries whole supervisor sessions between replicas. The
498+
/// chart renders a plaintext peer endpoint only when the gateway itself serves
499+
fn validate_peer_endpoint_scheme(config: &Config, peer_endpoint: &str) -> Result<()> {
500+
if peer_endpoint.starts_with("https://") {
501+
return Ok(());
502+
}
503+
if config.tls.is_some() {
504+
return Err(Error::config(format!(
505+
"gateway peer endpoint {peer_endpoint} is plaintext but this gateway serves TLS; \
506+
set an https:// OPENSHELL_PEER_ENDPOINT so peer relay traffic is not downgraded"
507+
)));
508+
}
509+
warn!(
510+
peer_endpoint,
511+
"gateway peer relay traffic is plaintext because this gateway does not serve TLS"
512+
);
513+
Ok(())
514+
}
515+
495516
/// Run the `OpenShell` server.
496517
///
497518
/// This starts a multiplexed gRPC/HTTP server on the configured bind address.
@@ -753,6 +774,20 @@ pub(crate) async fn run_server(
753774
);
754775
}
755776

777+
let peer_routing_expected = state.peer_endpoint.is_some() && !state.store.is_single_replica();
778+
if let Some(peer_endpoint) = state.peer_endpoint.as_deref()
779+
&& peer_routing_expected
780+
{
781+
validate_peer_endpoint_scheme(&state.config, peer_endpoint)?;
782+
}
783+
if state.peer_endpoint.is_none() && !state.store.is_single_replica() {
784+
warn!(
785+
"no gateway peer endpoint configured; this replica owns its supervisor sessions but \
786+
peers cannot reach it. Single-gateway deployments are unaffected; set \
787+
OPENSHELL_PEER_ENDPOINT on every replica when running more than one."
788+
);
789+
}
790+
756791
if std::env::var_os("KUBERNETES_SERVICE_HOST").is_some() {
757792
let namespace = std::env::var("OPENSHELL_POD_NAMESPACE").ok();
758793
let service_account = std::env::var("OPENSHELL_SERVICE_ACCOUNT_NAME").ok();
@@ -787,20 +822,41 @@ pub(crate) async fn run_server(
787822
"gateway peer ServiceAccount TokenReview authentication enabled"
788823
);
789824
}
825+
Err(err) if peer_routing_expected => {
826+
return Err(Error::config(format!(
827+
"in-cluster K8s client construction failed ({err}); \
828+
gateway peer authentication is required because \
829+
OPENSHELL_PEER_ENDPOINT is configured"
830+
)));
831+
}
790832
Err(err) => warn!(
791833
error = %err,
792834
"in-cluster K8s client construction failed; \
793835
gateway peer ServiceAccount authentication is disabled"
794836
),
795837
}
796838
}
839+
_ if peer_routing_expected => {
840+
return Err(Error::config(
841+
"OPENSHELL_POD_NAMESPACE or OPENSHELL_SERVICE_ACCOUNT_NAME missing; \
842+
both are required for gateway peer authentication because \
843+
OPENSHELL_PEER_ENDPOINT is configured"
844+
.to_string(),
845+
));
846+
}
797847
_ => {
798848
debug!(
799849
"OPENSHELL_POD_NAMESPACE or OPENSHELL_SERVICE_ACCOUNT_NAME missing; \
800850
gateway peer ServiceAccount authentication disabled"
801851
);
802852
}
803853
}
854+
} else if peer_routing_expected {
855+
return Err(Error::config(
856+
"OPENSHELL_PEER_ENDPOINT is configured but the gateway is not running in a \
857+
Kubernetes cluster, so gateway peer authentication is unavailable"
858+
.to_string(),
859+
));
804860
}
805861

806862
let state = Arc::new(state);
@@ -942,13 +998,13 @@ pub(crate) async fn run_server(
942998
}
943999

9441000
startup_tx.send_replace(true);
945-
// The poller observes writes made by other replicas. Single-replica
946-
// backends already see every write through the local update bus.
1001+
// The poller exists to observe writes made by other replicas. Single-
1002+
// replica backends have none, so it would only add load.
9471003
if !store.is_single_replica() {
9481004
sandbox_watch::spawn_store_poller(
9491005
store.clone(),
9501006
state.sandbox_watch_bus.clone(),
951-
Duration::from_secs(1),
1007+
sandbox_watch::DEFAULT_STORE_POLL_INTERVAL,
9521008
shutdown_rx.clone(),
9531009
);
9541010
}
@@ -1862,7 +1918,7 @@ mod tests {
18621918
MultiplexService, ServerState, TlsAcceptor, allow_plaintext_service_http,
18631919
bind_gateway_listener, classify_initial_bytes, configured_compute_driver,
18641920
extension_token_ttl, is_benign_tls_handshake_failure, mint_gateway_extension_credential,
1865-
serve_gateway_listener,
1921+
serve_gateway_listener, validate_peer_endpoint_scheme,
18661922
};
18671923
use openshell_core::{
18681924
Config,
@@ -1882,6 +1938,40 @@ mod tests {
18821938

18831939
use crate::tls_test_utils::generate_test_certs_with_ca;
18841940

1941+
fn tls_enabled_config() -> Config {
1942+
Config::new(Some(openshell_core::TlsConfig {
1943+
cert_path: "/tmp/cert.pem".into(),
1944+
key_path: "/tmp/key.pem".into(),
1945+
client_ca_path: None,
1946+
require_client_auth: false,
1947+
external_cert_path: None,
1948+
external_key_path: None,
1949+
external_server_names: Vec::new(),
1950+
}))
1951+
}
1952+
1953+
#[test]
1954+
fn plaintext_peer_endpoint_is_rejected_on_a_tls_gateway() {
1955+
let error = validate_peer_endpoint_scheme(&tls_enabled_config(), "http://10.0.0.1:8080")
1956+
.expect_err("plaintext peer endpoint must not be accepted alongside gateway TLS");
1957+
assert!(
1958+
error.to_string().contains("plaintext"),
1959+
"error should name the downgrade: {error}"
1960+
);
1961+
}
1962+
1963+
#[test]
1964+
fn https_peer_endpoint_is_accepted_on_a_tls_gateway() {
1965+
validate_peer_endpoint_scheme(&tls_enabled_config(), "https://10.0.0.1:8080").unwrap();
1966+
}
1967+
1968+
#[test]
1969+
fn plaintext_peer_endpoint_is_allowed_on_a_plaintext_gateway() {
1970+
let config = Config::new(None);
1971+
assert!(config.tls.is_none());
1972+
validate_peer_endpoint_scheme(&config, "http://10.0.0.1:8080").unwrap();
1973+
}
1974+
18851975
static DETECTION_PROBE_ORDER: LazyLock<Mutex<Vec<&'static str>>> =
18861976
LazyLock::new(|| Mutex::new(Vec::new()));
18871977

‎crates/openshell-server/src/persistence/postgres.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,11 @@ pub struct PostgresStore {
4040
// 64-bit advisory-lock key space.
4141
const CROSS_OBJECT_ADVISORY_LOCK_KEY: i64 = 0x4f50_454e_5348_4c4c;
4242

43+
// Bounds the wait for the cross-object lock. The holder only validates and
44+
// writes, so a wait this long means a stuck replica; failing beats blocking
45+
// every sandbox and provider mutation in the fleet indefinitely.
46+
const CROSS_OBJECT_ADVISORY_LOCK_TIMEOUT: &str = "10s";
47+
4348
pub(super) struct PostgresAdvisoryLockGuard {
4449
// `close_on_drop` is set before this guard is constructed. Closing the
4550
// dedicated session releases the session-level advisory lock even when a
@@ -114,6 +119,11 @@ impl PostgresStore {
114119
) -> PersistenceResult<PostgresAdvisoryLockGuard> {
115120
let mut connection = self.pool.acquire().await.map_err(|e| map_db_error(&e))?;
116121
connection.close_on_drop();
122+
sqlx::query("SELECT set_config('lock_timeout', $1, false)")
123+
.bind(CROSS_OBJECT_ADVISORY_LOCK_TIMEOUT)
124+
.execute(&mut *connection)
125+
.await
126+
.map_err(|e| map_db_error(&e))?;
117127
sqlx::query("SELECT pg_advisory_lock($1)")
118128
.bind(CROSS_OBJECT_ADVISORY_LOCK_KEY)
119129
.execute(&mut *connection)

‎crates/openshell-server/src/sandbox_watch.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ use tonic::Status;
1313
use crate::persistence::Store;
1414
use openshell_core::proto::Sandbox;
1515

16+
/// How often [`spawn_store_poller`] rechecks watched sandboxes for writes made
17+
/// by other gateway replicas.
18+
pub const DEFAULT_STORE_POLL_INTERVAL: Duration = Duration::from_secs(1);
19+
1620
/// Broadcast bus of sandbox updates keyed by sandbox id.
1721
///
1822
/// Producers call [`SandboxWatchBus::notify`] whenever the persisted sandbox record changes.

‎crates/openshell-server/src/service_routing.rs‎

Lines changed: 85 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,14 @@ impl ServiceUpstreamPool {
100100
let mut inner = self.inner.lock().unwrap();
101101
let entries = inner.idle.get_mut(key)?;
102102
entries.retain(|entry| is_reusable(entry, now));
103-
let index = entries.iter().position(|entry| entry.sender.is_ready())?;
104-
Some(entries.swap_remove(index).sender)
103+
let taken = entries
104+
.iter()
105+
.position(|entry| entry.sender.is_ready())
106+
.map(|index| entries.swap_remove(index).sender);
107+
if entries.is_empty() {
108+
inner.idle.remove(key);
109+
}
110+
taken
105111
}
106112

107113
fn put(&self, key: &str, sender: UpstreamSender) {
@@ -435,19 +441,34 @@ async fn proxy_to_endpoint(
435441
state.service_upstreams.take(&pool_key)
436442
};
437443

444+
let reused = pooled.is_some();
438445
let mut sender = match pooled {
439446
Some(sender) => sender,
440447
None => open_upstream(&state, &sandbox, &endpoint, target_port, websocket_upgrade).await?,
441448
};
442449

443450
let upstream = build_upstream_request(req, target_port, websocket_upgrade)?;
444-
let mut response = sender.send_request(upstream).await.map_err(|err| {
445-
warn!(error = %err, "sandbox service routing: upstream HTTP request failed");
446-
state.service_upstreams.evict(&pool_key);
447-
let route_err = ServiceRouteError::service_unreachable();
448-
emit_service_relay_failure(&endpoint, target_port, route_err.reason);
449-
route_err
450-
})?;
451+
let replay = reused.then(|| replayable_request(&upstream)).flatten();
452+
let mut response = match sender.send_request(upstream).await {
453+
Ok(response) => response,
454+
Err(err) => {
455+
warn!(error = %err, "sandbox service routing: upstream HTTP request failed");
456+
state.service_upstreams.evict(&pool_key);
457+
let Some(replay) = replay else {
458+
let route_err = ServiceRouteError::service_unreachable();
459+
emit_service_relay_failure(&endpoint, target_port, route_err.reason);
460+
return Err(route_err);
461+
};
462+
sender =
463+
open_upstream(&state, &sandbox, &endpoint, target_port, websocket_upgrade).await?;
464+
sender.send_request(replay).await.map_err(|err| {
465+
warn!(error = %err, "sandbox service routing: upstream HTTP retry failed");
466+
let route_err = ServiceRouteError::service_unreachable();
467+
emit_service_relay_failure(&endpoint, target_port, route_err.reason);
468+
route_err
469+
})?
470+
}
471+
};
451472

452473
if !websocket_upgrade {
453474
state.service_upstreams.put(&pool_key, sender);
@@ -574,6 +595,23 @@ async fn load_endpoint(
574595
.ok_or_else(ServiceRouteError::endpoint_not_found)
575596
}
576597

598+
/// Copies a request that can be sent again on a fresh connection.
599+
///
600+
/// A pooled connection can be closed by the sandbox between the liveness check
601+
/// and the send. Only bodyless methods are replayable, because the original
602+
/// body is consumed by the failed attempt.
603+
fn replayable_request(request: &Request<Body>) -> Option<Request<Body>> {
604+
if !matches!(*request.method(), Method::GET | Method::HEAD) {
605+
return None;
606+
}
607+
let mut replay = Request::new(Body::empty());
608+
*replay.method_mut() = request.method().clone();
609+
*replay.uri_mut() = request.uri().clone();
610+
*replay.version_mut() = request.version();
611+
*replay.headers_mut() = request.headers().clone();
612+
Some(replay)
613+
}
614+
577615
fn build_upstream_request(
578616
req: Request<Body>,
579617
target_port: u16,
@@ -1403,6 +1441,44 @@ mod tests {
14031441
);
14041442
}
14051443

1444+
#[tokio::test]
1445+
async fn take_removes_an_endpoint_left_with_no_upstreams() {
1446+
let pool = ServiceUpstreamPool::default();
1447+
let (sender, _sandbox) = test_upstream().await;
1448+
pool.put("ep-a|8080", sender);
1449+
1450+
assert!(pool.take("ep-a|8080").is_some());
1451+
assert!(
1452+
!pool.inner.lock().unwrap().idle.contains_key("ep-a|8080"),
1453+
"an emptied endpoint must not linger until the next sweep"
1454+
);
1455+
}
1456+
1457+
#[test]
1458+
fn only_bodyless_requests_are_replayable() {
1459+
for method in [Method::GET, Method::HEAD] {
1460+
let mut request = Request::new(Body::empty());
1461+
*request.method_mut() = method.clone();
1462+
*request.uri_mut() = "/health".parse().unwrap();
1463+
request
1464+
.headers_mut()
1465+
.insert(header::HOST, HeaderValue::from_static("svc"));
1466+
1467+
let replay =
1468+
replayable_request(&request).expect("bodyless method should be replayable");
1469+
assert_eq!(*replay.method(), method);
1470+
assert_eq!(replay.uri().path(), "/health");
1471+
assert_eq!(replay.headers().get(header::HOST).unwrap(), "svc");
1472+
}
1473+
1474+
let mut post = Request::new(Body::empty());
1475+
*post.method_mut() = Method::POST;
1476+
assert!(
1477+
replayable_request(&post).is_none(),
1478+
"a request with a body cannot be replayed"
1479+
);
1480+
}
1481+
14061482
#[tokio::test]
14071483
async fn pool_does_not_hand_out_another_endpoints_upstream() {
14081484
let pool = ServiceUpstreamPool::default();

0 commit comments

Comments
 (0)