From 612e7c035bc5e3c6938140102d5f930ef70c8401 Mon Sep 17 00:00:00 2001 From: divesh Date: Thu, 29 Oct 2026 13:34:40 -0700 Subject: [PATCH 01/11] feat(server): place pool sandbox sessions with their supervisor's replica A shared supervisor pool sandbox's session now goes to the gateway replica serving its owning member's stream, instead of the replica the ring picks by sandbox ID. Each pool reconciliation pass publishes the sandbox-to-replica map; sessions fall back to the ring when the parent is unknown, not a live peer, or this replica is shutting down. Signed-off-by: divesh --- crates/openshell-server/src/lib.rs | 7 ++ .../openshell-server/src/supervisor_pool.rs | 94 +++++++++++++++++++ .../src/supervisor_session.rs | 89 +++++++++++++++++- 3 files changed, 186 insertions(+), 4 deletions(-) diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 12ca9223ea..6a9e5ea992 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -379,6 +379,11 @@ pub struct ServerState { /// alongside `gateway_ring` so a redirect can name a dialable address. pub gateway_peers: Arc>>, + /// Replica serving the shared supervisor pool member that owns each + /// sandbox, keyed by sandbox ID. Refreshed by every pool reconciliation + /// pass so a pool sandbox's session is placed with its supervisor. + pub pool_session_parents: supervisor_pool::SessionParents, + /// Validated built-in and operator-registered supervisor middleware. pub middleware_registry: Arc, @@ -519,6 +524,7 @@ impl ServerState { service_route_index, gateway_ring: Arc::new(std::sync::RwLock::new(gateway_ring::GatewayRing::default())), gateway_peers: Arc::new(std::sync::RwLock::new(HashMap::new())), + pool_session_parents: supervisor_pool::SessionParents::default(), extension_mint_limiter: auth::extension_mint_limit::ExtensionMintLimiter::default(), middleware_registry: Arc::new(MiddlewareRegistry::default()), oidc_cache, @@ -842,6 +848,7 @@ pub(crate) async fn run_server( authority.clone(), state.credentials.clone(), state.replica_id.clone(), + state.pool_session_parents.clone(), ) .start(shutdown_rx.clone()); spawn_supervisor_token_refresh(authority, ttl, controller.clone(), shutdown_rx.clone()); diff --git a/crates/openshell-server/src/supervisor_pool.rs b/crates/openshell-server/src/supervisor_pool.rs index cd2dc478bc..ba7ff3ba18 100644 --- a/crates/openshell-server/src/supervisor_pool.rs +++ b/crates/openshell-server/src/supervisor_pool.rs @@ -181,8 +181,14 @@ pub struct LiveMember { pub process_id: String, pub runtime_address: String, pub leaving: bool, + /// The gateway replica serving the member's stream. + pub gateway_replica: String, } +/// The gateway replica serving the member that owns each pool sandbox, keyed +/// by sandbox ID. +pub type SessionParents = Arc>>; + #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] enum OwnerPhase { @@ -409,6 +415,8 @@ pub struct SupervisorPool { withdrawn_since: Mutex>, /// Where the next pass resumes checking records for ended generations. sweep_from: Mutex, + /// Published after every pass for session placement. + session_parents: SessionParents, /// See [`MATERIAL_CANDIDATE_GRACE`]. candidate_grace: Duration, /// See [`RESERVATION_GRACE`]. @@ -566,6 +574,30 @@ fn recharge( } /// Whether the process that owns `owner` still serves its member. +/// The replica serving each granted, live owner's member, by sandbox. +fn session_parents_of( + owners: &BTreeMap, + members: &[LiveMember], +) -> HashMap { + owners + .values() + .filter(|owner| { + owner.payload.phase == OwnerPhase::Granted && owner_alive(&owner.payload, members) + }) + .filter_map(|owner| { + let member = members + .iter() + .find(|member| member.member_id == owner.payload.member_id)?; + (!member.gateway_replica.is_empty()).then(|| { + ( + owner.payload.sandbox_id.clone(), + member.gateway_replica.clone(), + ) + }) + }) + .collect() +} + fn owner_alive(owner: &OwnerPayload, members: &[LiveMember]) -> bool { members.iter().any(|member| { member.member_id == owner.member_id @@ -681,6 +713,7 @@ impl SupervisorPool { authority: Arc, credentials: CredentialRuntime, replica_id: String, + session_parents: SessionParents, ) -> Arc { Arc::new(Self { store, @@ -697,6 +730,7 @@ impl SupervisorPool { materials: Mutex::new(HashMap::new()), withdrawn_since: Mutex::new(HashMap::new()), sweep_from: Mutex::new(0), + session_parents, candidate_grace: MATERIAL_CANDIDATE_GRACE, reservation_grace: RESERVATION_GRACE, #[cfg(test)] @@ -1079,6 +1113,7 @@ impl SupervisorPool { process_id: payload.process_id, runtime_address: payload.runtime_address, leaving: payload.leaving, + gateway_replica: payload.gateway_replica, }); } members.sort_by(|left, right| left.member_id.cmp(&right.member_id)); @@ -2152,6 +2187,9 @@ impl SupervisorPool { } } self.remove_withdrawn_owners(&desired, &mut owners).await; + if let Ok(mut parents) = self.session_parents.write() { + *parents = session_parents_of(&owners, &members); + } self.sweep_ended_instances(&desired).await; self.record_reports(&owners, &report_rows).await; self.remove_stale_reports(&owners, &report_rows).await; @@ -2560,6 +2598,7 @@ impl SupervisorPool { process_id: payload.process_id, runtime_address: payload.runtime_address, leaving: payload.leaving, + gateway_replica: payload.gateway_replica, }); Ok(takeover_cause(owner, member.as_ref())) } @@ -2926,6 +2965,7 @@ mod tests { self.authority.clone(), self.credentials.clone(), replica_id.to_string(), + SessionParents::default(), ); let unique = Arc::get_mut(&mut pool).expect("unique"); unique.unowned_grace = Duration::ZERO; @@ -2958,6 +2998,7 @@ mod tests { self.authority.clone(), self.credentials.clone(), replica_id.to_string(), + SessionParents::default(), ) } @@ -5248,6 +5289,57 @@ mod tests { assert!(fixture.report_rows(&route).await.is_empty()); } + #[test] + fn session_parents_name_the_replica_serving_each_live_owner() { + let owner = |sandbox: &str, member: &str, phase| OwnerRecord { + payload: OwnerPayload { + route: format!("rt-{sandbox}"), + sandbox_id: sandbox.to_string(), + generation: "generation-1".to_string(), + member_id: member.to_string(), + process_id: format!("p-{member}"), + auth_epoch: 1, + supervisor_instance_id: None, + gateway_token_id: None, + phase, + previous_member_id: None, + }, + resource_version: 1, + }; + let member = |id: &str, gateway: &str, leaving| LiveMember { + member_id: id.to_string(), + process_id: format!("p-{id}"), + runtime_address: "10.0.0.1:5500".to_string(), + leaving, + gateway_replica: gateway.to_string(), + }; + let owners = [ + owner("sb-1", "m1", OwnerPhase::Granted), + owner("sb-2", "m2", OwnerPhase::Granted), + owner("sb-3", "m3", OwnerPhase::Granted), + owner("sb-4", "m1", OwnerPhase::Granting), + owner("sb-5", "m-dead", OwnerPhase::Granted), + ] + .into_iter() + .map(|record| (record.payload.route.clone(), record)) + .collect::>(); + let members = vec![ + member("m1", "gw-0", false), + member("m2", "gw-1", false), + member("m3", "gw-1", true), + ]; + + let parents = session_parents_of(&owners, &members); + + assert_eq!( + parents, + HashMap::from([ + ("sb-1".to_string(), "gw-0".to_string()), + ("sb-2".to_string(), "gw-1".to_string()), + ]) + ); + } + #[test] fn a_takeover_names_why_the_owner_lost_its_instance() { let owner = OwnerPayload { @@ -5267,6 +5359,7 @@ mod tests { process_id: process_id.to_string(), runtime_address: "10.0.0.1:5500".to_string(), leaving, + gateway_replica: "gw-a".to_string(), }; assert_eq!(takeover_cause(&owner, None), Some("expired")); assert_eq!( @@ -5291,6 +5384,7 @@ mod tests { process_id: format!("p-{id}"), runtime_address: "10.0.0.1:5500".to_string(), leaving, + gateway_replica: "gw-a".to_string(), }; let members = vec![member("m1", false), member("m2", false), member("m3", true)]; let ring = ring_of(&members); diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 248f5fab7b..5f7a7c0154 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -1859,7 +1859,11 @@ fn local_owner_endpoint(replica_id: &str) -> String { ) } -/// Where the ring says this sandbox belongs, when that is not this replica. +/// Where this sandbox's session belongs, when that is not this replica. +/// +/// A shared supervisor pool sandbox belongs with the replica serving its +/// owning member; any other sandbox, or a pool sandbox whose parent is not +/// known to be live, goes where the ring says. /// /// Returns `None` whenever this replica should serve the session itself, which /// covers single-replica mode, an empty ring, the ring naming this replica, and @@ -1871,9 +1875,25 @@ fn preferred_peer_redirect(state: &Arc, sandbox_id: &str) -> Option return None; } + let shutting_down = state.gateway_shutting_down.load(Ordering::Acquire); + let parent = { + let parents = state.pool_session_parents.read().ok()?; + let peers = state.gateway_peers.read().ok()?; + pool_session_parent( + &parents, + &peers, + &state.replica_id, + sandbox_id, + shutting_down, + ) + }; + if let Some((parent, peer_endpoint)) = parent { + return ring_redirect(&state.replica_id, &parent, peer_endpoint); + } + let preferred = { let ring = state.gateway_ring.read().ok()?; - if state.gateway_shutting_down.load(Ordering::Acquire) { + if shutting_down { ring.without(&state.replica_id) .owner_for(sandbox_id)? .to_string() @@ -1889,6 +1909,24 @@ fn preferred_peer_redirect(state: &Arc, sandbox_id: &str) -> Option ring_redirect(&state.replica_id, &preferred, peer_endpoint) } +/// The replica serving the pool member that owns `sandbox_id`, with its peer +/// endpoint, when that replica can take the session: this replica unless it +/// is shutting down, or a live peer. +fn pool_session_parent( + parents: &HashMap, + peers: &HashMap, + self_replica_id: &str, + sandbox_id: &str, + shutting_down: bool, +) -> Option<(String, Option)> { + let parent = parents.get(sandbox_id)?; + if parent == self_replica_id { + return (!shutting_down).then(|| (parent.clone(), None)); + } + let peer_endpoint = peers.get(parent)?; + Some((parent.clone(), Some(peer_endpoint.clone()))) +} + /// Decide whether a preferred replica is worth redirecting to. /// /// Split out from the lock reads so the rules are testable on their own. @@ -2555,7 +2593,7 @@ pub async fn handle_connect_supervisor( // Validate readiness identities before replacing a healthy session. let provider_readiness = ProviderReadinessEvidence::from_hello(&hello)?; - // If the ring places this sandbox on a different live replica, send the + // If this sandbox belongs on a different live replica, send the // supervisor there instead of claiming ownership here. Decided before the // session is tracked so a redirect never consumes a session slot. An // established session is only relocated when this replica shuts down. @@ -2566,7 +2604,7 @@ pub async fn handle_connect_supervisor( sandbox_id = %sandbox_id, owner_replica_id = %redirect.owner_replica_id, replica_id = %state.replica_id, - "supervisor session: redirecting to ring-preferred replica" + "supervisor session: redirecting to preferred replica" ); if state.supervisor_sessions.has_session(&sandbox_id) { let owner_index = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL); @@ -5584,6 +5622,49 @@ mod tests { assert!(ring_redirect("gw-0", "gw-1", Some("local://gw-1".to_string())).is_none()); } + fn pool_maps() -> (HashMap, HashMap) { + let parents = HashMap::from([ + ("sb-a".to_string(), "gw-0".to_string()), + ("sb-b".to_string(), "gw-1".to_string()), + ("sb-gone".to_string(), "gw-9".to_string()), + ]); + let peers = HashMap::from([("gw-1".to_string(), "https://gw-1:8443".to_string())]); + (parents, peers) + } + + #[test] + fn pool_session_follows_its_members_replica() { + let (parents, peers) = pool_maps(); + assert_eq!( + pool_session_parent(&parents, &peers, "gw-0", "sb-b", false), + Some(("gw-1".to_string(), Some("https://gw-1:8443".to_string()))) + ); + assert_eq!( + pool_session_parent(&parents, &peers, "gw-0", "sb-a", false), + Some(("gw-0".to_string(), None)) + ); + } + + #[test] + fn pool_session_falls_back_to_the_ring_without_a_live_parent() { + let (parents, peers) = pool_maps(); + // Not a pool sandbox. + assert_eq!( + pool_session_parent(&parents, &peers, "gw-0", "sb-dedicated", false), + None + ); + // Parent replica is not a live peer. + assert_eq!( + pool_session_parent(&parents, &peers, "gw-0", "sb-gone", false), + None + ); + // This replica is the parent but is shutting down. + assert_eq!( + pool_session_parent(&parents, &peers, "gw-0", "sb-a", true), + None + ); + } + // ---- ownership terms ---- fn term(value: u64) -> CredentialEpoch { From d3910f0de97f3a5f596d782bfbd3ba76e4bf0238 Mon Sep 17 00:00:00 2001 From: divesh Date: Thu, 29 Oct 2026 14:32:32 -0700 Subject: [PATCH 02/11] feat(server): move live pool sessions to their supervisor's replica A pool sandbox session served by a replica other than the one serving its owning member is redirected there on its next heartbeat, reusing the shutdown handoff so the sandbox stays Ready. This re-homes sessions left on the wrong replica by a rolling restart. Signed-off-by: divesh --- .../src/supervisor_session.rs | 64 ++++++++++++++++++- 1 file changed, 62 insertions(+), 2 deletions(-) diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 5f7a7c0154..7029fb494a 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -1927,6 +1927,28 @@ fn pool_session_parent( Some((parent.clone(), Some(peer_endpoint.clone()))) } +/// Where to move a live pool sandbox session that this replica serves but +/// whose owning member's stream is served by another live replica. +fn pool_relocation(state: &ServerState, sandbox_id: &str) -> Option { + if state.store.is_single_replica() || state.gateway_shutting_down.load(Ordering::Acquire) { + return None; + } + let parents = state.pool_session_parents.read().ok()?; + let peers = state.gateway_peers.read().ok()?; + pool_relocation_target(&parents, &peers, &state.replica_id, sandbox_id) +} + +fn pool_relocation_target( + parents: &HashMap, + peers: &HashMap, + self_replica_id: &str, + sandbox_id: &str, +) -> Option { + let (parent, peer_endpoint) = + pool_session_parent(parents, peers, self_replica_id, sandbox_id, false)?; + ring_redirect(self_replica_id, &parent, peer_endpoint) +} + /// Decide whether a preferred replica is worth redirecting to. /// /// Split out from the lock reads so the rules are testable on their own. @@ -2788,6 +2810,8 @@ enum SessionLoopExit { /// This gateway is shutting down. `redirected` records whether a /// `SessionRedirect` was queued for the supervisor. GatewayShutdown { redirected: bool }, + /// The supervisor was redirected to the replica serving its pool member. + Relocated, /// The session ended for any other reason. Other, } @@ -2810,8 +2834,10 @@ async fn finish_supervisor_session( // owner and the sandbox stays Ready through the handoff. Without a // redirect the supervisor reconnects with backoff through the Service, so // waiting would only delay the demotion. - if exit == (SessionLoopExit::GatewayShutdown { redirected: true }) - && terminal_finalized == Some(false) + if matches!( + exit, + SessionLoopExit::GatewayShutdown { redirected: true } | SessionLoopExit::Relocated + ) && terminal_finalized == Some(false) && sandbox_may_be_ready(state, sandbox_id).await && wait_for_replacement_owner( &owner_index, @@ -3024,6 +3050,21 @@ async fn run_session_loop( } } _ = heartbeat_timer.tick() => { + if *supports_session_redirect + && bootstrap_complete + && let Some(redirect) = pool_relocation(state, sandbox_id) + { + let owner_replica_id = redirect.owner_replica_id.clone(); + if tx + .try_send(GatewayMessage { + payload: Some(gateway_message::Payload::SessionRedirect(redirect)), + }) + .is_ok() + { + info!(sandbox_id = %sandbox_id, session_id = %session_id, owner_replica_id, "supervisor session: moving to its supervisor's replica"); + return SessionLoopExit::Relocated; + } + } let hb = GatewayMessage { payload: Some(gateway_message::Payload::Heartbeat( openshell_core::proto::GatewayHeartbeat {}, @@ -5665,6 +5706,25 @@ mod tests { ); } + #[test] + fn a_live_pool_session_moves_only_to_its_parents_live_peer() { + let (parents, peers) = pool_maps(); + assert_eq!( + pool_relocation_target(&parents, &peers, "gw-0", "sb-b"), + Some(SessionRedirect { + peer_endpoint: "https://gw-1:8443".to_string(), + owner_replica_id: "gw-1".to_string(), + }) + ); + for sandbox_id in ["sb-a", "sb-gone", "sb-dedicated"] { + assert_eq!( + pool_relocation_target(&parents, &peers, "gw-0", sandbox_id), + None, + "{sandbox_id}" + ); + } + } + // ---- ownership terms ---- fn term(value: u64) -> CredentialEpoch { From 465aa71eb3b64179cf95037a6afa4fec2aa619f4 Mon Sep 17 00:00:00 2001 From: divesh Date: Fri, 30 Oct 2026 14:41:59 -0700 Subject: [PATCH 03/11] refactor(server): drop gateway-side pool session placement Removes the per-sandbox parent map and the heartbeat relocation; pool sessions will dial their supervisor's gateway directly instead. Signed-off-by: divesh --- crates/openshell-server/src/lib.rs | 7 - .../openshell-server/src/supervisor_pool.rs | 94 ----------- .../src/supervisor_session.rs | 153 +----------------- 3 files changed, 6 insertions(+), 248 deletions(-) diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 6a9e5ea992..12ca9223ea 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -379,11 +379,6 @@ pub struct ServerState { /// alongside `gateway_ring` so a redirect can name a dialable address. pub gateway_peers: Arc>>, - /// Replica serving the shared supervisor pool member that owns each - /// sandbox, keyed by sandbox ID. Refreshed by every pool reconciliation - /// pass so a pool sandbox's session is placed with its supervisor. - pub pool_session_parents: supervisor_pool::SessionParents, - /// Validated built-in and operator-registered supervisor middleware. pub middleware_registry: Arc, @@ -524,7 +519,6 @@ impl ServerState { service_route_index, gateway_ring: Arc::new(std::sync::RwLock::new(gateway_ring::GatewayRing::default())), gateway_peers: Arc::new(std::sync::RwLock::new(HashMap::new())), - pool_session_parents: supervisor_pool::SessionParents::default(), extension_mint_limiter: auth::extension_mint_limit::ExtensionMintLimiter::default(), middleware_registry: Arc::new(MiddlewareRegistry::default()), oidc_cache, @@ -848,7 +842,6 @@ pub(crate) async fn run_server( authority.clone(), state.credentials.clone(), state.replica_id.clone(), - state.pool_session_parents.clone(), ) .start(shutdown_rx.clone()); spawn_supervisor_token_refresh(authority, ttl, controller.clone(), shutdown_rx.clone()); diff --git a/crates/openshell-server/src/supervisor_pool.rs b/crates/openshell-server/src/supervisor_pool.rs index ba7ff3ba18..cd2dc478bc 100644 --- a/crates/openshell-server/src/supervisor_pool.rs +++ b/crates/openshell-server/src/supervisor_pool.rs @@ -181,14 +181,8 @@ pub struct LiveMember { pub process_id: String, pub runtime_address: String, pub leaving: bool, - /// The gateway replica serving the member's stream. - pub gateway_replica: String, } -/// The gateway replica serving the member that owns each pool sandbox, keyed -/// by sandbox ID. -pub type SessionParents = Arc>>; - #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] enum OwnerPhase { @@ -415,8 +409,6 @@ pub struct SupervisorPool { withdrawn_since: Mutex>, /// Where the next pass resumes checking records for ended generations. sweep_from: Mutex, - /// Published after every pass for session placement. - session_parents: SessionParents, /// See [`MATERIAL_CANDIDATE_GRACE`]. candidate_grace: Duration, /// See [`RESERVATION_GRACE`]. @@ -574,30 +566,6 @@ fn recharge( } /// Whether the process that owns `owner` still serves its member. -/// The replica serving each granted, live owner's member, by sandbox. -fn session_parents_of( - owners: &BTreeMap, - members: &[LiveMember], -) -> HashMap { - owners - .values() - .filter(|owner| { - owner.payload.phase == OwnerPhase::Granted && owner_alive(&owner.payload, members) - }) - .filter_map(|owner| { - let member = members - .iter() - .find(|member| member.member_id == owner.payload.member_id)?; - (!member.gateway_replica.is_empty()).then(|| { - ( - owner.payload.sandbox_id.clone(), - member.gateway_replica.clone(), - ) - }) - }) - .collect() -} - fn owner_alive(owner: &OwnerPayload, members: &[LiveMember]) -> bool { members.iter().any(|member| { member.member_id == owner.member_id @@ -713,7 +681,6 @@ impl SupervisorPool { authority: Arc, credentials: CredentialRuntime, replica_id: String, - session_parents: SessionParents, ) -> Arc { Arc::new(Self { store, @@ -730,7 +697,6 @@ impl SupervisorPool { materials: Mutex::new(HashMap::new()), withdrawn_since: Mutex::new(HashMap::new()), sweep_from: Mutex::new(0), - session_parents, candidate_grace: MATERIAL_CANDIDATE_GRACE, reservation_grace: RESERVATION_GRACE, #[cfg(test)] @@ -1113,7 +1079,6 @@ impl SupervisorPool { process_id: payload.process_id, runtime_address: payload.runtime_address, leaving: payload.leaving, - gateway_replica: payload.gateway_replica, }); } members.sort_by(|left, right| left.member_id.cmp(&right.member_id)); @@ -2187,9 +2152,6 @@ impl SupervisorPool { } } self.remove_withdrawn_owners(&desired, &mut owners).await; - if let Ok(mut parents) = self.session_parents.write() { - *parents = session_parents_of(&owners, &members); - } self.sweep_ended_instances(&desired).await; self.record_reports(&owners, &report_rows).await; self.remove_stale_reports(&owners, &report_rows).await; @@ -2598,7 +2560,6 @@ impl SupervisorPool { process_id: payload.process_id, runtime_address: payload.runtime_address, leaving: payload.leaving, - gateway_replica: payload.gateway_replica, }); Ok(takeover_cause(owner, member.as_ref())) } @@ -2965,7 +2926,6 @@ mod tests { self.authority.clone(), self.credentials.clone(), replica_id.to_string(), - SessionParents::default(), ); let unique = Arc::get_mut(&mut pool).expect("unique"); unique.unowned_grace = Duration::ZERO; @@ -2998,7 +2958,6 @@ mod tests { self.authority.clone(), self.credentials.clone(), replica_id.to_string(), - SessionParents::default(), ) } @@ -5289,57 +5248,6 @@ mod tests { assert!(fixture.report_rows(&route).await.is_empty()); } - #[test] - fn session_parents_name_the_replica_serving_each_live_owner() { - let owner = |sandbox: &str, member: &str, phase| OwnerRecord { - payload: OwnerPayload { - route: format!("rt-{sandbox}"), - sandbox_id: sandbox.to_string(), - generation: "generation-1".to_string(), - member_id: member.to_string(), - process_id: format!("p-{member}"), - auth_epoch: 1, - supervisor_instance_id: None, - gateway_token_id: None, - phase, - previous_member_id: None, - }, - resource_version: 1, - }; - let member = |id: &str, gateway: &str, leaving| LiveMember { - member_id: id.to_string(), - process_id: format!("p-{id}"), - runtime_address: "10.0.0.1:5500".to_string(), - leaving, - gateway_replica: gateway.to_string(), - }; - let owners = [ - owner("sb-1", "m1", OwnerPhase::Granted), - owner("sb-2", "m2", OwnerPhase::Granted), - owner("sb-3", "m3", OwnerPhase::Granted), - owner("sb-4", "m1", OwnerPhase::Granting), - owner("sb-5", "m-dead", OwnerPhase::Granted), - ] - .into_iter() - .map(|record| (record.payload.route.clone(), record)) - .collect::>(); - let members = vec![ - member("m1", "gw-0", false), - member("m2", "gw-1", false), - member("m3", "gw-1", true), - ]; - - let parents = session_parents_of(&owners, &members); - - assert_eq!( - parents, - HashMap::from([ - ("sb-1".to_string(), "gw-0".to_string()), - ("sb-2".to_string(), "gw-1".to_string()), - ]) - ); - } - #[test] fn a_takeover_names_why_the_owner_lost_its_instance() { let owner = OwnerPayload { @@ -5359,7 +5267,6 @@ mod tests { process_id: process_id.to_string(), runtime_address: "10.0.0.1:5500".to_string(), leaving, - gateway_replica: "gw-a".to_string(), }; assert_eq!(takeover_cause(&owner, None), Some("expired")); assert_eq!( @@ -5384,7 +5291,6 @@ mod tests { process_id: format!("p-{id}"), runtime_address: "10.0.0.1:5500".to_string(), leaving, - gateway_replica: "gw-a".to_string(), }; let members = vec![member("m1", false), member("m2", false), member("m3", true)]; let ring = ring_of(&members); diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 7029fb494a..248f5fab7b 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -1859,11 +1859,7 @@ fn local_owner_endpoint(replica_id: &str) -> String { ) } -/// Where this sandbox's session belongs, when that is not this replica. -/// -/// A shared supervisor pool sandbox belongs with the replica serving its -/// owning member; any other sandbox, or a pool sandbox whose parent is not -/// known to be live, goes where the ring says. +/// Where the ring says this sandbox belongs, when that is not this replica. /// /// Returns `None` whenever this replica should serve the session itself, which /// covers single-replica mode, an empty ring, the ring naming this replica, and @@ -1875,25 +1871,9 @@ fn preferred_peer_redirect(state: &Arc, sandbox_id: &str) -> Option return None; } - let shutting_down = state.gateway_shutting_down.load(Ordering::Acquire); - let parent = { - let parents = state.pool_session_parents.read().ok()?; - let peers = state.gateway_peers.read().ok()?; - pool_session_parent( - &parents, - &peers, - &state.replica_id, - sandbox_id, - shutting_down, - ) - }; - if let Some((parent, peer_endpoint)) = parent { - return ring_redirect(&state.replica_id, &parent, peer_endpoint); - } - let preferred = { let ring = state.gateway_ring.read().ok()?; - if shutting_down { + if state.gateway_shutting_down.load(Ordering::Acquire) { ring.without(&state.replica_id) .owner_for(sandbox_id)? .to_string() @@ -1909,46 +1889,6 @@ fn preferred_peer_redirect(state: &Arc, sandbox_id: &str) -> Option ring_redirect(&state.replica_id, &preferred, peer_endpoint) } -/// The replica serving the pool member that owns `sandbox_id`, with its peer -/// endpoint, when that replica can take the session: this replica unless it -/// is shutting down, or a live peer. -fn pool_session_parent( - parents: &HashMap, - peers: &HashMap, - self_replica_id: &str, - sandbox_id: &str, - shutting_down: bool, -) -> Option<(String, Option)> { - let parent = parents.get(sandbox_id)?; - if parent == self_replica_id { - return (!shutting_down).then(|| (parent.clone(), None)); - } - let peer_endpoint = peers.get(parent)?; - Some((parent.clone(), Some(peer_endpoint.clone()))) -} - -/// Where to move a live pool sandbox session that this replica serves but -/// whose owning member's stream is served by another live replica. -fn pool_relocation(state: &ServerState, sandbox_id: &str) -> Option { - if state.store.is_single_replica() || state.gateway_shutting_down.load(Ordering::Acquire) { - return None; - } - let parents = state.pool_session_parents.read().ok()?; - let peers = state.gateway_peers.read().ok()?; - pool_relocation_target(&parents, &peers, &state.replica_id, sandbox_id) -} - -fn pool_relocation_target( - parents: &HashMap, - peers: &HashMap, - self_replica_id: &str, - sandbox_id: &str, -) -> Option { - let (parent, peer_endpoint) = - pool_session_parent(parents, peers, self_replica_id, sandbox_id, false)?; - ring_redirect(self_replica_id, &parent, peer_endpoint) -} - /// Decide whether a preferred replica is worth redirecting to. /// /// Split out from the lock reads so the rules are testable on their own. @@ -2615,7 +2555,7 @@ pub async fn handle_connect_supervisor( // Validate readiness identities before replacing a healthy session. let provider_readiness = ProviderReadinessEvidence::from_hello(&hello)?; - // If this sandbox belongs on a different live replica, send the + // If the ring places this sandbox on a different live replica, send the // supervisor there instead of claiming ownership here. Decided before the // session is tracked so a redirect never consumes a session slot. An // established session is only relocated when this replica shuts down. @@ -2626,7 +2566,7 @@ pub async fn handle_connect_supervisor( sandbox_id = %sandbox_id, owner_replica_id = %redirect.owner_replica_id, replica_id = %state.replica_id, - "supervisor session: redirecting to preferred replica" + "supervisor session: redirecting to ring-preferred replica" ); if state.supervisor_sessions.has_session(&sandbox_id) { let owner_index = SupervisorOwnerIndex::new(state.store.clone(), OWNER_TTL); @@ -2810,8 +2750,6 @@ enum SessionLoopExit { /// This gateway is shutting down. `redirected` records whether a /// `SessionRedirect` was queued for the supervisor. GatewayShutdown { redirected: bool }, - /// The supervisor was redirected to the replica serving its pool member. - Relocated, /// The session ended for any other reason. Other, } @@ -2834,10 +2772,8 @@ async fn finish_supervisor_session( // owner and the sandbox stays Ready through the handoff. Without a // redirect the supervisor reconnects with backoff through the Service, so // waiting would only delay the demotion. - if matches!( - exit, - SessionLoopExit::GatewayShutdown { redirected: true } | SessionLoopExit::Relocated - ) && terminal_finalized == Some(false) + if exit == (SessionLoopExit::GatewayShutdown { redirected: true }) + && terminal_finalized == Some(false) && sandbox_may_be_ready(state, sandbox_id).await && wait_for_replacement_owner( &owner_index, @@ -3050,21 +2986,6 @@ async fn run_session_loop( } } _ = heartbeat_timer.tick() => { - if *supports_session_redirect - && bootstrap_complete - && let Some(redirect) = pool_relocation(state, sandbox_id) - { - let owner_replica_id = redirect.owner_replica_id.clone(); - if tx - .try_send(GatewayMessage { - payload: Some(gateway_message::Payload::SessionRedirect(redirect)), - }) - .is_ok() - { - info!(sandbox_id = %sandbox_id, session_id = %session_id, owner_replica_id, "supervisor session: moving to its supervisor's replica"); - return SessionLoopExit::Relocated; - } - } let hb = GatewayMessage { payload: Some(gateway_message::Payload::Heartbeat( openshell_core::proto::GatewayHeartbeat {}, @@ -5663,68 +5584,6 @@ mod tests { assert!(ring_redirect("gw-0", "gw-1", Some("local://gw-1".to_string())).is_none()); } - fn pool_maps() -> (HashMap, HashMap) { - let parents = HashMap::from([ - ("sb-a".to_string(), "gw-0".to_string()), - ("sb-b".to_string(), "gw-1".to_string()), - ("sb-gone".to_string(), "gw-9".to_string()), - ]); - let peers = HashMap::from([("gw-1".to_string(), "https://gw-1:8443".to_string())]); - (parents, peers) - } - - #[test] - fn pool_session_follows_its_members_replica() { - let (parents, peers) = pool_maps(); - assert_eq!( - pool_session_parent(&parents, &peers, "gw-0", "sb-b", false), - Some(("gw-1".to_string(), Some("https://gw-1:8443".to_string()))) - ); - assert_eq!( - pool_session_parent(&parents, &peers, "gw-0", "sb-a", false), - Some(("gw-0".to_string(), None)) - ); - } - - #[test] - fn pool_session_falls_back_to_the_ring_without_a_live_parent() { - let (parents, peers) = pool_maps(); - // Not a pool sandbox. - assert_eq!( - pool_session_parent(&parents, &peers, "gw-0", "sb-dedicated", false), - None - ); - // Parent replica is not a live peer. - assert_eq!( - pool_session_parent(&parents, &peers, "gw-0", "sb-gone", false), - None - ); - // This replica is the parent but is shutting down. - assert_eq!( - pool_session_parent(&parents, &peers, "gw-0", "sb-a", true), - None - ); - } - - #[test] - fn a_live_pool_session_moves_only_to_its_parents_live_peer() { - let (parents, peers) = pool_maps(); - assert_eq!( - pool_relocation_target(&parents, &peers, "gw-0", "sb-b"), - Some(SessionRedirect { - peer_endpoint: "https://gw-1:8443".to_string(), - owner_replica_id: "gw-1".to_string(), - }) - ); - for sandbox_id in ["sb-a", "sb-gone", "sb-dedicated"] { - assert_eq!( - pool_relocation_target(&parents, &peers, "gw-0", sandbox_id), - None, - "{sandbox_id}" - ); - } - } - // ---- ownership terms ---- fn term(value: u64) -> CredentialEpoch { From 96aa7cf80ede5c8184863b63dc1686db621d1f15 Mon Sep 17 00:00:00 2001 From: divesh Date: Fri, 30 Oct 2026 14:52:44 -0700 Subject: [PATCH 04/11] feat(supervisor): dial pool sandbox sessions to the host's gateway replica HostAccepted now carries the accepting replica's peer endpoint. A shared host's sandbox sessions dial it directly, so each session runs on the replica serving its supervisor. A stopping replica sheds its host streams, and redirected sessions follow the host to its next replica. Closes #5051. Signed-off-by: divesh --- .../src/grpc/supervisor_host.rs | 9 +++ crates/openshell-server/src/lib.rs | 1 + .../src/supervisor_session.rs | 5 ++ .../src/controller.rs | 11 +++ .../src/delegated.rs | 4 ++ .../src/supervisor_session.rs | 70 ++++++++++++++++++- .../src/backend_setup/tests.rs | 2 + .../src/legacy_startup_regression_tests.rs | 1 + crates/openshell-supervisor/src/lib.rs | 6 ++ crates/openshell-supervisor/src/main.rs | 1 + .../openshell-supervisor/src/shared_host.rs | 34 +++++++++ proto/supervisor_host.proto | 4 ++ 12 files changed, 147 insertions(+), 1 deletion(-) diff --git a/crates/openshell-server/src/grpc/supervisor_host.rs b/crates/openshell-server/src/grpc/supervisor_host.rs index e51ca87097..9760c4c97b 100644 --- a/crates/openshell-server/src/grpc/supervisor_host.rs +++ b/crates/openshell-server/src/grpc/supervisor_host.rs @@ -59,7 +59,16 @@ impl SupervisorHost for SupervisorHostService { Status::failed_precondition("this gateway does not host a shared supervisor") })? .clone(); + // A stopping replica sheds its hosts so they reattach to a live one + // before their sandbox sessions follow them there. + let mut shutdown = self.state.supervisor_sessions.subscribe_shutdown(); + if *shutdown.borrow() { + return Err(Status::unavailable("gateway is shutting down")); + } let outbound = controller.attach(identity, request.into_inner()).await?; + let outbound = futures::StreamExt::take_until(outbound, async move { + let _ = shutdown.wait_for(|shutdown| *shutdown).await; + }); Ok(Response::new(Box::pin(outbound))) } } diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 12ca9223ea..7522d5cbe7 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -836,6 +836,7 @@ pub(crate) async fn run_server( "shared supervisor assignment stream enabled; minting supervisor tokens" ); info!("shared supervisor pool enabled; reconciling pool membership and ownership"); + controller.set_peer_endpoint(state.peer_endpoint.clone().unwrap_or_default()); supervisor_pool::SupervisorPool::new( state.store.clone(), controller.clone(), diff --git a/crates/openshell-server/src/supervisor_session.rs b/crates/openshell-server/src/supervisor_session.rs index 248f5fab7b..2cbe63de28 100644 --- a/crates/openshell-server/src/supervisor_session.rs +++ b/crates/openshell-server/src/supervisor_session.rs @@ -435,6 +435,11 @@ impl SupervisorSessionRegistry { self.admission_closed.store(true, Ordering::Release); } + /// Watch for [`Self::shutdown`] beginning. + pub(crate) fn subscribe_shutdown(&self) -> watch::Receiver { + self.shutdown.subscribe() + } + /// Close control sessions and wait for tracked ownership cleanup. Call /// after compute shutdown so supervisors can finish normal stop reporting. pub(crate) async fn shutdown(&self, timeout: Duration) -> Result<(), String> { diff --git a/crates/openshell-supervisor-control/src/controller.rs b/crates/openshell-supervisor-control/src/controller.rs index d05c5ba728..a5a87507cb 100644 --- a/crates/openshell-supervisor-control/src/controller.rs +++ b/crates/openshell-supervisor-control/src/controller.rs @@ -356,6 +356,8 @@ struct Inner { /// stream attached or ended, or a member announced leaving. membership: watch::Sender, observed: Mutex, + /// This replica's peer endpoint, sent in every `HostAccepted`. + peer_endpoint: Mutex, } /// Gateway-owned desired assignments and observed status of the shared @@ -406,10 +408,17 @@ impl SupervisorHostController { revision: watch::Sender::new(0), membership: watch::Sender::new(0), observed: Mutex::new(Observed::default()), + peer_endpoint: Mutex::new(String::new()), }), } } + /// Tell attaching hosts to dial their sandbox sessions to `peer_endpoint`. + /// Call it before any host attaches. + pub fn set_peer_endpoint(&self, peer_endpoint: String) { + *lock(&self.inner.peer_endpoint) = peer_endpoint; + } + /// Offer every instance's runtime `recovery_timeout_secs` instead of /// [`POOL_RECOVERY_TIMEOUT_SECS`]. Call it before any host attaches. pub fn set_recovery_timeout_secs(&self, recovery_timeout_secs: u32) { @@ -843,6 +852,7 @@ impl SupervisorHostController { let (tx, rx) = mpsc::channel(OUTBOUND_QUEUE); let expires_at = identity.expires_at; let (session_id, superseded) = self.begin_session(identity, &hello)?; + let peer_endpoint = lock(&self.inner.peer_endpoint).clone(); tx.send(Ok(HostGatewayMessage { payload: Some(host_gateway_message::Payload::Accepted(HostAccepted { session_id: session_id.to_string(), @@ -850,6 +860,7 @@ impl SupervisorHostController { seconds: i64::try_from(heartbeat.as_secs()).unwrap_or(15), nanos: 0, }), + peer_endpoint, })), })) .await diff --git a/crates/openshell-supervisor-process/src/delegated.rs b/crates/openshell-supervisor-process/src/delegated.rs index 32e3fd4433..788d80c337 100644 --- a/crates/openshell-supervisor-process/src/delegated.rs +++ b/crates/openshell-supervisor-process/src/delegated.rs @@ -93,6 +93,7 @@ pub async fn start_boundary_access( config_apply_tx: tokio::sync::mpsc::Sender, config_push_updates: tokio::sync::watch::Sender, host_key: Option, + gateway_home: Option>>, ) -> Result { let instance_id = uuid::Uuid::new_v4().to_string(); let terminating = Arc::new(AtomicBool::new(false)); @@ -130,6 +131,7 @@ pub async fn start_boundary_access( session_id_updates: supervisor_session_updates, config_apply_tx, config_push_updates, + gateway_home, }, ); (Some(task), Some(accepted)) @@ -374,6 +376,7 @@ mod tests { tokio::sync::mpsc::channel(1).0, tokio::sync::watch::channel(false).0, None, + None, ) .await .expect("access without SSH is portable"); @@ -415,6 +418,7 @@ mod tests { tokio::sync::mpsc::channel(1).0, tokio::sync::watch::channel(false).0, None, + None, ) .await; assert!( diff --git a/crates/openshell-supervisor-process/src/supervisor_session.rs b/crates/openshell-supervisor-process/src/supervisor_session.rs index 71313e19c0..e3c32011d9 100644 --- a/crates/openshell-supervisor-process/src/supervisor_session.rs +++ b/crates/openshell-supervisor-process/src/supervisor_session.rs @@ -44,6 +44,9 @@ use openshell_core::transport_errors::is_expected_transport_close_status; const INITIAL_BACKOFF: Duration = Duration::from_secs(1); const MAX_BACKOFF: Duration = Duration::from_secs(30); +/// How long a shared host's session waits for the host to attach to a +/// gateway replica before dialing another address. +const GATEWAY_HOME_WAIT: Duration = Duration::from_secs(5); const CONFIG_APPLY_TIMEOUT: Duration = Duration::from_mins(1); /// A stream-delivered configuration update awaiting application by the @@ -129,6 +132,9 @@ pub struct SessionRuntimeContext { /// Publishes whether the current session pushes configuration, so the /// runtime stops polling while it does. pub config_push_updates: watch::Sender, + /// Peer endpoint of the gateway replica serving this supervisor's shared + /// host. When set, sessions dial it instead of the configured address. + pub gateway_home: Option>>, } /// Parse a gRPC endpoint URI into an OCSF `Endpoint` (host + port). Falls back @@ -423,6 +429,7 @@ pub fn spawn_with_readiness( session_id_updates: runtime.session_id_updates, config_apply_tx: runtime.config_apply_tx, config_push_updates: runtime.config_push_updates, + gateway_home: runtime.gateway_home, ready_tx, }; (tokio::spawn(run_session_loop(config)), ready_rx) @@ -442,6 +449,7 @@ struct SessionConfig { /// Publishes whether the current session pushes configuration. A /// reconnect to a gateway that does not push resumes polling. config_push_updates: watch::Sender, + gateway_home: Option>>, ready_tx: watch::Sender, } @@ -455,9 +463,22 @@ async fn run_session_loop(config: SessionConfig) { // back to the configured address, which load-balances across replicas. let mut target = config.gateway.clone(); let mut redirected = false; + // A shared host's sessions dial the replica serving the host, as + // redirected so that replica serves them, unless they are following a + // redirect away from it. + let mut gateway_home = config.gateway_home.clone(); + let mut follow_home = true; loop { attempt += 1; + if follow_home + && let Some(home) = gateway_home.as_mut() + && let Some(endpoint) = wait_for_home(home, None).await + { + target = config.gateway.with_url(endpoint); + redirected = true; + } + follow_home = true; let result = run_single_session(&config, &target, redirected, attempt).await; if let Some(updates) = &config.session_id_updates { @@ -485,8 +506,14 @@ async fn run_session_loop(config: SessionConfig) { peer_endpoint = %peer_endpoint, "supervisor session: following gateway redirect" ); - target = config.gateway.with_url(peer_endpoint); + // Prefer the replica the shared host moved to. + let home = match gateway_home.as_mut() { + Some(home) => wait_for_home(home, Some(target.url())).await, + None => None, + }; + target = config.gateway.with_url(home.unwrap_or(peer_endpoint)); redirected = true; + follow_home = false; // No backoff: this is an expected handoff, not a failure. } Err(e) => { @@ -555,6 +582,22 @@ impl ReconnectBackoff { } } +/// The peer endpoint of the replica serving the shared host, once it is one +/// other than `avoid`, waiting up to [`GATEWAY_HOME_WAIT`]. +async fn wait_for_home( + home: &mut watch::Receiver>, + avoid: Option<&str>, +) -> Option { + tokio::time::timeout( + GATEWAY_HOME_WAIT, + home.wait_for(|home| home.as_deref().is_some_and(|home| Some(home) != avoid)), + ) + .await + .ok()? + .ok() + .and_then(|home| (*home).clone()) +} + /// Keep asking not to be redirected when a redirected attempt never got a /// session, so a stale ring pointing at a dead replica cannot bounce us back /// to it. A session that was accepted and later dropped starts placement over. @@ -1636,6 +1679,31 @@ mod ocsf_event_tests { assert!(!redirect_survives_failure(false, true)); } + #[tokio::test(start_paused = true)] + async fn shared_host_sessions_wait_for_the_hosts_gateway() { + let (tx, mut home) = watch::channel(None); + assert_eq!(wait_for_home(&mut home, None).await, None); + + tx.send_replace(Some("https://gw-a:8443".to_string())); + assert_eq!( + wait_for_home(&mut home, None).await.as_deref(), + Some("https://gw-a:8443") + ); + assert_eq!( + wait_for_home(&mut home, Some("https://gw-a:8443")).await, + None + ); + + let mut moving = home.clone(); + let moved = + tokio::spawn( + async move { wait_for_home(&mut moving, Some("https://gw-a:8443")).await }, + ); + tokio::task::yield_now().await; + tx.send_replace(Some("https://gw-b:8443".to_string())); + assert_eq!(moved.await.unwrap().as_deref(), Some("https://gw-b:8443")); + } + #[test] fn reconnect_backoff_doubles_until_a_session_is_accepted() { let mut backoff = ReconnectBackoff::new(); diff --git a/crates/openshell-supervisor/src/backend_setup/tests.rs b/crates/openshell-supervisor/src/backend_setup/tests.rs index bdf4035c61..44ddcbf5ef 100644 --- a/crates/openshell-supervisor/src/backend_setup/tests.rs +++ b/crates/openshell-supervisor/src/backend_setup/tests.rs @@ -834,6 +834,7 @@ async fn shared_startup_uses_selected_backend_through_readiness_and_shutdown() { sandbox: None, container_image: String::new(), openshell_endpoint: None, + gateway_home: None, policy_rules: Some(rules.to_string_lossy().into_owned()), policy_data: Some(data.to_string_lossy().into_owned()), ssh_socket_path: None, @@ -1038,6 +1039,7 @@ async fn suspend_signal_exits_without_terminating_boundary() { sandbox: None, container_image: String::new(), openshell_endpoint: None, + gateway_home: None, policy_rules: Some(rules.to_string_lossy().into_owned()), policy_data: Some(data.to_string_lossy().into_owned()), ssh_socket_path: None, diff --git a/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs b/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs index 2d3894d563..e2d0518e24 100644 --- a/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs +++ b/crates/openshell-supervisor/src/legacy_startup_regression_tests.rs @@ -287,6 +287,7 @@ async fn config_push_regression_legacy_gateway_never_observes_ready_before_workl sandbox: Some("legacy-restart".into()), container_image: String::new(), openshell_endpoint: Some(endpoint), + gateway_home: None, policy_rules: None, policy_data: None, ssh_socket_path: None, diff --git a/crates/openshell-supervisor/src/lib.rs b/crates/openshell-supervisor/src/lib.rs index 8b1035beca..d2456a8e5e 100644 --- a/crates/openshell-supervisor/src/lib.rs +++ b/crates/openshell-supervisor/src/lib.rs @@ -896,6 +896,9 @@ pub struct SandboxInstanceConfig { /// Workload image recorded in this instance's OCSF events. pub container_image: String, pub openshell_endpoint: Option, + /// Peer endpoint of the gateway replica serving the shared host, which + /// the session dials instead of `openshell_endpoint` when set. + pub gateway_home: Option>>, pub policy_rules: Option, pub policy_data: Option, pub ssh_socket_path: Option, @@ -1060,6 +1063,7 @@ pub async fn run_sandbox_instance( sandbox, container_image: _, openshell_endpoint, + gateway_home, policy_rules, policy_data, ssh_socket_path, @@ -1704,6 +1708,7 @@ pub async fn run_sandbox_instance( config_apply_tx, config_push_updates, ssh_host_key, + gateway_home, ) .instrument(tracing::info_span!(parent: &startup, "supervisor.access.start")) .await?; @@ -4628,6 +4633,7 @@ mod tests { sandbox: None, container_image: String::new(), openshell_endpoint: None, + gateway_home: None, policy_rules: None, policy_data: None, ssh_socket_path: None, diff --git a/crates/openshell-supervisor/src/main.rs b/crates/openshell-supervisor/src/main.rs index 7c35b96826..e9b0f7cf82 100644 --- a/crates/openshell-supervisor/src/main.rs +++ b/crates/openshell-supervisor/src/main.rs @@ -607,6 +607,7 @@ fn main() -> Result<()> { sandbox: args.sandbox, container_image: std::env::var("OPENSHELL_CONTAINER_IMAGE").unwrap_or_default(), openshell_endpoint: args.openshell_endpoint, + gateway_home: None, policy_rules: args.policy_rules, policy_data: args.policy_data, ssh_socket_path: args.ssh_socket_path, diff --git a/crates/openshell-supervisor/src/shared_host.rs b/crates/openshell-supervisor/src/shared_host.rs index 28f49b426b..2541eabc70 100644 --- a/crates/openshell-supervisor/src/shared_host.rs +++ b/crates/openshell-supervisor/src/shared_host.rs @@ -209,6 +209,9 @@ pub struct SharedHost { leaving: watch::Sender, /// A gateway currently accepts this host's stream. gateway_connected: watch::Sender, + /// Peer endpoint of the gateway replica accepting this host's stream; + /// hosted instances dial their sessions there. + gateway_home: watch::Sender>, /// Recovery deadline (seconds) the latest snapshot offers the runtimes /// of instances this host starts or takes over; `None` keeps the runtime /// default. @@ -303,6 +306,7 @@ impl SharedHost { deferred_reports: Mutex::new(BTreeMap::new()), leaving: watch::Sender::new(false), gateway_connected: watch::Sender::new(false), + gateway_home: watch::Sender::new(None), recovery_timeout_secs: Mutex::new(None), #[cfg(test)] reconcile_probe: ReconcileProbe::default(), @@ -785,6 +789,7 @@ impl SharedHost { sandbox: Some(registration.sandbox_name), container_image: registration.container_image, openshell_endpoint: Some(self.gateway_endpoint.clone()), + gateway_home: Some(self.gateway_home.subscribe()), policy_rules: None, policy_data: None, ssh_socket_path: Some(ssh_socket.to_string_lossy().into_owned()), @@ -1417,6 +1422,9 @@ async fn run_stream( Some(host_gateway_message::Payload::Accepted(accepted_message)) => { accepted = true; host.gateway_connected.send_replace(true); + host.gateway_home.send_replace( + Some(accepted_message.peer_endpoint.clone()).filter(|endpoint| !endpoint.is_empty()), + ); tracing::info!(session = %accepted_message.session_id, "gateway accepted the supervisor host"); // A pool gateway detects failed members by silence, // so it asks for more frequent heartbeats. @@ -1517,6 +1525,7 @@ async fn run_stream( } }; host.gateway_connected.send_replace(false); + host.gateway_home.send_replace(None); end } @@ -1982,6 +1991,10 @@ mod tests { } fn accepted() -> HostGatewayMessage { + accepted_by("") + } + + fn accepted_by(peer_endpoint: &str) -> HostGatewayMessage { gateway_message(host_gateway_message::Payload::Accepted( openshell_core::proto::supervisor_host::v1::HostAccepted { session_id: "session".to_string(), @@ -1989,6 +2002,7 @@ mod tests { seconds: 1, nanos: 0, }), + peer_endpoint: peer_endpoint.to_string(), }, )) } @@ -2220,6 +2234,26 @@ mod tests { assert_eq!(gateway.calls.lock().unwrap().len(), 2); } + /// Hosted instances dial the replica that accepted the host, until its + /// stream is lost. + #[tokio::test(start_paused = true)] + async fn the_accepting_replica_is_the_hosted_instances_gateway() { + let (host, _gateway, mut streams, _root) = + attach_pool_host(vec![Attempt::Now, Attempt::Hang]); + let mut home = host.gateway_home.subscribe(); + let stream = streams.recv().await.unwrap(); + stream + .inbound + .send(Ok(accepted_by("https://gw-a:8443"))) + .await + .unwrap(); + home.wait_for(|home| home.as_deref() == Some("https://gw-a:8443")) + .await + .unwrap(); + drop(stream); + home.wait_for(Option::is_none).await.unwrap(); + } + /// A replica that keeps writing but stops reading fills the host's /// outbound queue; the blocked send gives up within the silence budget /// rather than the ordinary send bound. diff --git a/proto/supervisor_host.proto b/proto/supervisor_host.proto index a9c6782d16..862a5e3efa 100644 --- a/proto/supervisor_host.proto +++ b/proto/supervisor_host.proto @@ -84,6 +84,10 @@ message HostAccepted { // Scopes every snapshot revision and status report of this stream. string session_id = 1; google.protobuf.Duration heartbeat_interval = 2; + // Address of the accepting gateway replica. A pool host dials its sandbox + // instances' sessions here so they run on the replica serving this stream. + // Empty when the replica has no peer endpoint. + string peer_endpoint = 3; } // The gateway refuses the host; the stream ends after this message. From efda3de715ba68956004d49e319032a5e7331bba Mon Sep 17 00:00:00 2001 From: divesh Date: Fri, 30 Oct 2026 15:04:41 -0700 Subject: [PATCH 05/11] test(e2e): keep pool sessions on their supervisor's gateway through failover Add the shared-supervisor-tree lane: two gateway replicas on PostgreSQL with a two-member supervisor pool. After start, a gateway crash, a gateway rolling restart, and a supervisor crash, every pool sandbox's session runs on the replica serving its supervisor and exec reaches every sandbox. Signed-off-by: divesh --- crates/openshell-supervisor-control/README.md | 4 +- e2e/rust/Cargo.toml | 2 + .../kubernetes_shared_supervisor_tree.rs | 288 ++++++++++++++++++ tasks/test.toml | 5 + 4 files changed, 298 insertions(+), 1 deletion(-) create mode 100644 e2e/rust/tests/kubernetes_shared_supervisor_tree.rs diff --git a/crates/openshell-supervisor-control/README.md b/crates/openshell-supervisor-control/README.md index 58782a696a..1cb7f0d115 100644 --- a/crates/openshell-supervisor-control/README.md +++ b/crates/openshell-supervisor-control/README.md @@ -46,7 +46,9 @@ re-implementing the control contracts. stream attaches, ends, or announces leaving, so the gateway refreshes membership records at once. The gateway closes a member's stream after `pool_timing::GATEWAY_SILENCE_TIMEOUT` of silence, and drops a member that - stops reading as soon. + stops reading as soon. `HostAccepted` carries this replica's peer endpoint + (`set_peer_endpoint`); the host dials its instances' sandbox sessions + there, so each session runs on the replica serving its supervisor. - `pool_timing`: the timing contract between pool hosts and the gateway: the 15 s membership TTL, the 2 s record refresh, 1 s heartbeats, host silence detection and bounded reconnect attempts, and `reconnect_budget`, the worst diff --git a/e2e/rust/Cargo.toml b/e2e/rust/Cargo.toml index b40b170996..84a22b919c 100644 --- a/e2e/rust/Cargo.toml +++ b/e2e/rust/Cargo.toml @@ -44,6 +44,8 @@ e2e-kubernetes-shared-supervisor-ha = ["e2e-kubernetes", "e2e-host-gateway"] # Scaling a one-member high-availability pool of shared supervisors to two # members and back. It upgrades the release, so it runs on its own release. e2e-kubernetes-shared-supervisor-scaling = ["e2e-kubernetes", "e2e-host-gateway"] +# A shared supervisor pool behind two gateway replicas sharing PostgreSQL. +e2e-kubernetes-shared-supervisor-tree = ["e2e-kubernetes", "e2e-host-gateway"] e2e-podman = ["e2e", "e2e-host-gateway", "e2e-local-container-driver"] e2e-podman-gpu = ["e2e-podman", "e2e-gpu"] e2e-oidc-pkce = [] diff --git a/e2e/rust/tests/kubernetes_shared_supervisor_tree.rs b/e2e/rust/tests/kubernetes_shared_supervisor_tree.rs new file mode 100644 index 0000000000..da2c0d377f --- /dev/null +++ b/e2e/rust/tests/kubernetes_shared_supervisor_tree.rs @@ -0,0 +1,288 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![cfg(feature = "e2e-kubernetes-shared-supervisor-tree")] + +//! Two gateway replicas sharing a database and a two-member shared supervisor +//! pool. Every pool sandbox's session runs on the replica serving its +//! supervisor at start and after a gateway crash, a gateway rolling restart, +//! and a supervisor crash; exec reaches every sandbox after each step. + +mod shared_supervisor; + +use std::process::Stdio; +use std::time::Duration; + +use openshell_e2e::harness::sandbox::SandboxGuard; +use shared_supervisor::{ + EXEC_PROBE_TIMEOUT, GATEWAY_SELECTOR, kubectl, kubectl_command, release_namespace, + supervisor_pods, wait_until_attached, +}; + +const SANDBOXES: usize = 4; +const SETTLE_DEADLINE: Duration = Duration::from_mins(3); +const POSTGRES: &str = "deploy/openshell-e2e-postgres"; +/// Each granted pool sandbox with its owning member, the replica serving that +/// member's stream, and the replica owning the sandbox's session. +const TREE: &str = "\ +WITH io AS (SELECT convert_from(payload, 'UTF8')::json AS p FROM objects + WHERE object_type = 'supervisor_instance_owner'), + pm AS (SELECT convert_from(payload, 'UTF8')::json AS p FROM objects + WHERE object_type = 'supervisor_pool_member'), + so AS (SELECT substring(id FROM 'supervisor-owner:(.*)') AS sandbox_id, + convert_from(payload, 'UTF8')::json AS p FROM objects + WHERE object_type = 'supervisor_session_owner') +SELECT io.p->>'sandbox_id', io.p->>'member_id', coalesce(pm.p->>'gateway_replica', ''), + coalesce(so.p->>'owner_replica_id', '') +FROM io LEFT JOIN pm ON pm.p->>'member_id' = io.p->>'member_id' +LEFT JOIN so ON so.sandbox_id = io.p->>'sandbox_id' +WHERE io.p->>'phase' = 'granted'"; + +#[derive(Debug)] +struct Branch { + member: String, + supervisor_gateway: String, + session_gateway: String, +} + +async fn tree() -> Vec { + kubectl(&[ + "-n", + &release_namespace(), + "exec", + POSTGRES, + "--", + "psql", + "-U", + "openshell", + "-d", + "openshell", + "-XAt", + "-F", + "|", + "-c", + TREE, + ]) + .await + .lines() + .filter_map(|row| { + let mut fields = row.split('|').skip(1); + Some(Branch { + member: fields.next()?.to_string(), + supervisor_gateway: fields.next()?.to_string(), + session_gateway: fields.next()?.to_string(), + }) + }) + .collect() +} + +/// Wait until every sandbox's session runs on its supervisor's replica and no +/// record names a `gone` gateway or supervisor Pod. +async fn assert_tree(step: &str, gone: &[String]) -> Vec { + let mut last = Vec::new(); + let settled = tokio::time::timeout(SETTLE_DEADLINE, async { + loop { + last = tree().await; + if last.len() == SANDBOXES + && last.iter().all(|branch| { + !branch.session_gateway.is_empty() + && branch.session_gateway == branch.supervisor_gateway + && !gone.contains(&branch.member) + && !gone.contains(&branch.session_gateway) + }) + { + return; + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await; + assert!( + settled.is_ok(), + "{step}: sessions did not settle on their supervisor's gateway within {SETTLE_DEADLINE:?}: {last:#?}" + ); + last +} + +/// Reopens the CLI's port-forward when the gateway Pod behind it is gone. +struct GatewayForward { + port: u16, + child: Option, +} + +impl GatewayForward { + fn new() -> Self { + let config = std::env::var("XDG_CONFIG_HOME").expect("XDG_CONFIG_HOME"); + let gateway = std::env::var("OPENSHELL_GATEWAY").expect("OPENSHELL_GATEWAY"); + let metadata: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(format!( + "{config}/openshell/gateways/{gateway}/metadata.json" + )) + .expect("read gateway metadata"), + ) + .expect("parse gateway metadata"); + let port = metadata["gateway_port"] + .as_u64() + .and_then(|port| u16::try_from(port).ok()) + .expect("gateway_port"); + Self { port, child: None } + } + + async fn reachable(&self) -> bool { + tokio::net::TcpStream::connect(("127.0.0.1", self.port)) + .await + .is_ok() + } + + async fn ensure(&mut self) { + if self.reachable().await { + return; + } + self.child = Some( + kubectl_command() + .args([ + "-n", + &release_namespace(), + "port-forward", + &format!("svc/{}", release()), + &format!("{}:8080", self.port), + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .kill_on_drop(true) + .spawn() + .expect("spawn kubectl port-forward"), + ); + let opened = tokio::time::timeout(Duration::from_secs(30), async { + while !self.reachable().await { + tokio::time::sleep(Duration::from_millis(500)).await; + } + }) + .await; + assert!(opened.is_ok(), "port-forward on {} never opened", self.port); + } +} + +async fn assert_exec(step: &str, forward: &mut GatewayForward, sandboxes: &[SandboxGuard]) { + for sandbox in sandboxes { + let mut last = String::new(); + let reached = tokio::time::timeout(SETTLE_DEADLINE, async { + loop { + forward.ensure().await; + match tokio::time::timeout(EXEC_PROBE_TIMEOUT, sandbox.exec(&["true"])).await { + Ok(Ok(_)) => return, + Ok(Err(error)) => last = error, + Err(_) => last = "exec timed out".to_string(), + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await; + assert!(reached.is_ok(), "{step}: exec into {} failed: {last}", sandbox.name); + } +} + +fn release() -> String { + std::env::var("OPENSHELL_E2E_KUBE_RELEASE").expect("OPENSHELL_E2E_KUBE_RELEASE") +} + +async fn gateway_pods() -> Vec { + kubectl(&[ + "-n", + &release_namespace(), + "get", + "pods", + "-l", + GATEWAY_SELECTOR, + "-o", + "jsonpath={.items[*].metadata.name}", + ]) + .await + .split_whitespace() + .map(str::to_string) + .collect() +} + +/// The gateway's workload, a Deployment or a `StatefulSet`. +async fn gateway_workload() -> String { + kubectl(&[ + "-n", + &release_namespace(), + "get", + "deploy,statefulset", + "--field-selector", + &format!("metadata.name={}", release()), + "-o", + "name", + ]) + .await + .trim() + .to_string() +} + +async fn rollout(action: &str, workload: &str) { + kubectl(&[ + "-n", + &release_namespace(), + "rollout", + action, + workload, + "--timeout=300s", + ]) + .await; +} + +async fn force_delete(pod: &str) { + kubectl(&[ + "-n", + &release_namespace(), + "delete", + "pod", + pod, + "--grace-period=0", + "--force", + ]) + .await; +} + +#[tokio::test] +async fn pool_sessions_follow_their_supervisors_gateway() { + let gateway = gateway_workload().await; + let supervisor = format!("deployment.apps/{}-supervisor", release()); + assert_eq!(gateway_pods().await.len(), 2, "the lane runs two gateway replicas"); + wait_until_attached(&supervisor_pods().await).await; + + let mut forward = GatewayForward::new(); + forward.ensure().await; + let mut sandboxes = Vec::new(); + for _ in 0..SANDBOXES { + sandboxes.push( + SandboxGuard::create_detached_main_with_args(&[], &["sleep", "infinity"]) + .await + .expect("create sandbox"), + ); + } + let branches = assert_tree("start", &[]).await; + assert_exec("start", &mut forward, &sandboxes).await; + + let crashed = vec![branches[0].supervisor_gateway.clone()]; + force_delete(&crashed[0]).await; + rollout("status", &gateway).await; + assert_tree("gateway crash", &crashed).await; + assert_exec("gateway crash", &mut forward, &sandboxes).await; + + let restarted = gateway_pods().await; + kubectl(&["-n", &release_namespace(), "rollout", "restart", &gateway]).await; + rollout("status", &gateway).await; + assert_tree("gateway rolling restart", &restarted).await; + assert_exec("gateway rolling restart", &mut forward, &sandboxes).await; + + let crashed = vec![supervisor_pods().await.remove(0)]; + force_delete(&crashed[0]).await; + rollout("status", &supervisor).await; + assert_tree("supervisor crash", &crashed).await; + assert_exec("supervisor crash", &mut forward, &sandboxes).await; + + forward.ensure().await; + drop(sandboxes); +} diff --git a/tasks/test.toml b/tasks/test.toml index 6bafa95e9e..efa72e0dfb 100644 --- a/tasks/test.toml +++ b/tasks/test.toml @@ -307,6 +307,11 @@ env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_scaling", OPENSH # test. The shared supervisor requires gateway TLS, which the Skaffold values turn off. run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor-ha-scaling.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' +["e2e:kubernetes:shared-supervisor-tree"] +description = "Run Kubernetes e2e that crashes and restarts gateway replicas and a supervisor, checking each pool sandbox's session stays on its supervisor's gateway" +env = { OPENSHELL_E2E_KUBE_EXTERNAL_POSTGRES_SECRET = "openshell-ha-pg", OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_tree", OPENSHELL_E2E_KUBE_TLS = "1", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor-tree" } +run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-high-availability.yaml:deploy/helm/openshell/ci/values-supervisor-ha.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' + ["e2e:kubernetes:workspace-operator"] description = "Run Kubernetes e2e with operator workspace mode (pre-provisioned per-workspace namespaces)" env = { OPENSHELL_E2E_KUBE_EXTRA_VALUES = "deploy/helm/openshell/ci/values-workspace-operator.yaml", OPENSHELL_E2E_KUBE_TEST = "workspace_namespace_operator", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-kubernetes,e2e-kubernetes-workspace-operator" } From 3370bc8f0ca9c6307a91fdd7d6914ea54d6f8756 Mon Sep 17 00:00:00 2001 From: divesh Date: Sat, 31 Oct 2026 18:12:53 -0700 Subject: [PATCH 06/11] fix(supervisor): move established pool sessions when their host changes gateway A host whose stream ends, for example when its supervisor token expires, can reattach to another replica. Its sandbox sessions now watch for that and reconnect to the host's new replica instead of staying on the old one. Signed-off-by: divesh --- .../src/supervisor_session.rs | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/crates/openshell-supervisor-process/src/supervisor_session.rs b/crates/openshell-supervisor-process/src/supervisor_session.rs index e3c32011d9..427af34d12 100644 --- a/crates/openshell-supervisor-process/src/supervisor_session.rs +++ b/crates/openshell-supervisor-process/src/supervisor_session.rs @@ -516,6 +516,17 @@ async fn run_session_loop(config: SessionConfig) { follow_home = false; // No backoff: this is an expected handoff, not a failure. } + Ok(SessionOutcome::HomeMoved { peer_endpoint }) => { + backoff.on_redirect(config.ready_tx.send_replace(false)); + info!( + sandbox_id = %config.sandbox_id, + peer_endpoint = %peer_endpoint, + "supervisor session: following the shared host to another gateway" + ); + target = config.gateway.with_url(peer_endpoint); + redirected = true; + follow_home = false; + } Err(e) => { let accepted = config.ready_tx.send_replace(false); let event = session_failed_event( @@ -598,6 +609,23 @@ async fn wait_for_home( .and_then(|home| (*home).clone()) } +/// The peer endpoint the shared host moves to, once it changes to one other +/// than `current`. Pending forever without a shared host. +async fn home_moved(home: Option<&mut watch::Receiver>>, current: &str) -> String { + let Some(home) = home else { + return std::future::pending().await; + }; + loop { + if home.changed().await.is_err() { + return std::future::pending().await; + } + let endpoint = home.borrow_and_update().clone(); + if let Some(endpoint) = endpoint.filter(|endpoint| endpoint != current) { + return endpoint; + } + } +} + /// Keep asking not to be redirected when a redirected attempt never got a /// session, so a stale ring pointing at a dead replica cannot bounce us back /// to it. A session that was accepted and later dropped starts placement over. @@ -622,6 +650,9 @@ enum SessionOutcome { peer_endpoint: String, owner_replica_id: String, }, + /// The shared host reattached to another replica. The caller reconnects + /// there. + HomeMoved { peer_endpoint: String }, } async fn run_single_session( @@ -736,9 +767,16 @@ async fn run_single_session( // Main loop: receive gateway messages + send heartbeats. let mut heartbeat_interval = tokio::time::interval(Duration::from_secs(heartbeat_secs)); heartbeat_interval.tick().await; // skip immediate tick + let mut home = config.gateway_home.clone(); + if let Some(home) = home.as_mut() { + home.mark_unchanged(); + } loop { tokio::select! { + peer_endpoint = home_moved(home.as_mut(), target.url()) => { + return Ok(SessionOutcome::HomeMoved { peer_endpoint }); + } msg = inbound.message() => { let msg = match map_session_stream_message( msg, @@ -1704,6 +1742,25 @@ mod ocsf_event_tests { assert_eq!(moved.await.unwrap().as_deref(), Some("https://gw-b:8443")); } + #[tokio::test] + async fn an_established_session_follows_its_host_to_another_gateway() { + let (tx, mut home) = watch::channel(Some("https://gw-a:8443".to_string())); + home.mark_unchanged(); + let moved = + tokio::spawn(async move { home_moved(Some(&mut home), "https://gw-a:8443").await }); + + tx.send_replace(None); + tx.send_replace(Some("https://gw-a:8443".to_string())); + tokio::task::yield_now().await; + assert!( + !moved.is_finished(), + "a reattach to the same gateway is not a move" + ); + + tx.send_replace(Some("https://gw-b:8443".to_string())); + assert_eq!(moved.await.unwrap(), "https://gw-b:8443"); + } + #[test] fn reconnect_backoff_doubles_until_a_session_is_accepted() { let mut backoff = ReconnectBackoff::new(); From 68a95013b4ef832a7b33a8cfc03255dc628e8ad9 Mon Sep 17 00:00:00 2001 From: divesh Date: Sat, 31 Oct 2026 18:12:53 -0700 Subject: [PATCH 07/11] test(e2e): fold the supervisor tree checks into the shared-supervisor-ha lane Run the shared-supervisor-ha lane with two gateway replicas on PostgreSQL and move the gateway crash, rolling restart, and supervisor crash colocation test into it, replacing the separate shared-supervisor-tree lane. Signed-off-by: divesh --- .config/nextest.toml | 8 +- TESTING.md | 2 +- e2e/rust/Cargo.toml | 2 - .../tests/kubernetes_shared_supervisor_ha.rs | 295 +++++++++++++++++- .../kubernetes_shared_supervisor_tree.rs | 288 ----------------- tasks/test.toml | 11 +- 6 files changed, 300 insertions(+), 306 deletions(-) delete mode 100644 e2e/rust/tests/kubernetes_shared_supervisor_tree.rs diff --git a/.config/nextest.toml b/.config/nextest.toml index 1f3f5d0364..b159e1db2c 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -32,11 +32,11 @@ kubernetes-shared-supervisor-pool = { max-threads = 1 } filter = "test(/gateway_(scale_and_rollout|pod_rolls)$/)" test-group = "kubernetes-ha" -# Both change the shared supervisor pool of the release they run against: one -# kills a member, the other scales the pool. Never run them at once. Test -# names, not binary(), for the reason above. +# These change the shared supervisor pool of the release they run against: +# they kill members or gateways, or scale the pool. Never run them at once. +# Test names, not binary(), for the reason above. [[profile.e2e-kubernetes.overrides]] -filter = "test(/^(losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them|scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running)$/)" +filter = "test(/^(losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them|scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running|pool_sessions_follow_their_supervisors_gateway)$/)" test-group = "kubernetes-shared-supervisor-pool" # Scaling the pool runs two bounded Helm upgrades (up to 6 minutes each), waits diff --git a/TESTING.md b/TESTING.md index 9649ba2515..6b99187825 100644 --- a/TESTING.md +++ b/TESTING.md @@ -672,7 +672,7 @@ Available task variants: | `e2e:kubernetes:v1alpha1` | Agent Sandbox v1alpha1 compatibility | | `e2e:kubernetes:external-driver` | External Kubernetes driver sidecar | | `e2e:kubernetes:shared-supervisor` | One-member shared supervisor pool for the gateway's sandboxes (gateway TLS through the port-forward) | -| `e2e:kubernetes:shared-supervisor-ha` | Two-replica shared supervisor pool: kills the member that owns a sandbox and checks the sandbox is adopted, not restarted | +| `e2e:kubernetes:shared-supervisor-ha` | Two-replica shared supervisor pool behind two gateway replicas on PostgreSQL: kills the member that owns a sandbox and checks the sandbox is adopted, not restarted; crashes and restarts gateways and checks each sandbox's session stays on its supervisor's gateway | | `e2e:kubernetes:shared-supervisor-scaling` | One-member shared supervisor pool: scales it to two members and back with `helm upgrade`, checking sandboxes stay running and the removed member hands its sandboxes off | Kubernetes e2e environment variables: diff --git a/e2e/rust/Cargo.toml b/e2e/rust/Cargo.toml index 84a22b919c..b40b170996 100644 --- a/e2e/rust/Cargo.toml +++ b/e2e/rust/Cargo.toml @@ -44,8 +44,6 @@ e2e-kubernetes-shared-supervisor-ha = ["e2e-kubernetes", "e2e-host-gateway"] # Scaling a one-member high-availability pool of shared supervisors to two # members and back. It upgrades the release, so it runs on its own release. e2e-kubernetes-shared-supervisor-scaling = ["e2e-kubernetes", "e2e-host-gateway"] -# A shared supervisor pool behind two gateway replicas sharing PostgreSQL. -e2e-kubernetes-shared-supervisor-tree = ["e2e-kubernetes", "e2e-host-gateway"] e2e-podman = ["e2e", "e2e-host-gateway", "e2e-local-container-driver"] e2e-podman-gpu = ["e2e-podman", "e2e-gpu"] e2e-oidc-pkce = [] diff --git a/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs b/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs index 896549305b..95265e8e14 100644 --- a/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs +++ b/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs @@ -13,16 +13,23 @@ //! process's in-memory state survive; exec and policy-mediated networking work //! again; and a sandbox on another member keeps running and answering //! throughout, under the same owner. +//! +//! The pool sits behind two gateway replicas sharing a database. Every pool +//! sandbox's session runs on the replica serving its supervisor at start and +//! after a gateway crash, a gateway rolling restart, and a supervisor crash; +//! exec reaches every sandbox after each step. mod shared_supervisor; use std::process::Stdio; use std::time::Duration; +use openshell_e2e::harness::sandbox::SandboxGuard; use shared_supervisor::{ - HostEndpoint, SUPERVISOR_CONTAINER, call_host, counts_past, create, gateway_log_lines, - generation, is_takeover, kubectl, kubectl_command, kubectl_json, main_process, owner_of, - policy, release_namespace, supervisor_pods, wait_for_owner, wait_until_attached, + EXEC_PROBE_TIMEOUT, GATEWAY_SELECTOR, HostEndpoint, SUPERVISOR_CONTAINER, call_host, + counts_past, create, gateway_log_lines, generation, is_takeover, kubectl, kubectl_command, + kubectl_json, main_process, owner_of, policy, release_namespace, supervisor_pods, + wait_for_owner, wait_until_attached, }; use tokio::io::AsyncBufReadExt as _; @@ -387,3 +394,285 @@ async fn losing_a_supervisor_pod_moves_its_sandboxes_without_restarting_them() { first.cleanup().await; second.cleanup().await; } + +const SANDBOXES: usize = 4; +const SETTLE_DEADLINE: Duration = Duration::from_mins(3); +const POSTGRES: &str = "deploy/openshell-e2e-postgres"; +/// Each granted pool sandbox's route with its owning member, the replica +/// serving that member's stream, and the replica owning the sandbox's session. +const TREE: &str = "\ +WITH io AS (SELECT convert_from(payload, 'UTF8')::json AS p FROM objects + WHERE object_type = 'supervisor_instance_owner'), + pm AS (SELECT convert_from(payload, 'UTF8')::json AS p FROM objects + WHERE object_type = 'supervisor_pool_member'), + so AS (SELECT substring(id FROM 'supervisor-owner:(.*)') AS sandbox_id, + convert_from(payload, 'UTF8')::json AS p FROM objects + WHERE object_type = 'supervisor_session_owner') +SELECT io.p->>'route', io.p->>'member_id', coalesce(pm.p->>'gateway_replica', ''), + coalesce(so.p->>'owner_replica_id', '') +FROM io LEFT JOIN pm ON pm.p->>'member_id' = io.p->>'member_id' +LEFT JOIN so ON so.sandbox_id = io.p->>'sandbox_id' +WHERE io.p->>'phase' = 'granted'"; + +#[derive(Debug)] +struct Branch { + member: String, + supervisor_gateway: String, + session_gateway: String, +} + +async fn tree(routes: &[String]) -> Vec { + kubectl(&[ + "-n", + &release_namespace(), + "exec", + POSTGRES, + "--", + "psql", + "-U", + "openshell", + "-d", + "openshell", + "-XAt", + "-F", + "|", + "-c", + TREE, + ]) + .await + .lines() + .filter_map(|row| { + let mut fields = row.split('|'); + let route = fields.next()?; + if !routes.iter().any(|ours| ours == route) { + return None; + } + Some(Branch { + member: fields.next()?.to_string(), + supervisor_gateway: fields.next()?.to_string(), + session_gateway: fields.next()?.to_string(), + }) + }) + .collect() +} + +/// Wait until the session of every sandbox on `routes` runs on its +/// supervisor's replica and no record names a `gone` gateway or supervisor Pod. +async fn assert_tree(step: &str, routes: &[String], gone: &[String]) -> Vec { + let mut last = Vec::new(); + let settled = tokio::time::timeout(SETTLE_DEADLINE, async { + loop { + last = tree(routes).await; + if last.len() == routes.len() + && last.iter().all(|branch| { + !branch.session_gateway.is_empty() + && branch.session_gateway == branch.supervisor_gateway + && !gone.contains(&branch.member) + && !gone.contains(&branch.session_gateway) + }) + { + return; + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await; + assert!( + settled.is_ok(), + "{step}: sessions did not settle on their supervisor's gateway within {SETTLE_DEADLINE:?}: {last:#?}" + ); + last +} + +/// Reopens the CLI's port-forward when the gateway Pod behind it is gone. The +/// new port-forward outlives the test, so later tests keep the CLI's address. +struct GatewayForward { + port: u16, +} + +impl GatewayForward { + fn new() -> Self { + let config = std::env::var("XDG_CONFIG_HOME").expect("XDG_CONFIG_HOME"); + let gateway = std::env::var("OPENSHELL_GATEWAY").expect("OPENSHELL_GATEWAY"); + let metadata: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(format!( + "{config}/openshell/gateways/{gateway}/metadata.json" + )) + .expect("read gateway metadata"), + ) + .expect("parse gateway metadata"); + let port = metadata["gateway_port"] + .as_u64() + .and_then(|port| u16::try_from(port).ok()) + .expect("gateway_port"); + Self { port } + } + + async fn reachable(&self) -> bool { + tokio::net::TcpStream::connect(("127.0.0.1", self.port)) + .await + .is_ok() + } + + async fn ensure(&self) { + if self.reachable().await { + return; + } + kubectl_command() + .args([ + "-n", + &release_namespace(), + "port-forward", + &format!("svc/{}", release()), + &format!("{}:8080", self.port), + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .kill_on_drop(false) + .spawn() + .expect("spawn kubectl port-forward"); + let opened = tokio::time::timeout(Duration::from_secs(30), async { + while !self.reachable().await { + tokio::time::sleep(Duration::from_millis(500)).await; + } + }) + .await; + assert!(opened.is_ok(), "port-forward on {} never opened", self.port); + } +} + +async fn assert_exec(step: &str, forward: &GatewayForward, sandboxes: &[SandboxGuard]) { + for sandbox in sandboxes { + let mut last = String::new(); + let reached = tokio::time::timeout(SETTLE_DEADLINE, async { + loop { + forward.ensure().await; + match tokio::time::timeout(EXEC_PROBE_TIMEOUT, sandbox.exec(&["true"])).await { + Ok(Ok(_)) => return, + Ok(Err(error)) => last = error, + Err(_) => last = "exec timed out".to_string(), + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + .await; + assert!( + reached.is_ok(), + "{step}: exec into {} failed: {last}", + sandbox.name + ); + } +} + +fn release() -> String { + std::env::var("OPENSHELL_E2E_KUBE_RELEASE").expect("OPENSHELL_E2E_KUBE_RELEASE") +} + +async fn gateway_pods() -> Vec { + kubectl(&[ + "-n", + &release_namespace(), + "get", + "pods", + "-l", + GATEWAY_SELECTOR, + "-o", + "jsonpath={.items[*].metadata.name}", + ]) + .await + .split_whitespace() + .map(str::to_string) + .collect() +} + +/// The gateway's workload, a Deployment or a `StatefulSet`. +async fn gateway_workload() -> String { + kubectl(&[ + "-n", + &release_namespace(), + "get", + "deploy,statefulset", + "--field-selector", + &format!("metadata.name={}", release()), + "-o", + "name", + ]) + .await + .trim() + .to_string() +} + +async fn rollout(action: &str, workload: &str) { + kubectl(&[ + "-n", + &release_namespace(), + "rollout", + action, + workload, + "--timeout=300s", + ]) + .await; +} + +async fn force_delete(pod: &str) { + kubectl(&[ + "-n", + &release_namespace(), + "delete", + "pod", + pod, + "--grace-period=0", + "--force", + ]) + .await; +} + +#[tokio::test] +async fn pool_sessions_follow_their_supervisors_gateway() { + let gateway = gateway_workload().await; + let supervisor = format!("deployment.apps/{}-supervisor", release()); + assert_eq!( + gateway_pods().await.len(), + 2, + "the lane runs two gateway replicas" + ); + wait_until_attached(&supervisor_pods().await).await; + + let forward = GatewayForward::new(); + forward.ensure().await; + let mut sandboxes = Vec::new(); + for _ in 0..SANDBOXES { + sandboxes.push( + SandboxGuard::create_detached_main_with_args(&[], &["sleep", "infinity"]) + .await + .expect("create sandbox"), + ); + } + let mut routes = Vec::new(); + for sandbox in &sandboxes { + routes.push(generation(&sandbox.name).await.route); + } + let branches = assert_tree("start", &routes, &[]).await; + assert_exec("start", &forward, &sandboxes).await; + + let crashed = vec![branches[0].supervisor_gateway.clone()]; + force_delete(&crashed[0]).await; + rollout("status", &gateway).await; + assert_tree("gateway crash", &routes, &crashed).await; + assert_exec("gateway crash", &forward, &sandboxes).await; + + let restarted = gateway_pods().await; + kubectl(&["-n", &release_namespace(), "rollout", "restart", &gateway]).await; + rollout("status", &gateway).await; + let branches = assert_tree("gateway rolling restart", &routes, &restarted).await; + assert_exec("gateway rolling restart", &forward, &sandboxes).await; + + let crashed = vec![branches[0].member.clone()]; + force_delete(&crashed[0]).await; + rollout("status", &supervisor).await; + assert_tree("supervisor crash", &routes, &crashed).await; + assert_exec("supervisor crash", &forward, &sandboxes).await; + + forward.ensure().await; + drop(sandboxes); +} diff --git a/e2e/rust/tests/kubernetes_shared_supervisor_tree.rs b/e2e/rust/tests/kubernetes_shared_supervisor_tree.rs deleted file mode 100644 index da2c0d377f..0000000000 --- a/e2e/rust/tests/kubernetes_shared_supervisor_tree.rs +++ /dev/null @@ -1,288 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -#![cfg(feature = "e2e-kubernetes-shared-supervisor-tree")] - -//! Two gateway replicas sharing a database and a two-member shared supervisor -//! pool. Every pool sandbox's session runs on the replica serving its -//! supervisor at start and after a gateway crash, a gateway rolling restart, -//! and a supervisor crash; exec reaches every sandbox after each step. - -mod shared_supervisor; - -use std::process::Stdio; -use std::time::Duration; - -use openshell_e2e::harness::sandbox::SandboxGuard; -use shared_supervisor::{ - EXEC_PROBE_TIMEOUT, GATEWAY_SELECTOR, kubectl, kubectl_command, release_namespace, - supervisor_pods, wait_until_attached, -}; - -const SANDBOXES: usize = 4; -const SETTLE_DEADLINE: Duration = Duration::from_mins(3); -const POSTGRES: &str = "deploy/openshell-e2e-postgres"; -/// Each granted pool sandbox with its owning member, the replica serving that -/// member's stream, and the replica owning the sandbox's session. -const TREE: &str = "\ -WITH io AS (SELECT convert_from(payload, 'UTF8')::json AS p FROM objects - WHERE object_type = 'supervisor_instance_owner'), - pm AS (SELECT convert_from(payload, 'UTF8')::json AS p FROM objects - WHERE object_type = 'supervisor_pool_member'), - so AS (SELECT substring(id FROM 'supervisor-owner:(.*)') AS sandbox_id, - convert_from(payload, 'UTF8')::json AS p FROM objects - WHERE object_type = 'supervisor_session_owner') -SELECT io.p->>'sandbox_id', io.p->>'member_id', coalesce(pm.p->>'gateway_replica', ''), - coalesce(so.p->>'owner_replica_id', '') -FROM io LEFT JOIN pm ON pm.p->>'member_id' = io.p->>'member_id' -LEFT JOIN so ON so.sandbox_id = io.p->>'sandbox_id' -WHERE io.p->>'phase' = 'granted'"; - -#[derive(Debug)] -struct Branch { - member: String, - supervisor_gateway: String, - session_gateway: String, -} - -async fn tree() -> Vec { - kubectl(&[ - "-n", - &release_namespace(), - "exec", - POSTGRES, - "--", - "psql", - "-U", - "openshell", - "-d", - "openshell", - "-XAt", - "-F", - "|", - "-c", - TREE, - ]) - .await - .lines() - .filter_map(|row| { - let mut fields = row.split('|').skip(1); - Some(Branch { - member: fields.next()?.to_string(), - supervisor_gateway: fields.next()?.to_string(), - session_gateway: fields.next()?.to_string(), - }) - }) - .collect() -} - -/// Wait until every sandbox's session runs on its supervisor's replica and no -/// record names a `gone` gateway or supervisor Pod. -async fn assert_tree(step: &str, gone: &[String]) -> Vec { - let mut last = Vec::new(); - let settled = tokio::time::timeout(SETTLE_DEADLINE, async { - loop { - last = tree().await; - if last.len() == SANDBOXES - && last.iter().all(|branch| { - !branch.session_gateway.is_empty() - && branch.session_gateway == branch.supervisor_gateway - && !gone.contains(&branch.member) - && !gone.contains(&branch.session_gateway) - }) - { - return; - } - tokio::time::sleep(Duration::from_secs(2)).await; - } - }) - .await; - assert!( - settled.is_ok(), - "{step}: sessions did not settle on their supervisor's gateway within {SETTLE_DEADLINE:?}: {last:#?}" - ); - last -} - -/// Reopens the CLI's port-forward when the gateway Pod behind it is gone. -struct GatewayForward { - port: u16, - child: Option, -} - -impl GatewayForward { - fn new() -> Self { - let config = std::env::var("XDG_CONFIG_HOME").expect("XDG_CONFIG_HOME"); - let gateway = std::env::var("OPENSHELL_GATEWAY").expect("OPENSHELL_GATEWAY"); - let metadata: serde_json::Value = serde_json::from_str( - &std::fs::read_to_string(format!( - "{config}/openshell/gateways/{gateway}/metadata.json" - )) - .expect("read gateway metadata"), - ) - .expect("parse gateway metadata"); - let port = metadata["gateway_port"] - .as_u64() - .and_then(|port| u16::try_from(port).ok()) - .expect("gateway_port"); - Self { port, child: None } - } - - async fn reachable(&self) -> bool { - tokio::net::TcpStream::connect(("127.0.0.1", self.port)) - .await - .is_ok() - } - - async fn ensure(&mut self) { - if self.reachable().await { - return; - } - self.child = Some( - kubectl_command() - .args([ - "-n", - &release_namespace(), - "port-forward", - &format!("svc/{}", release()), - &format!("{}:8080", self.port), - ]) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .kill_on_drop(true) - .spawn() - .expect("spawn kubectl port-forward"), - ); - let opened = tokio::time::timeout(Duration::from_secs(30), async { - while !self.reachable().await { - tokio::time::sleep(Duration::from_millis(500)).await; - } - }) - .await; - assert!(opened.is_ok(), "port-forward on {} never opened", self.port); - } -} - -async fn assert_exec(step: &str, forward: &mut GatewayForward, sandboxes: &[SandboxGuard]) { - for sandbox in sandboxes { - let mut last = String::new(); - let reached = tokio::time::timeout(SETTLE_DEADLINE, async { - loop { - forward.ensure().await; - match tokio::time::timeout(EXEC_PROBE_TIMEOUT, sandbox.exec(&["true"])).await { - Ok(Ok(_)) => return, - Ok(Err(error)) => last = error, - Err(_) => last = "exec timed out".to_string(), - } - tokio::time::sleep(Duration::from_secs(2)).await; - } - }) - .await; - assert!(reached.is_ok(), "{step}: exec into {} failed: {last}", sandbox.name); - } -} - -fn release() -> String { - std::env::var("OPENSHELL_E2E_KUBE_RELEASE").expect("OPENSHELL_E2E_KUBE_RELEASE") -} - -async fn gateway_pods() -> Vec { - kubectl(&[ - "-n", - &release_namespace(), - "get", - "pods", - "-l", - GATEWAY_SELECTOR, - "-o", - "jsonpath={.items[*].metadata.name}", - ]) - .await - .split_whitespace() - .map(str::to_string) - .collect() -} - -/// The gateway's workload, a Deployment or a `StatefulSet`. -async fn gateway_workload() -> String { - kubectl(&[ - "-n", - &release_namespace(), - "get", - "deploy,statefulset", - "--field-selector", - &format!("metadata.name={}", release()), - "-o", - "name", - ]) - .await - .trim() - .to_string() -} - -async fn rollout(action: &str, workload: &str) { - kubectl(&[ - "-n", - &release_namespace(), - "rollout", - action, - workload, - "--timeout=300s", - ]) - .await; -} - -async fn force_delete(pod: &str) { - kubectl(&[ - "-n", - &release_namespace(), - "delete", - "pod", - pod, - "--grace-period=0", - "--force", - ]) - .await; -} - -#[tokio::test] -async fn pool_sessions_follow_their_supervisors_gateway() { - let gateway = gateway_workload().await; - let supervisor = format!("deployment.apps/{}-supervisor", release()); - assert_eq!(gateway_pods().await.len(), 2, "the lane runs two gateway replicas"); - wait_until_attached(&supervisor_pods().await).await; - - let mut forward = GatewayForward::new(); - forward.ensure().await; - let mut sandboxes = Vec::new(); - for _ in 0..SANDBOXES { - sandboxes.push( - SandboxGuard::create_detached_main_with_args(&[], &["sleep", "infinity"]) - .await - .expect("create sandbox"), - ); - } - let branches = assert_tree("start", &[]).await; - assert_exec("start", &mut forward, &sandboxes).await; - - let crashed = vec![branches[0].supervisor_gateway.clone()]; - force_delete(&crashed[0]).await; - rollout("status", &gateway).await; - assert_tree("gateway crash", &crashed).await; - assert_exec("gateway crash", &mut forward, &sandboxes).await; - - let restarted = gateway_pods().await; - kubectl(&["-n", &release_namespace(), "rollout", "restart", &gateway]).await; - rollout("status", &gateway).await; - assert_tree("gateway rolling restart", &restarted).await; - assert_exec("gateway rolling restart", &mut forward, &sandboxes).await; - - let crashed = vec![supervisor_pods().await.remove(0)]; - force_delete(&crashed[0]).await; - rollout("status", &supervisor).await; - assert_tree("supervisor crash", &crashed).await; - assert_exec("supervisor crash", &mut forward, &sandboxes).await; - - forward.ensure().await; - drop(sandboxes); -} diff --git a/tasks/test.toml b/tasks/test.toml index efa72e0dfb..724d0faa31 100644 --- a/tasks/test.toml +++ b/tasks/test.toml @@ -295,10 +295,10 @@ env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor", OPENSHELL_E2E_ run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' ["e2e:kubernetes:shared-supervisor-ha"] -description = "Run Kubernetes e2e that kills the owning member of a two-replica shared supervisor pool and checks its sandboxes are adopted, not restarted" -env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_ha", OPENSHELL_E2E_KUBE_TLS = "1", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor-ha" } +description = "Run Kubernetes e2e for a two-replica shared supervisor pool behind two gateway replicas: a killed owning member's sandboxes are adopted, not restarted, and each pool sandbox's session stays on its supervisor's gateway through gateway and supervisor failures" +env = { OPENSHELL_E2E_KUBE_EXTERNAL_POSTGRES_SECRET = "openshell-ha-pg", OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_ha", OPENSHELL_E2E_KUBE_TLS = "1", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor-ha" } # The shared supervisor requires gateway TLS, which the Skaffold values turn off. -run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor-ha.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' +run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-high-availability.yaml:deploy/helm/openshell/ci/values-supervisor-ha.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' ["e2e:kubernetes:shared-supervisor-scaling"] description = "Run Kubernetes e2e that scales a one-member shared supervisor pool to two members and back to one with Helm, checking sandboxes are handed off, not restarted" @@ -307,11 +307,6 @@ env = { OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_scaling", OPENSH # test. The shared supervisor requires gateway TLS, which the Skaffold values turn off. run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-supervisor-ha-scaling.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' -["e2e:kubernetes:shared-supervisor-tree"] -description = "Run Kubernetes e2e that crashes and restarts gateway replicas and a supervisor, checking each pool sandbox's session stays on its supervisor's gateway" -env = { OPENSHELL_E2E_KUBE_EXTERNAL_POSTGRES_SECRET = "openshell-ha-pg", OPENSHELL_E2E_KUBE_TEST = "kubernetes_shared_supervisor_tree", OPENSHELL_E2E_KUBE_TLS = "1", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-host-gateway,e2e-kubernetes,e2e-kubernetes-shared-supervisor-tree" } -run = 'OPENSHELL_E2E_KUBE_EXTRA_VALUES="deploy/helm/openshell/ci/values-high-availability.yaml:deploy/helm/openshell/ci/values-supervisor-ha.yaml:deploy/helm/openshell/ci/values-tls-port-forward.yaml${OPENSHELL_E2E_KUBE_EXTRA_VALUES:+:$OPENSHELL_E2E_KUBE_EXTRA_VALUES}" e2e/rust/e2e-kubernetes.sh' - ["e2e:kubernetes:workspace-operator"] description = "Run Kubernetes e2e with operator workspace mode (pre-provisioned per-workspace namespaces)" env = { OPENSHELL_E2E_KUBE_EXTRA_VALUES = "deploy/helm/openshell/ci/values-workspace-operator.yaml", OPENSHELL_E2E_KUBE_TEST = "workspace_namespace_operator", OPENSHELL_E2E_KUBERNETES_FEATURES = "e2e,e2e-kubernetes,e2e-kubernetes-workspace-operator" } From 403902950974723bf6cac613193d24ede862fe9a Mon Sep 17 00:00:00 2001 From: divesh Date: Sat, 31 Oct 2026 18:56:52 -0700 Subject: [PATCH 08/11] fix(supervisor): redial the accepting gateway replica after a host stream ends A host's stream ends routinely, for example when its supervisor token expires. Redialing the replica that accepted it, before falling back to the Service, keeps the host and its sandbox sessions on that replica instead of moving them all on every reconnect. Signed-off-by: divesh --- .../openshell-supervisor/src/shared_host.rs | 87 +++++++++++++++---- 1 file changed, 69 insertions(+), 18 deletions(-) diff --git a/crates/openshell-supervisor/src/shared_host.rs b/crates/openshell-supervisor/src/shared_host.rs index 2541eabc70..d912965ca3 100644 --- a/crates/openshell-supervisor/src/shared_host.rs +++ b/crates/openshell-supervisor/src/shared_host.rs @@ -989,8 +989,13 @@ fn status_message(host: &SharedHost, applied_revision: u64) -> HostMessage { /// Why one gateway stream ended. enum StreamEnd { /// The gateway rejected or closed the stream; reconnect after a backoff. - /// `accepted` records whether the gateway had accepted the host. - Lost { accepted: bool }, + /// `accepted` records whether the gateway had accepted the host, and + /// `redial` is the replica that accepted it: the next attempt goes there + /// first, as after a token expiry, so the hosted sessions stay in place. + Lost { + accepted: bool, + redial: Option, + }, } /// What a gateway sends a host on its assignment stream. @@ -1001,28 +1006,32 @@ type GatewayInbound = std::pin::Pin< /// Opens a host's assignment stream: hands the gateway the host's outbound /// messages and returns the gateway's. trait GatewayConnector: Send + Sync { - /// Connect and attach, each stage bounded by `deadline`. - fn open( - &self, + /// Connect to `endpoint`, or the configured endpoint when `None`, and + /// attach, each stage bounded by `deadline`. + fn open<'a>( + &'a self, outbound: tokio::sync::mpsc::Receiver, deadline: tokio::time::Instant, - ) -> std::pin::Pin> + Send + '_>>; + endpoint: Option<&'a str>, + ) -> std::pin::Pin> + Send + 'a>>; } impl GatewayConnector for GatewayLink { - fn open( - &self, + fn open<'a>( + &'a self, outbound: tokio::sync::mpsc::Receiver, deadline: tokio::time::Instant, - ) -> std::pin::Pin> + Send + '_>> { + endpoint: Option<&'a str>, + ) -> std::pin::Pin> + Send + 'a>> { Box::pin(async move { - if !self.endpoint.starts_with("https://") { + let endpoint = endpoint.unwrap_or(&self.endpoint); + if !endpoint.starts_with("https://") { return Err("the supervisor token requires an https gateway endpoint".to_string()); } let token = self .read_token() .map_err(|error| format!("no usable supervisor token: {error}"))?; - let connect = openshell_core::grpc_client::connect_gateway_channel(&self.endpoint); + let connect = openshell_core::grpc_client::connect_gateway_channel(endpoint); let channel = tokio::time::timeout_at(deadline, connect) .await .map_err(|_| "connecting to the gateway timed out".to_string())? @@ -1072,6 +1081,7 @@ async fn attach_with( // attempt bound depends on it, so the attempt after the covered // unanswered ones is still quick. let mut unanswered = 0_u32; + let mut redial = None; tokio::pin!(shutdown); loop { connection_epoch += 1; @@ -1084,10 +1094,15 @@ async fn attach_with( }), }; let attempt_timeout = timings.attempt_timeout(unanswered); - let StreamEnd::Lost { accepted } = tokio::select! { - end = run_stream(&host, connector.as_ref(), &reconciler, hello, &timings, attempt_timeout) => end, + let end = tokio::select! { + end = run_stream(&host, connector.as_ref(), &reconciler, hello, &timings, attempt_timeout, redial.as_deref()) => end, () = &mut shutdown => return reconciler.shut_down().await, }; + let StreamEnd::Lost { + accepted, + redial: next, + } = end; + redial = next; if accepted { failed = 0; unanswered = 0; @@ -1098,6 +1113,7 @@ async fn attach_with( let backoff = timings.backoff(failed, jitter_sample()); tracing::info!( retry_in_ms = u64::try_from(backoff.as_millis()).unwrap_or(u64::MAX), + redial = redial.as_deref().unwrap_or_default(), "gateway assignment stream lost; hosted sandbox instances keep running" ); tokio::select! { @@ -1346,8 +1362,14 @@ async fn run_stream( hello: HostHello, timings: &StreamTimings, attempt_timeout: Duration, + endpoint: Option<&str>, ) -> StreamEnd { - let lost = |accepted| StreamEnd::Lost { accepted }; + let lost = |accepted: bool| StreamEnd::Lost { + accepted, + redial: accepted + .then(|| host.gateway_home.borrow().clone()) + .flatten(), + }; // One deadline covers saying hello, connecting, attaching, and being // accepted, so a slow stage leaves less time for the next; the reconnect // budget counts on it. @@ -1377,7 +1399,7 @@ async fn run_stream( { return lost(false); } - let opened = tokio::time::timeout_at(deadline, connector.open(rx, deadline)) + let opened = tokio::time::timeout_at(deadline, connector.open(rx, deadline, endpoint)) .await .unwrap_or_else(|_| Err("connecting to the gateway timed out".to_string())); let mut inbound = match opened { @@ -1938,6 +1960,8 @@ mod tests { struct Opened { /// When the host asked for it. at: tokio::time::Instant, + /// The replica the host redialed, or `None` for the Service. + endpoint: Option, outbound: tokio::sync::mpsc::Receiver, inbound: tokio::sync::mpsc::Sender>, } @@ -1951,13 +1975,15 @@ mod tests { } impl GatewayConnector for FakeGateway { - fn open( - &self, + fn open<'a>( + &'a self, outbound: tokio::sync::mpsc::Receiver, _deadline: tokio::time::Instant, - ) -> std::pin::Pin> + Send + '_>> + endpoint: Option<&'a str>, + ) -> std::pin::Pin> + Send + 'a>> { let at = tokio::time::Instant::now(); + let endpoint = endpoint.map(str::to_string); self.calls.lock().unwrap().push(at); let attempt = self .plan @@ -1974,6 +2000,7 @@ mod tests { let (inbound, stream) = tokio::sync::mpsc::channel(1); let _ = self.opened.send(Opened { at, + endpoint, outbound, inbound, }); @@ -2254,6 +2281,30 @@ mod tests { home.wait_for(Option::is_none).await.unwrap(); } + /// A host whose accepted stream is lost, as on token expiry, redials the + /// replica that accepted it, so the hosted sessions stay in place; when + /// that replica does not accept it, the host goes back to the Service. + #[tokio::test(start_paused = true)] + async fn a_lost_stream_redials_its_replica_once() { + let (_host, _gateway, mut streams, _root) = + attach_pool_host(vec![Attempt::Now, Attempt::Now, Attempt::Now]); + let first = streams.recv().await.unwrap(); + assert_eq!(first.endpoint, None); + first + .inbound + .send(Ok(accepted_by("https://gw-a:8443"))) + .await + .unwrap(); + drop(first); + + let refused = streams.recv().await.unwrap(); + assert_eq!(refused.endpoint.as_deref(), Some("https://gw-a:8443")); + drop(refused); + + let next = streams.recv().await.unwrap(); + assert_eq!(next.endpoint, None); + } + /// A replica that keeps writing but stops reading fills the host's /// outbound queue; the blocked send gives up within the silence budget /// rather than the ordinary send bound. From e2cae14371df99ec3230904ede525a593f75ca1f Mon Sep 17 00:00:00 2001 From: divesh Date: Mon, 2 Nov 2026 14:20:30 -0800 Subject: [PATCH 09/11] fix(supervisor): follow a host move that lands during session setup An established session marked its host's current gateway as seen before watching it, so a move that landed while the session was connecting was never followed. Compare the current gateway first, then wait for changes. Signed-off-by: divesh --- .../src/supervisor_session.rs | 28 +++++++++++++------ 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/crates/openshell-supervisor-process/src/supervisor_session.rs b/crates/openshell-supervisor-process/src/supervisor_session.rs index 427af34d12..848d7fddf1 100644 --- a/crates/openshell-supervisor-process/src/supervisor_session.rs +++ b/crates/openshell-supervisor-process/src/supervisor_session.rs @@ -609,20 +609,21 @@ async fn wait_for_home( .and_then(|home| (*home).clone()) } -/// The peer endpoint the shared host moves to, once it changes to one other -/// than `current`. Pending forever without a shared host. +/// The peer endpoint of the shared host, once it is one other than `current`, +/// including a move that landed before this is first polled. Pending forever +/// without a shared host. async fn home_moved(home: Option<&mut watch::Receiver>>, current: &str) -> String { let Some(home) = home else { return std::future::pending().await; }; loop { - if home.changed().await.is_err() { - return std::future::pending().await; - } let endpoint = home.borrow_and_update().clone(); if let Some(endpoint) = endpoint.filter(|endpoint| endpoint != current) { return endpoint; } + if home.changed().await.is_err() { + return std::future::pending().await; + } } } @@ -768,9 +769,6 @@ async fn run_single_session( let mut heartbeat_interval = tokio::time::interval(Duration::from_secs(heartbeat_secs)); heartbeat_interval.tick().await; // skip immediate tick let mut home = config.gateway_home.clone(); - if let Some(home) = home.as_mut() { - home.mark_unchanged(); - } loop { tokio::select! { @@ -1743,9 +1741,21 @@ mod ocsf_event_tests { } #[tokio::test] - async fn an_established_session_follows_its_host_to_another_gateway() { + async fn a_host_move_during_session_setup_is_followed() { let (tx, mut home) = watch::channel(Some("https://gw-a:8443".to_string())); + // The session dialed A; the host moved to B before it was accepted, + // and the move was already observed. + tx.send_replace(Some("https://gw-b:8443".to_string())); home.mark_unchanged(); + assert_eq!( + home_moved(Some(&mut home), "https://gw-a:8443").await, + "https://gw-b:8443" + ); + } + + #[tokio::test] + async fn an_established_session_follows_its_host_to_another_gateway() { + let (tx, mut home) = watch::channel(Some("https://gw-a:8443".to_string())); let moved = tokio::spawn(async move { home_moved(Some(&mut home), "https://gw-a:8443").await }); From f7593dbd0a9dda47dbab6f5890bf74d34a174368 Mon Sep 17 00:00:00 2001 From: divesh Date: Tue, 3 Nov 2026 14:28:27 -0800 Subject: [PATCH 10/11] fix(supervisor): dial the configured gateway at once when the host's replica has no peer endpoint Validate the advertised peer endpoint scheme with the shared supervisor, and document that pool sessions follow their supervisor's gateway. Signed-off-by: divesh --- crates/openshell-server/src/lib.rs | 4 ++++ .../src/supervisor_session.rs | 23 +++++++++++++++++-- crates/openshell-supervisor/src/lib.rs | 3 ++- .../openshell-supervisor/src/shared_host.rs | 10 ++++---- docs/kubernetes/high-availability.mdx | 4 ++++ proto/openshell.proto | 9 ++++---- sdk/go/proto/openshellv1/openshell.pb.go | 9 ++++---- 7 files changed, 46 insertions(+), 16 deletions(-) diff --git a/crates/openshell-server/src/lib.rs b/crates/openshell-server/src/lib.rs index 7522d5cbe7..6a8fdfd500 100644 --- a/crates/openshell-server/src/lib.rs +++ b/crates/openshell-server/src/lib.rs @@ -836,6 +836,10 @@ pub(crate) async fn run_server( "shared supervisor assignment stream enabled; minting supervisor tokens" ); info!("shared supervisor pool enabled; reconciling pool membership and ownership"); + // Sandbox sessions dial this endpoint with their tokens. + if let Some(peer_endpoint) = state.peer_endpoint.as_deref() { + validate_peer_endpoint_scheme(&state.config, peer_endpoint)?; + } controller.set_peer_endpoint(state.peer_endpoint.clone().unwrap_or_default()); supervisor_pool::SupervisorPool::new( state.store.clone(), diff --git a/crates/openshell-supervisor-process/src/supervisor_session.rs b/crates/openshell-supervisor-process/src/supervisor_session.rs index 848d7fddf1..6e86ebd9c1 100644 --- a/crates/openshell-supervisor-process/src/supervisor_session.rs +++ b/crates/openshell-supervisor-process/src/supervisor_session.rs @@ -594,7 +594,8 @@ impl ReconnectBackoff { } /// The peer endpoint of the replica serving the shared host, once it is one -/// other than `avoid`, waiting up to [`GATEWAY_HOME_WAIT`]. +/// other than `avoid`, waiting up to [`GATEWAY_HOME_WAIT`] while no replica +/// serves the host. `None` at once when the serving replica has no endpoint. async fn wait_for_home( home: &mut watch::Receiver>, avoid: Option<&str>, @@ -607,6 +608,7 @@ async fn wait_for_home( .ok()? .ok() .and_then(|home| (*home).clone()) + .filter(|home| !home.is_empty()) } /// The peer endpoint of the shared host, once it is one other than `current`, @@ -618,7 +620,9 @@ async fn home_moved(home: Option<&mut watch::Receiver>>, current: }; loop { let endpoint = home.borrow_and_update().clone(); - if let Some(endpoint) = endpoint.filter(|endpoint| endpoint != current) { + if let Some(endpoint) = + endpoint.filter(|endpoint| !endpoint.is_empty() && endpoint != current) + { return endpoint; } if home.changed().await.is_err() { @@ -1740,6 +1744,21 @@ mod ocsf_event_tests { assert_eq!(moved.await.unwrap().as_deref(), Some("https://gw-b:8443")); } + #[tokio::test] + async fn a_host_gateway_without_a_peer_endpoint_is_not_waited_for() { + let (_tx, mut home) = watch::channel(Some(String::new())); + let started = tokio::time::Instant::now(); + assert_eq!(wait_for_home(&mut home, None).await, None); + assert!(started.elapsed() < GATEWAY_HOME_WAIT); + + let moved = tokio::time::timeout( + Duration::from_millis(50), + home_moved(Some(&mut home), "https://gw-a:8443"), + ) + .await; + assert!(moved.is_err(), "an endpointless gateway is not a move"); + } + #[tokio::test] async fn a_host_move_during_session_setup_is_followed() { let (tx, mut home) = watch::channel(Some("https://gw-a:8443".to_string())); diff --git a/crates/openshell-supervisor/src/lib.rs b/crates/openshell-supervisor/src/lib.rs index d2456a8e5e..9f6edd34a8 100644 --- a/crates/openshell-supervisor/src/lib.rs +++ b/crates/openshell-supervisor/src/lib.rs @@ -897,7 +897,8 @@ pub struct SandboxInstanceConfig { pub container_image: String, pub openshell_endpoint: Option, /// Peer endpoint of the gateway replica serving the shared host, which - /// the session dials instead of `openshell_endpoint` when set. + /// the session dials instead of `openshell_endpoint` when set and + /// nonempty. pub gateway_home: Option>>, pub policy_rules: Option, pub policy_data: Option, diff --git a/crates/openshell-supervisor/src/shared_host.rs b/crates/openshell-supervisor/src/shared_host.rs index d912965ca3..6467343d1a 100644 --- a/crates/openshell-supervisor/src/shared_host.rs +++ b/crates/openshell-supervisor/src/shared_host.rs @@ -210,7 +210,8 @@ pub struct SharedHost { /// A gateway currently accepts this host's stream. gateway_connected: watch::Sender, /// Peer endpoint of the gateway replica accepting this host's stream; - /// hosted instances dial their sessions there. + /// hosted instances dial their sessions there. Empty when that replica + /// has none, and `None` while no replica accepts the stream. gateway_home: watch::Sender>, /// Recovery deadline (seconds) the latest snapshot offers the runtimes /// of instances this host starts or takes over; `None` keeps the runtime @@ -1368,7 +1369,8 @@ async fn run_stream( accepted, redial: accepted .then(|| host.gateway_home.borrow().clone()) - .flatten(), + .flatten() + .filter(|endpoint| !endpoint.is_empty()), }; // One deadline covers saying hello, connecting, attaching, and being // accepted, so a slow stage leaves less time for the next; the reconnect @@ -1444,9 +1446,7 @@ async fn run_stream( Some(host_gateway_message::Payload::Accepted(accepted_message)) => { accepted = true; host.gateway_connected.send_replace(true); - host.gateway_home.send_replace( - Some(accepted_message.peer_endpoint.clone()).filter(|endpoint| !endpoint.is_empty()), - ); + host.gateway_home.send_replace(Some(accepted_message.peer_endpoint.clone())); tracing::info!(session = %accepted_message.session_id, "gateway accepted the supervisor host"); // A pool gateway detects failed members by silence, // so it asks for more frequent heartbeats. diff --git a/docs/kubernetes/high-availability.mdx b/docs/kubernetes/high-availability.mdx index 73f40cc6f4..c21150f762 100644 --- a/docs/kubernetes/high-availability.mdx +++ b/docs/kubernetes/high-availability.mdx @@ -191,6 +191,10 @@ stays `Ready` through a graceful handoff. Supervisors reconnect with exponential backoff that restarts after each accepted session, so repeated gateway restarts do not lengthen later reconnects. +A sandbox on a shared supervisor keeps its session on the gateway replica that +serves its supervisor, and reconnects there when the supervisor moves to +another replica. + Established streams do not move between gateway replicas. An interrupted exec, forwarding, or service stream must reconnect. File sync retries with a new SSH session instead of resuming the interrupted byte stream. Rolling updates can diff --git a/proto/openshell.proto b/proto/openshell.proto index e498a84af9..721f815aac 100644 --- a/proto/openshell.proto +++ b/proto/openshell.proto @@ -3229,10 +3229,11 @@ message SupervisorHello { uint64 connection_epoch = 3; // The supervisor can report credential, policy, and launch-environment installation. bool supports_provider_readiness = 4; - // Set when this connection is the result of following a SessionRedirect. The - // receiving gateway serves a redirected hello locally instead of redirecting - // again, so a disagreement about membership cannot bounce the supervisor - // between replicas. A gateway that is shutting down may still redirect it. + // Set when this connection follows a SessionRedirect, or dials the replica + // serving the sandbox's shared supervisor. The receiving gateway serves a + // redirected hello locally instead of redirecting again, so a disagreement + // about membership cannot bounce the supervisor between replicas. A gateway + // that is shutting down may still redirect it. bool redirected = 5; // The supervisor understands SessionRedirect. Gateways never redirect a // supervisor that does not set this. diff --git a/sdk/go/proto/openshellv1/openshell.pb.go b/sdk/go/proto/openshellv1/openshell.pb.go index 1ae1800784..bb665f05ad 100644 --- a/sdk/go/proto/openshellv1/openshell.pb.go +++ b/sdk/go/proto/openshellv1/openshell.pb.go @@ -14345,10 +14345,11 @@ type SupervisorHello struct { ConnectionEpoch uint64 `protobuf:"varint,3,opt,name=connection_epoch,json=connectionEpoch,proto3" json:"connection_epoch,omitempty"` // The supervisor can report credential, policy, and launch-environment installation. SupportsProviderReadiness bool `protobuf:"varint,4,opt,name=supports_provider_readiness,json=supportsProviderReadiness,proto3" json:"supports_provider_readiness,omitempty"` - // Set when this connection is the result of following a SessionRedirect. The - // receiving gateway serves a redirected hello locally instead of redirecting - // again, so a disagreement about membership cannot bounce the supervisor - // between replicas. A gateway that is shutting down may still redirect it. + // Set when this connection follows a SessionRedirect, or dials the replica + // serving the sandbox's shared supervisor. The receiving gateway serves a + // redirected hello locally instead of redirecting again, so a disagreement + // about membership cannot bounce the supervisor between replicas. A gateway + // that is shutting down may still redirect it. Redirected bool `protobuf:"varint,5,opt,name=redirected,proto3" json:"redirected,omitempty"` // The supervisor understands SessionRedirect. Gateways never redirect a // supervisor that does not set this. From 3e2dfab579465726d928f2a1d318c953efd9eca4 Mon Sep 17 00:00:00 2001 From: divesh Date: Tue, 3 Nov 2026 14:28:27 -0800 Subject: [PATCH 11/11] test(e2e): harden the supervisor tree check's port-forward and time limit Signed-off-by: divesh --- .config/nextest.toml | 5 ++ .../tests/kubernetes_shared_supervisor_ha.rs | 79 ++++++++++++------- 2 files changed, 54 insertions(+), 30 deletions(-) diff --git a/.config/nextest.toml b/.config/nextest.toml index b159e1db2c..fef1880a35 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -46,6 +46,11 @@ test-group = "kubernetes-shared-supervisor-pool" filter = "test(/^scaling_a_supervisor_pool_up_and_down_keeps_its_sandboxes_running$/)" slow-timeout = { period = "60s", terminate-after = 20 } +# Three bounded rollouts (up to 5 minutes each) and four settle and exec rounds. +[[profile.e2e-kubernetes.overrides]] +filter = "test(/^pool_sessions_follow_their_supervisors_gateway$/)" +slow-timeout = { period = "60s", terminate-after = 20 } + # Relative to the profile store dir (`e2e/rust/target/nextest/e2e-kubernetes/`). [profile.e2e-kubernetes.junit] path = "../../../../../results/e2e-kubernetes.xml" diff --git a/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs b/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs index 95265e8e14..36b16cf698 100644 --- a/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs +++ b/e2e/rust/tests/kubernetes_shared_supervisor_ha.rs @@ -477,6 +477,10 @@ async fn assert_tree(step: &str, routes: &[String], gone: &[String]) -> Vec Vec, } impl GatewayForward { @@ -505,7 +510,7 @@ impl GatewayForward { .as_u64() .and_then(|port| u16::try_from(port).ok()) .expect("gateway_port"); - Self { port } + Self { port, child: None } } async fn reachable(&self) -> bool { @@ -514,34 +519,48 @@ impl GatewayForward { .is_ok() } - async fn ensure(&self) { - if self.reachable().await { - return; - } - kubectl_command() - .args([ - "-n", - &release_namespace(), - "port-forward", - &format!("svc/{}", release()), - &format!("{}:8080", self.port), - ]) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .kill_on_drop(false) - .spawn() - .expect("spawn kubectl port-forward"); - let opened = tokio::time::timeout(Duration::from_secs(30), async { - while !self.reachable().await { - tokio::time::sleep(Duration::from_millis(500)).await; + /// Returns once the port answers through a live port-forward: the + /// harness's, or one this started that has not exited, for example after + /// losing the port to a dying predecessor. + async fn ensure(&mut self) { + let deadline = tokio::time::Instant::now() + Duration::from_mins(1); + loop { + let ours_alive = self + .child + .as_mut() + .map(|child| matches!(child.try_wait(), Ok(None))); + match (ours_alive, self.reachable().await) { + (None | Some(true), true) => return, + (Some(true), false) => {} + _ => { + self.child = Some( + kubectl_command() + .args([ + "-n", + &release_namespace(), + "port-forward", + &format!("svc/{}", release()), + &format!("{}:8080", self.port), + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .kill_on_drop(false) + .spawn() + .expect("spawn kubectl port-forward"), + ); + } } - }) - .await; - assert!(opened.is_ok(), "port-forward on {} never opened", self.port); + assert!( + tokio::time::Instant::now() < deadline, + "no live port-forward on {}", + self.port + ); + tokio::time::sleep(Duration::from_millis(500)).await; + } } } -async fn assert_exec(step: &str, forward: &GatewayForward, sandboxes: &[SandboxGuard]) { +async fn assert_exec(step: &str, forward: &mut GatewayForward, sandboxes: &[SandboxGuard]) { for sandbox in sandboxes { let mut last = String::new(); let reached = tokio::time::timeout(SETTLE_DEADLINE, async { @@ -638,7 +657,7 @@ async fn pool_sessions_follow_their_supervisors_gateway() { ); wait_until_attached(&supervisor_pods().await).await; - let forward = GatewayForward::new(); + let mut forward = GatewayForward::new(); forward.ensure().await; let mut sandboxes = Vec::new(); for _ in 0..SANDBOXES { @@ -653,25 +672,25 @@ async fn pool_sessions_follow_their_supervisors_gateway() { routes.push(generation(&sandbox.name).await.route); } let branches = assert_tree("start", &routes, &[]).await; - assert_exec("start", &forward, &sandboxes).await; + assert_exec("start", &mut forward, &sandboxes).await; let crashed = vec![branches[0].supervisor_gateway.clone()]; force_delete(&crashed[0]).await; rollout("status", &gateway).await; assert_tree("gateway crash", &routes, &crashed).await; - assert_exec("gateway crash", &forward, &sandboxes).await; + assert_exec("gateway crash", &mut forward, &sandboxes).await; let restarted = gateway_pods().await; kubectl(&["-n", &release_namespace(), "rollout", "restart", &gateway]).await; rollout("status", &gateway).await; let branches = assert_tree("gateway rolling restart", &routes, &restarted).await; - assert_exec("gateway rolling restart", &forward, &sandboxes).await; + assert_exec("gateway rolling restart", &mut forward, &sandboxes).await; let crashed = vec![branches[0].member.clone()]; force_delete(&crashed[0]).await; rollout("status", &supervisor).await; assert_tree("supervisor crash", &routes, &crashed).await; - assert_exec("supervisor crash", &forward, &sandboxes).await; + assert_exec("supervisor crash", &mut forward, &sandboxes).await; forward.ensure().await; drop(sandboxes);