From e104a7eb2b9674b7ecf81b7f670c1d1b33990879 Mon Sep 17 00:00:00 2001 From: Nikhil Unni Date: Mon, 9 Nov 2026 08:42:51 -0800 Subject: [PATCH 1/2] upgrade-drain: retry a coordinator call that answers Unavailable The host-fleet pre-upgrade drain failed on its first call: list_hosts got `Unavailable` (connection refused). The coordinator rollout had finished seconds before, so the Service still routed to a stopping pod. The drain loop retries its actions, but a failed list_hosts or list_sessions at the top of a pass ended the hook. Every call of the drain's Coordinator client now goes through one retry: on `Unavailable` it tries again every 2 s, up to 30 attempts (about a minute). Any other status fails at once. The retry counts attempts, so it reads no clock (ADR 0098). Both hooks (`upgrade-drain` in the coordinator and host-operator images, and `upgrade-roll`) use this client. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_017XkqDqWbQdPXdy94AdFV5e --- crates/engram-upgrade-drain/src/client.rs | 199 +++++++++++++++++----- 1 file changed, 152 insertions(+), 47 deletions(-) diff --git a/crates/engram-upgrade-drain/src/client.rs b/crates/engram-upgrade-drain/src/client.rs index e4469bee6..2eb4d875f 100644 --- a/crates/engram-upgrade-drain/src/client.rs +++ b/crates/engram-upgrade-drain/src/client.rs @@ -2,6 +2,7 @@ //! coordinator that still serves, so it uses only RPCs the previous release //! has. +use std::future::Future; use std::time::Duration; use engram_core::HostId; @@ -55,6 +56,49 @@ pub fn channel(endpoint: &str, call_timeout: Duration) -> Channel { .connect_lazy() } +/// How often one call tries a coordinator that answers `Unavailable`, and the +/// wait between tries: about a minute in all. The hook can start while the +/// coordinator rolls, when the Service still routes to a stopping pod; the +/// Helm hook timeout bounds the whole drain. +const UNAVAILABLE_ATTEMPTS: u32 = 30; +const UNAVAILABLE_BACKOFF: Duration = Duration::from_secs(2); + +/// Run one coordinator call, again on `Unavailable` (connection refused or +/// reset, no ready endpoint) up to `attempts` times. Any other status fails +/// at once. The drain re-reads the fleet each pass, so a call that reached +/// the coordinator before the error is safe to send again. +async fn retry_unavailable( + op: &'static str, + attempts: u32, + backoff: Duration, + mut call: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future, tonic::Status>>, +{ + let mut attempt = 1; + loop { + match call().await { + Ok(resp) => return Ok(resp.into_inner()), + Err(status) if status.code() == tonic::Code::Unavailable && attempt < attempts => { + tracing::warn!(op, attempt, error = %status, "coordinator unavailable; retrying"); + attempt += 1; + tokio::time::sleep(backoff).await; + } + Err(status) => return Err(DrainError::rpc(op, status)), + } + } +} + +async fn call(op: &'static str, call: F) -> Result +where + F: FnMut() -> Fut, + Fut: Future, tonic::Status>>, +{ + retry_unavailable(op, UNAVAILABLE_ATTEMPTS, UNAVAILABLE_BACKOFF, call).await +} + pub fn parse_host_id(id: &str) -> Result { id.parse() .map_err(|_| DrainError::Invalid(format!("coordinator returned malformed host id {id:?}"))) @@ -103,14 +147,12 @@ impl Coordinator { } pub async fn list_hosts(&self) -> Result, DrainError> { - let resp = self - .fleet - .clone() - .list_hosts(app::ListHostsRequest {}) - .await - .map_err(|status| DrainError::rpc("list_hosts", status))?; - resp.into_inner() - .hosts + let resp = call("list_hosts", || { + let mut fleet = self.fleet.clone(); + async move { fleet.list_hosts(app::ListHostsRequest {}).await } + }) + .await?; + resp.hosts .into_iter() .map(|v| { Ok(UpgradeHost { @@ -125,26 +167,25 @@ impl Coordinator { } pub async fn cordon(&self, host: HostId, owner: &str) -> Result<(), DrainError> { - self.fleet - .clone() - .cordon_host(app::CordonHostRequest { + call("cordon", || { + let mut fleet = self.fleet.clone(); + let req = app::CordonHostRequest { owner: owner.into(), host_id: host.to_string(), - }) - .await - .map_err(|status| DrainError::rpc("cordon", status))?; + }; + async move { fleet.cordon_host(req).await } + }) + .await?; Ok(()) } pub async fn list_sessions(&self) -> Result, DrainError> { - let resp = self - .sessions - .clone() - .list_sessions(app::ListSessionsRequest {}) - .await - .map_err(|status| DrainError::rpc("list_sessions", status))?; - resp.into_inner() - .sessions + let resp = call("list_sessions", || { + let mut sessions = self.sessions.clone(); + async move { sessions.list_sessions(app::ListSessionsRequest {}).await } + }) + .await?; + resp.sessions .into_iter() .filter_map(|item| item.session) .map(|s| { @@ -159,49 +200,113 @@ impl Coordinator { /// Whether the session has a checkpoint it can resume from. pub async fn has_checkpoint(&self, id: &str) -> Result { - let resp = self - .sessions - .clone() - .list_checkpoints(app::ListCheckpointsRequest { + let resp = call("list_checkpoints", || { + let mut sessions = self.sessions.clone(); + let req = app::ListCheckpointsRequest { session_id: id.into(), - }) - .await - .map_err(|status| DrainError::rpc("list_checkpoints", status))?; - Ok(resp.into_inner().checkpoints.iter().any(|c| c.recoverable)) + }; + async move { sessions.list_checkpoints(req).await } + }) + .await?; + Ok(resp.checkpoints.iter().any(|c| c.recoverable)) } pub async fn snapshot(&self, id: &str) -> Result<(), DrainError> { - self.sessions - .clone() - .snapshot(app::SnapshotRequest { + call("snapshot", || { + let mut sessions = self.sessions.clone(); + let req = app::SnapshotRequest { session_id: id.into(), - }) - .await - .map_err(|status| DrainError::rpc("snapshot", status))?; + }; + async move { sessions.snapshot(req).await } + }) + .await?; Ok(()) } pub async fn resume(&self, id: &str) -> Result<(), DrainError> { - self.sessions - .clone() - .resume(app::ResumeRequest { + call("resume", || { + let mut sessions = self.sessions.clone(); + let req = app::ResumeRequest { session_id: id.into(), - }) - .await - .map_err(|status| DrainError::rpc("resume", status))?; + }; + async move { sessions.resume(req).await } + }) + .await?; Ok(()) } /// The eviction that does not park: the session ends `idle` behind a /// fresh checkpoint, its VM destroyed. pub async fn evict(&self, id: &str) -> Result<(), DrainError> { - self.sessions - .clone() - .evict_local(app::EvictLocalRequest { + call("evict_local", || { + let mut sessions = self.sessions.clone(); + let req = app::EvictLocalRequest { session_id: id.into(), + }; + async move { sessions.evict_local(req).await } + }) + .await?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicU32, Ordering}; + + use super::*; + + /// A call that answers each status of `script` in turn, then `Ok(7)`. + fn scripted( + script: Vec, + calls: &AtomicU32, + ) -> impl FnMut() -> std::future::Ready, tonic::Status>> + '_ { + move || { + let n = calls.fetch_add(1, Ordering::SeqCst) as usize; + std::future::ready(match script.get(n) { + Some(code) => Err(tonic::Status::new(*code, "scripted")), + None => Ok(tonic::Response::new(7)), }) + } + } + + #[tokio::test] + async fn an_unavailable_coordinator_is_tried_again_until_it_answers() { + let calls = AtomicU32::new(0); + let script = vec![tonic::Code::Unavailable, tonic::Code::Unavailable]; + let got = retry_unavailable("list_hosts", 5, Duration::ZERO, scripted(script, &calls)) .await - .map_err(|status| DrainError::rpc("evict_local", status))?; - Ok(()) + .unwrap(); + assert_eq!((got, calls.load(Ordering::SeqCst)), (7, 3)); + } + + #[tokio::test] + async fn any_other_status_fails_at_once() { + let calls = AtomicU32::new(0); + let script = vec![tonic::Code::PermissionDenied]; + let err = retry_unavailable("list_hosts", 5, Duration::ZERO, scripted(script, &calls)) + .await + .unwrap_err(); + assert!(matches!( + err, + DrainError::Rpc { + op: "list_hosts", + .. + } + )); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn the_retries_stop_after_the_last_attempt() { + let calls = AtomicU32::new(0); + let script = vec![tonic::Code::Unavailable; 10]; + let err = retry_unavailable("list_hosts", 3, Duration::ZERO, scripted(script, &calls)) + .await + .unwrap_err(); + assert!( + matches!(&err, DrainError::Rpc { status, .. } if status.code() == tonic::Code::Unavailable) + ); + assert_eq!(calls.load(Ordering::SeqCst), 3); } } From 1790e48d601d198b2e5ef1f407a19297c301e1f2 Mon Sep 17 00:00:00 2001 From: Nikhil Unni Date: Tue, 10 Nov 2026 06:34:23 -0800 Subject: [PATCH 2/2] upgrade-drain: only the fleet reads wait for the coordinator The coordinator answers `Unavailable` for two different things: it is not reachable, and its own call to a host failed. The retry treated both the same, so an evict, snapshot or resume of a session on a dead host waited a minute per call inside its pass, and held the other actions of that pass, where it failed in seconds before. The drain already sends a failed action again on its next pass. `list_hosts` and `list_sessions` keep the retry: their failure ends the hook. Cordon, the checkpoint read, snapshot, resume and evict get one try. Co-Authored-By: Claude Fable 5.2 Claude-Session: https://claude.ai/code/session_01HNhpmz4QPYaVCE7tFna54D --- crates/engram-upgrade-drain/src/client.rs | 54 +++++++++++++++++------ 1 file changed, 40 insertions(+), 14 deletions(-) diff --git a/crates/engram-upgrade-drain/src/client.rs b/crates/engram-upgrade-drain/src/client.rs index 2eb4d875f..ef0c1b57f 100644 --- a/crates/engram-upgrade-drain/src/client.rs +++ b/crates/engram-upgrade-drain/src/client.rs @@ -56,17 +56,16 @@ pub fn channel(endpoint: &str, call_timeout: Duration) -> Channel { .connect_lazy() } -/// How often one call tries a coordinator that answers `Unavailable`, and the -/// wait between tries: about a minute in all. The hook can start while the -/// coordinator rolls, when the Service still routes to a stopping pod; the -/// Helm hook timeout bounds the whole drain. +/// How often a fleet read tries a coordinator that answers `Unavailable`, +/// and the wait between tries: about a minute in all. The hook can start +/// while the coordinator rolls, when the Service still routes to a stopping +/// pod; the Helm hook timeout bounds the whole drain. const UNAVAILABLE_ATTEMPTS: u32 = 30; const UNAVAILABLE_BACKOFF: Duration = Duration::from_secs(2); /// Run one coordinator call, again on `Unavailable` (connection refused or /// reset, no ready endpoint) up to `attempts` times. Any other status fails -/// at once. The drain re-reads the fleet each pass, so a call that reached -/// the coordinator before the error is safe to send again. +/// at once. async fn retry_unavailable( op: &'static str, attempts: u32, @@ -91,7 +90,9 @@ where } } -async fn call(op: &'static str, call: F) -> Result +/// A read of the fleet. Its failure ends the hook, so it waits for a +/// coordinator that answers `Unavailable`. +async fn read(op: &'static str, call: F) -> Result where F: FnMut() -> Fut, Fut: Future, tonic::Status>>, @@ -99,6 +100,18 @@ where retry_unavailable(op, UNAVAILABLE_ATTEMPTS, UNAVAILABLE_BACKOFF, call).await } +/// One step of a drain action: one try. The coordinator also answers +/// `Unavailable` when its own call to a host fails, and the drain sends a +/// failed action again on its next pass; a wait here would hold the other +/// actions of the pass. +async fn act(op: &'static str, call: F) -> Result +where + F: FnMut() -> Fut, + Fut: Future, tonic::Status>>, +{ + retry_unavailable(op, 1, Duration::ZERO, call).await +} + pub fn parse_host_id(id: &str) -> Result { id.parse() .map_err(|_| DrainError::Invalid(format!("coordinator returned malformed host id {id:?}"))) @@ -147,7 +160,7 @@ impl Coordinator { } pub async fn list_hosts(&self) -> Result, DrainError> { - let resp = call("list_hosts", || { + let resp = read("list_hosts", || { let mut fleet = self.fleet.clone(); async move { fleet.list_hosts(app::ListHostsRequest {}).await } }) @@ -167,7 +180,7 @@ impl Coordinator { } pub async fn cordon(&self, host: HostId, owner: &str) -> Result<(), DrainError> { - call("cordon", || { + act("cordon", || { let mut fleet = self.fleet.clone(); let req = app::CordonHostRequest { owner: owner.into(), @@ -180,7 +193,7 @@ impl Coordinator { } pub async fn list_sessions(&self) -> Result, DrainError> { - let resp = call("list_sessions", || { + let resp = read("list_sessions", || { let mut sessions = self.sessions.clone(); async move { sessions.list_sessions(app::ListSessionsRequest {}).await } }) @@ -200,7 +213,7 @@ impl Coordinator { /// Whether the session has a checkpoint it can resume from. pub async fn has_checkpoint(&self, id: &str) -> Result { - let resp = call("list_checkpoints", || { + let resp = act("list_checkpoints", || { let mut sessions = self.sessions.clone(); let req = app::ListCheckpointsRequest { session_id: id.into(), @@ -212,7 +225,7 @@ impl Coordinator { } pub async fn snapshot(&self, id: &str) -> Result<(), DrainError> { - call("snapshot", || { + act("snapshot", || { let mut sessions = self.sessions.clone(); let req = app::SnapshotRequest { session_id: id.into(), @@ -224,7 +237,7 @@ impl Coordinator { } pub async fn resume(&self, id: &str) -> Result<(), DrainError> { - call("resume", || { + act("resume", || { let mut sessions = self.sessions.clone(); let req = app::ResumeRequest { session_id: id.into(), @@ -238,7 +251,7 @@ impl Coordinator { /// The eviction that does not park: the session ends `idle` behind a /// fresh checkpoint, its VM destroyed. pub async fn evict(&self, id: &str) -> Result<(), DrainError> { - call("evict_local", || { + act("evict_local", || { let mut sessions = self.sessions.clone(); let req = app::EvictLocalRequest { session_id: id.into(), @@ -309,4 +322,17 @@ mod tests { ); assert_eq!(calls.load(Ordering::SeqCst), 3); } + + #[tokio::test] + async fn an_action_gets_one_try() { + let calls = AtomicU32::new(0); + let script = vec![tonic::Code::Unavailable]; + let err = act("evict_local", scripted(script, &calls)) + .await + .unwrap_err(); + assert!( + matches!(&err, DrainError::Rpc { status, .. } if status.code() == tonic::Code::Unavailable) + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } }