Skip to content

Commit 9b5bcdd

Browse files
authored
fix(ha): keep sandboxes ready across gateway pod rolls (#4321)
* fix(ha): keep sandboxes ready across gateway pod rolls A gateway pod roll left sandboxes not ready long enough for clients to fail with "sandbox is not ready", which made the HA e2e test sandbox_file_sync_survives_gateway_pod_rolls flaky. Two causes: - The supervisor never reset its reconnect backoff after an accepted session, so each gateway restart over a sandbox's lifetime doubled the delay until every reconnect waited the full 30s maximum. Reset it after an accepted session. - A stopping gateway replica redirected its supervisors and then immediately demoted their sandboxes to Provisioning, before the redirected supervisor published its replacement session. Wait up to 5s for the replacement before demoting; the existing disconnect path keeps the sandbox Ready once a replacement exists. Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com> * fix(ha): wait for handoff only when a redirect went out Address review feedback on the shutdown handoff: - Wait for a replacement session only when the stopping replica queued a SessionRedirect. Single-replica gateways, supervisors without redirect support, and full outbound channels no longer pay the grace period. - Close the session stream in both directions before cleanup so a supervisor without a redirect sees EOF and reconnects at once. - Bound the replacement wait with timeout_at and the shared session-wait backoff, so the grace is a hard ceiling. - Name the supervisor session shutdown timeout and assert at compile time that it leaves room after the grace period. - Demote the sandbox when session setup fails after the owner publish, so a failed takeover cannot leave it Ready with no session. - Move the supervisor reconnect delays into ReconnectBackoff so both loop branches reset the same way, and test reconnect sequences. Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com> * fix(ha): refresh gateway membership before shutdown redirects A stopping replica redirected its supervisors using its cached ring, which the membership worker refreshes only every 10s. During a rolling update the replacement for the previous pod is often only seconds old, so the cached ring missed it and still named the departed pod. The redirect then went nowhere or to a dead pod, whose 10s connect timeout outlasted the handoff grace, and the sandbox was demoted mid-roll. Re-read live membership, bounded to 2s, just before session shutdown computes redirects. Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com> * fix(ha): hold ownership through the shutdown handoff The stopping replica released its owner record before waiting for the replacement session. A reconciliation sweep on another replica in that window saw no owner, and with the Kubernetes driver leaving readiness to the gateway it demoted the sandbox to Provisioning anyway. Keep the record until a different session supersedes it or the grace expires, then release it only if it is still ours. Also from review: - Keep polling the owner record through store errors until the deadline. - Skip the handoff wait for sandboxes that are not Ready. - Re-read membership in the membership worker's shutdown branch and await the worker before session shutdown, so the worker stays the only writer of the ring and peer map. - Share the remove, release, and demote steps of failed session setup. - Test redirect queuing directly, including unsupported supervisors and a full outbound channel. Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com> --------- Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
1 parent 277f922 commit 9b5bcdd

6 files changed

Lines changed: 657 additions & 103 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4760,7 +4760,7 @@ impl ComputeRuntime {
47604760
Ok(())
47614761
}
47624762

4763-
async fn supervisor_session_ready(&self, sandbox_id: &str) -> Result<bool, String> {
4763+
pub(crate) async fn supervisor_session_ready(&self, sandbox_id: &str) -> Result<bool, String> {
47644764
if self.supervisor_sessions.has_session(sandbox_id) {
47654765
return Ok(true);
47664766
}

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

Lines changed: 97 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,13 @@ pub const MEMBER_REFRESH_INTERVAL: Duration = Duration::from_secs(10);
2828
/// Upper bound on replicas read back when building the ring.
2929
const MEMBER_LIST_LIMIT: u32 = 1024;
3030

31+
/// How long the worker waits on each store call while leaving at shutdown.
32+
const SHUTDOWN_STORE_TIMEOUT: Duration = Duration::from_secs(2);
33+
34+
/// How long gateway shutdown waits for the membership worker to finish
35+
/// leaving: one membership read and one deregister.
36+
pub const SHUTDOWN_MEMBERSHIP_TIMEOUT: Duration = Duration::from_secs(5);
37+
3138
fn member_object_id(replica_id: &str) -> String {
3239
format!("gateway-member:{replica_id}")
3340
}
@@ -152,25 +159,28 @@ impl GatewayMemberIndex {
152159
/// One task does both so the ring is never newer than our own liveness claim.
153160
/// The ring starts empty, which means "serve locally" — so a store outage
154161
/// degrades to today's behaviour rather than breaking placement.
162+
///
163+
/// Returns the worker task, which finishes after this replica leaves the ring
164+
/// at shutdown. Shutdown awaits it before redirecting supervisor sessions.
155165
pub fn spawn_membership_worker(
156166
state: Arc<crate::ServerState>,
157167
interval: Duration,
158168
mut shutdown: tokio::sync::watch::Receiver<bool>,
159-
) {
169+
) -> Option<tokio::task::JoinHandle<()>> {
160170
if state.store.is_single_replica() {
161-
return;
171+
return None;
162172
}
163173
let Some(peer_endpoint) = state.peer_endpoint.clone() else {
164174
// Without a peer endpoint other replicas cannot dial us, so we must
165175
// not advertise ourselves as a placement target.
166176
tracing::debug!("gateway membership: no peer endpoint, not joining the ring");
167-
return;
177+
return None;
168178
};
169179

170180
let index = GatewayMemberIndex::new(state.store.clone(), MEMBER_TTL);
171181
let replica_id = state.replica_id.clone();
172182

173-
tokio::spawn(async move {
183+
Some(tokio::spawn(async move {
174184
let mut ticker = tokio::time::interval(interval);
175185
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
176186
let mut last_refreshed = tokio::time::Instant::now();
@@ -180,14 +190,7 @@ pub fn spawn_membership_worker(
180190
_ = ticker.tick() => {}
181191
_ = shutdown.changed() => {
182192
if *shutdown.borrow() {
183-
// Sessions closing during shutdown redirect to the
184-
// replica that owns them once this one is gone.
185-
if let Ok(mut slot) = state.gateway_ring.write() {
186-
*slot = slot.without(&replica_id);
187-
}
188-
if let Err(error) = index.deregister(&replica_id).await {
189-
tracing::warn!(%error, "gateway membership: deregister failed");
190-
}
193+
leave_ring(&state, &index, &replica_id).await;
191194
return;
192195
}
193196
continue;
@@ -220,7 +223,50 @@ pub fn spawn_membership_worker(
220223
}
221224
}
222225
}
223-
});
226+
}))
227+
}
228+
229+
/// Point the ring at the replicas that remain once this one is gone, then
230+
/// remove this replica's record.
231+
///
232+
/// Sessions closing during shutdown redirect to the replica that owns them
233+
/// next. The periodic refresh can be up to one interval old, and during a
234+
/// rolling update that ring can miss a replacement that just started and still
235+
/// name a replica that already left. Redirects built from it send supervisors
236+
/// nowhere or to a dead pod, so re-read membership first. If the read fails,
237+
/// keep the cached ring without this replica.
238+
async fn leave_ring(state: &crate::ServerState, index: &GatewayMemberIndex, replica_id: &str) {
239+
match tokio::time::timeout(SHUTDOWN_STORE_TIMEOUT, index.ring()).await {
240+
Ok(Ok((ring, members))) => {
241+
let peers: HashMap<String, String> = members
242+
.into_iter()
243+
.map(|member| (member.replica_id, member.peer_endpoint))
244+
.collect();
245+
if let Ok(mut slot) = state.gateway_ring.write() {
246+
*slot = ring.without(replica_id);
247+
}
248+
if let Ok(mut slot) = state.gateway_peers.write() {
249+
*slot = peers;
250+
}
251+
}
252+
Ok(Err(error)) => {
253+
tracing::warn!(%error, "gateway membership: shutdown ring refresh failed");
254+
if let Ok(mut slot) = state.gateway_ring.write() {
255+
*slot = slot.without(replica_id);
256+
}
257+
}
258+
Err(_) => {
259+
tracing::warn!("gateway membership: shutdown ring refresh timed out");
260+
if let Ok(mut slot) = state.gateway_ring.write() {
261+
*slot = slot.without(replica_id);
262+
}
263+
}
264+
}
265+
match tokio::time::timeout(SHUTDOWN_STORE_TIMEOUT, index.deregister(replica_id)).await {
266+
Ok(Ok(())) => {}
267+
Ok(Err(error)) => tracing::warn!(%error, "gateway membership: deregister failed"),
268+
Err(_) => tracing::warn!("gateway membership: deregister timed out"),
269+
}
224270
}
225271

226272
fn clear_ring_if_stale(state: &crate::ServerState, last_refreshed: tokio::time::Instant) {
@@ -259,6 +305,44 @@ mod tests {
259305
);
260306
}
261307

308+
#[tokio::test]
309+
async fn leaving_points_ring_at_new_replicas_and_not_departed_ones() {
310+
let state = crate::grpc::test_support::test_server_state().await;
311+
let index = GatewayMemberIndex::new(state.store.clone(), MEMBER_TTL);
312+
index.register("gw-old", "https://gw-old").await.unwrap();
313+
index
314+
.register("gw-stopping", "https://gw-stopping")
315+
.await
316+
.unwrap();
317+
*state.gateway_ring.write().unwrap() = GatewayRing::new(["gw-old", "gw-stopping"]);
318+
// A rolling update replaces gw-old after the stopping replica last
319+
// refreshed its cached ring.
320+
index.deregister("gw-old").await.unwrap();
321+
index.register("gw-new", "https://gw-new").await.unwrap();
322+
323+
leave_ring(&state, &index, "gw-stopping").await;
324+
325+
let ring = state.gateway_ring.read().unwrap().clone();
326+
assert_eq!(ring.owner_for("sb-1"), Some("gw-new"));
327+
assert_eq!(
328+
state
329+
.gateway_peers
330+
.read()
331+
.unwrap()
332+
.get("gw-new")
333+
.map(String::as_str),
334+
Some("https://gw-new")
335+
);
336+
let remaining: Vec<String> = index
337+
.live_members()
338+
.await
339+
.unwrap()
340+
.into_iter()
341+
.map(|member| member.replica_id)
342+
.collect();
343+
assert_eq!(remaining, vec!["gw-new"]);
344+
}
345+
262346
#[tokio::test]
263347
async fn register_is_idempotent_and_updates_the_endpoint() {
264348
let index = test_index(MEMBER_TTL).await;

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -967,7 +967,7 @@ pub(crate) async fn run_server(
967967
shutdown_rx.clone(),
968968
);
969969
}
970-
gateway_members::spawn_membership_worker(
970+
let membership_worker = gateway_members::spawn_membership_worker(
971971
state.clone(),
972972
gateway_members::MEMBER_REFRESH_INTERVAL,
973973
shutdown_rx.clone(),
@@ -987,11 +987,20 @@ pub(crate) async fn run_server(
987987
}
988988

989989
let compute_cleanup = state.compute.cleanup_on_shutdown().await;
990+
// Closing sessions redirect their supervisors through the ring. Let the
991+
// membership worker finish pointing it at the remaining replicas first.
992+
if let Some(worker) = membership_worker
993+
&& tokio::time::timeout(gateway_members::SHUTDOWN_MEMBERSHIP_TIMEOUT, worker)
994+
.await
995+
.is_err()
996+
{
997+
warn!("Gateway membership worker did not finish leaving the ring before session shutdown");
998+
}
990999
// A stopped supervisor may still have a detached task deleting its owner
9911000
// record. Drain it even when compute cleanup failed before exiting Tokio.
9921001
let session_cleanup = state
9931002
.supervisor_sessions
994-
.shutdown(Duration::from_secs(10))
1003+
.shutdown(supervisor_session::SUPERVISOR_SESSION_SHUTDOWN_TIMEOUT)
9951004
.await;
9961005
if let Err(err) = &session_cleanup {
9971006
warn!(error = %err, "Gateway supervisor session cleanup incomplete");

0 commit comments

Comments
 (0)