diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 119760422..b4cdbe818 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -2241,6 +2241,7 @@ dependencies = [ "safe-write", "schnorrkel", "scopeguard", + "sd-notify", "semver", "serde", "serde-human-bytes", diff --git a/dstack/dstack-util/Cargo.toml b/dstack/dstack-util/Cargo.toml index 9a603bff6..41c3e01cd 100644 --- a/dstack/dstack-util/Cargo.toml +++ b/dstack/dstack-util/Cargo.toml @@ -25,6 +25,7 @@ serde.workspace = true serde-human-bytes.workspace = true semver.workspace = true serde_json.workspace = true +sd-notify.workspace = true sha2.workspace = true tokio = { workspace = true, features = ["full"] } tracing.workspace = true diff --git a/dstack/dstack-util/src/gateway_checker.rs b/dstack/dstack-util/src/gateway_checker.rs new file mode 100644 index 000000000..e373c46d6 --- /dev/null +++ b/dstack/dstack-util/src/gateway_checker.rs @@ -0,0 +1,676 @@ +// SPDX-FileCopyrightText: © 2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Keeps this CVM's dstack-gateway registration alive after boot. +//! +//! Boot registers the CVM once, but that registration is not durable: the +//! gateway can be restarted, the WireGuard peers can change, and — since boot +//! no longer treats a gateway outage as fatal — the CVM may reach the +//! application start with no route at all. This supervisor closes that gap by +//! re-running the same registration whenever the observable state says it is +//! needed. +//! +//! There are exactly three reasons to refresh. The order matters, because +//! [`HANDSHAKE_TIMEOUT`] equals [`REFRESH_INTERVAL`] and the deadlines collide: +//! +//! 1. **No WireGuard config.** Boot-time registration never succeeded, so the +//! CVM has no route. Retried on a backoff (see [`Backoff`]). +//! 2. **Stale WireGuard handshake.** The tunnel exists but the peer stopped +//! answering for [`HANDSHAKE_TIMEOUT`], which usually means the gateway +//! restarted and forgot us. Checked *before* the periodic refresh and rate +//! limited to one attempt per timeout, because only this forced path can +//! rebuild a tunnel whose config is unchanged, and it is the expensive one. +//! 3. **Periodic re-registration.** The gateway expires idle registrations, so +//! re-register every [`REFRESH_INTERVAL`] even when everything looks fine. +//! +//! The decision logic is a pure function of an [`Observation`] so it can be +//! unit tested without a gateway, a KMS, or a WireGuard interface; all I/O +//! lives in [`cmd_gateway_checker`]. + +use std::path::{Path, PathBuf}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use anyhow::{Context, Result}; +use cmd_lib::run_fun as cmd; +use sd_notify::NotifyState; +use tracing::{error, info, warn}; + +use crate::system_setup::{GatewayRefresher, WG_CONFIG_PATH, WG_INTERFACE}; + +/// How often the loop samples the world. +const POLL_INTERVAL: Duration = Duration::from_secs(10); +/// Unconditional re-registration interval. +const REFRESH_INTERVAL: i64 = 180; +/// A tunnel with no handshake for this long is considered dead. +const HANDSHAKE_TIMEOUT: i64 = 180; +/// First retry delay when the CVM has no WireGuard config at all. +const MISSING_CONFIG_RETRY_INTERVAL: i64 = 30; +/// Upper bound for the retry backoff. +const MAX_RETRY_INTERVAL: i64 = 120; + +/// Exit code meaning "the gateway config is broken in a way retrying cannot +/// fix". Pinned by `RestartPreventExitStatus` in dstack-gateway-checker.service, so +/// changing it requires changing the unit too. +const EXIT_MISCONFIGURED: i32 = 3; + +#[derive(clap::Parser)] +/// Keep the dstack-gateway registration fresh +pub struct GatewayCheckerArgs { + /// dstack work directory + #[arg(long)] + work_dir: PathBuf, +} + +/// Exponential backoff over consecutive refresh failures. +/// +/// A refresh is not cheap: on a cold cache it costs a KMS round-trip, two +/// certificate signing requests and a TDX quote. A gateway outage is typically +/// fleet-wide, so a fixed short retry interval would have every CVM hammering +/// the KMS in lockstep and turn a gateway outage into a KMS outage. Backing off +/// to [`MAX_RETRY_INTERVAL`] keeps recovery prompt while bounding that load. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +struct Backoff { + consecutive_failures: u32, +} + +impl Backoff { + /// Minimum seconds that must elapse after the last attempt before the next + /// one. Zero while the last attempt succeeded. + fn delay(&self) -> i64 { + match self.consecutive_failures { + 0 => 0, + n => { + let shift = (n - 1).min(u32::BITS - 1); + MISSING_CONFIG_RETRY_INTERVAL + .checked_shl(shift) + .unwrap_or(MAX_RETRY_INTERVAL) + .min(MAX_RETRY_INTERVAL) + } + } + } + + fn record(&mut self, succeeded: bool) { + self.consecutive_failures = if succeeded { + 0 + } else { + self.consecutive_failures.saturating_add(1) + }; + } +} + +/// Everything the decision logic is allowed to look at. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Observation { + /// Seconds since the UNIX epoch. + now: i64, + /// Whether `/etc/wireguard/dstack-wg0.conf` exists. + config_present: bool, + /// Most recent handshake as a UNIX timestamp; `None` if the interface has + /// never completed one (or does not exist yet). + latest_handshake: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Refresh { + /// Force `wg-quick` to be reapplied even if the rendered config is + /// byte-identical to what is already on disk. + force: bool, + reason: &'static str, +} + +#[derive(Debug, Default)] +struct Checker { + /// Timestamp of the last refresh attempt; `None` if none has run yet. + last_attempt: Option, + /// Timestamp of the last *forced* refresh; `None` if none has run yet. + /// Tracked separately from `last_attempt` so a cheap periodic refresh does + /// not re-arm the rate limit that protects the expensive forced path. + last_force: Option, + /// When the tunnel was first seen without any handshake; `None` while a + /// handshake exists or the config is absent. Means exactly one thing: when + /// we first observed a peer with no handshake. + handshake_missing_since: Option, + backoff: Backoff, +} + +const MISSING_CONFIG: Refresh = Refresh { + force: true, + reason: "WireGuard config is missing", +}; +const STALE_HANDSHAKE: Refresh = Refresh { + force: true, + reason: "WireGuard handshake is stale", +}; +const PERIODIC: Refresh = Refresh { + force: false, + reason: "periodic re-registration", +}; + +impl Checker { + /// Build the starting state for a checker coming up at `now`. + /// + /// If the WireGuard config is already on disk, boot registered this CVM + /// moments ago: `/etc` is a volatile overlay, so the file can only exist + /// because *this* boot wrote it. Start the periodic clock from now instead + /// of re-registering immediately. Otherwise a fleet rebooting together + /// would hit the gateway with a second full round of registrations seconds + /// after the first — piling onto the component this checker exists to + /// tolerate the loss of. With no config, boot's registration failed and + /// recovering fast is the whole point, so leave the clock unset and act on + /// the first poll. + fn starting(now: i64, config_present: bool) -> Self { + Self { + last_attempt: config_present.then_some(now), + ..Self::default() + } + } + + /// Decide whether to refresh now. Pure: same state plus same observation + /// always yields the same answer. + fn decide(&mut self, obs: Observation) -> Option { + if !obs.config_present { + // No config means no interface, so there is no handshake to age. + self.handshake_missing_since = None; + let Some(last) = self.last_attempt else { + return Some(MISSING_CONFIG); + }; + let delay = MISSING_CONFIG_RETRY_INTERVAL.max(self.backoff.delay()); + return (obs.now.saturating_sub(last) >= delay).then_some(MISSING_CONFIG); + } + + // Staleness is checked BEFORE the periodic refresh, and the + // missing-handshake timer is cleared ONLY by an observed handshake -- + // never by a refresh. HANDSHAKE_TIMEOUT equals REFRESH_INTERVAL, so if + // a periodic refresh reset the timer it would re-arm at 0 while the + // timer had only reached REFRESH_INTERVAL - POLL_INTERVAL, and the + // forced branch would be unreachable for a peer that never handshakes. + // That matters because gateway setup short-circuits on an unchanged + // config unless force is set: a CVM whose WireGuard config is correct + // but whose tunnel is dead can only recover through a forced refresh. + let silent_since = match obs.latest_handshake { + Some(handshake) => { + self.handshake_missing_since = None; + handshake + } + // No handshake yet. Time it from when we first noticed rather than + // from process start, so a freshly created interface gets a full + // HANDSHAKE_TIMEOUT to complete its first handshake. + None => *self.handshake_missing_since.get_or_insert(obs.now), + }; + if obs.now.saturating_sub(silent_since) >= HANDSHAKE_TIMEOUT { + // A forced refresh bounces the interface and re-requests + // certificates, so cap it at one attempt per HANDSHAKE_TIMEOUT. + // Unthrottled, a gateway that stays down would be hit on every + // poll: a self-inflicted flood aimed at something already broken. + let due = match self.last_force { + None => true, + Some(last) => obs.now.saturating_sub(last) >= HANDSHAKE_TIMEOUT, + }; + // Return either way. While the tunnel is dead a periodic refresh + // is pointless, since only the forced path can rebuild it. + return due.then_some(STALE_HANDSHAKE); + } + + // The first poll always re-registers: boot may have left the CVM + // unregistered, and re-registering a healthy CVM is cheap. + let Some(last) = self.last_attempt else { + return Some(PERIODIC); + }; + (obs.now.saturating_sub(last) >= REFRESH_INTERVAL).then_some(PERIODIC) + } + + /// Record the outcome of a refresh triggered by [`Checker::decide`]. + fn record(&mut self, now: i64, refresh: Refresh, succeeded: bool) { + self.last_attempt = Some(now); + if refresh.force { + self.last_force = Some(now); + } + // handshake_missing_since is deliberately NOT cleared here; see decide(). + self.backoff.record(succeeded); + } +} + +fn now_secs() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0) +} + +/// Read the most recent handshake across all peers of the interface. +/// +/// `wg show latest-handshakes` prints one `\t` line per +/// peer, with `0` meaning "never". Returns `None` when the interface is absent, +/// has no peers, or no peer has ever completed a handshake — all of which the +/// caller treats the same way. +fn latest_handshake(interface: &str) -> Option { + let output = cmd!(wg show $interface latest-handshakes).ok()?; + parse_latest_handshake(&output) +} + +fn parse_latest_handshake(output: &str) -> Option { + output + .lines() + .filter_map(|line| line.split_whitespace().nth(1)) + .filter_map(|ts| ts.parse::().ok()) + .filter(|ts| *ts > 0) + .max() +} + +fn wg_config_present() -> bool { + Path::new(WG_CONFIG_PATH).exists() +} + +fn observe(now: i64) -> Observation { + let config_present = wg_config_present(); + Observation { + now, + config_present, + // Without a config the interface was never brought up, so there is no + // handshake to read and `decide` would not look at one anyway. Skipping + // the probe matters because that is precisely the state a gateway + // outage parks the CVM in: otherwise we would fork `wg` every poll for + // the entire outage to answer a question nobody asks. + latest_handshake: config_present + .then(|| latest_handshake(WG_INTERFACE)) + .flatten(), + } +} + +/// systemd liveness reporting, inert when the unit has no `WatchdogSec`. +/// +/// The recovery paths in this loop only work while the loop runs, and nothing +/// below it can guarantee that: a wedged refresh leaves a healthy-looking +/// process that systemd will never restart. Handing liveness to systemd covers +/// a hang wherever it comes from, including causes not anticipated here. +struct Watchdog { + enabled: bool, +} + +impl Watchdog { + /// Report readiness and arm the watchdog. Readiness is sent before the + /// checker decides whether it has anything to do, so the paths that exit + /// straight away are still a started service that then stopped, not a + /// service that failed to start. + fn arm() -> Self { + let mut usec = 0; + let enabled = sd_notify::watchdog_enabled(false, &mut usec); + if let Err(error) = sd_notify::notify(false, &[NotifyState::Ready]) { + warn!("failed to report readiness to systemd: {error}"); + } + if enabled { + info!("systemd watchdog armed, timeout={usec}us"); + } + Self { enabled } + } + + fn ping(&self) { + if !self.enabled { + return; + } + if let Err(error) = sd_notify::notify(false, &[NotifyState::Watchdog]) { + warn!("failed to ping the systemd watchdog: {error}"); + } + } +} + +pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { + let watchdog = Watchdog::arm(); + let refresher = + GatewayRefresher::load(&args.work_dir).context("failed to load gateway configuration")?; + + // Nothing to supervise for an app that never asked for a gateway. Exit + // successfully instead of polling forever; the unit is Restart=on-failure + // so systemd leaves the service alone. gateway_enabled is fixed by the + // measured app-compose and cannot change without a reboot. + if !refresher.gateway_enabled() { + info!("dstack-gateway is not enabled; nothing to check"); + return Ok(()); + } + + // A missing app id or gateway URL is a deployment mistake, not an outage. + // Both come from data fixed for the lifetime of the VM (app keys and the + // host-shared copy taken at setup), so no amount of retrying can fix it. + // Returning a plain error would have systemd restart us every RestartSec + // forever, so exit with the code the unit pins in RestartPreventExitStatus: + // that stops the respawn while still leaving the unit in `failed` state, + // which is what makes the mistake visible to the operator. + if let Err(error) = refresher.check_config() { + error!("dstack-gateway is enabled but misconfigured: {error:#}"); + error!("not retrying; this cannot be fixed without redeploying the CVM"); + std::process::exit(EXIT_MISCONFIGURED); + } + + info!("watching dstack-gateway registration"); + // The checker does not report gateway state to the host. Boot already + // reports the one signal that matters -- this CVM came up without a route + // -- and mirroring every later transition would mean tracking what the host + // has been told, which is state this loop should not have to carry. The + // consequence is that a boot error stays on the VMM after the checker + // recovers, until the VM restarts. + let mut checker = Checker::starting(now_secs(), wg_config_present()); + loop { + // Ping before the work, not after, so a refresh that never returns + // stops the pings. Nothing else can do this for us: a refresh spends + // most of its time in blocking `cmd!` shell-outs (`wg-quick up` alone + // resolves peer endpoints), which occupy a runtime worker with no await + // point. tokio::time::timeout cannot cancel that, and a watchdog task + // on another worker would happily keep pinging while this loop is + // wedged. Only the loop itself can prove the loop is alive. + watchdog.ping(); + + let now = now_secs(); + if let Some(refresh) = checker.decide(observe(now)) { + info!("refreshing dstack-gateway: {}", refresh.reason); + let succeeded = match refresher.refresh(refresh.force).await { + Ok(()) => { + info!("dstack-gateway refresh succeeded"); + true + } + Err(error) => { + warn!("dstack-gateway refresh failed: {error:#}"); + false + } + }; + // now_secs() is re-read here rather than reusing `now`: a refresh can + // block on network timeouts for a long time, and the backoff has to + // count from when the attempt ended, not when it started. + checker.record(now_secs(), refresh, succeeded); + } + tokio::time::sleep(POLL_INTERVAL).await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const T0: i64 = 1_700_000_000; + + fn obs(now_offset: i64, config_present: bool, latest_handshake: Option) -> Observation { + Observation { + now: T0 + now_offset, + config_present, + latest_handshake, + } + } + + /// Replay the checker against a scripted world for `duration` seconds at the + /// production poll interval, returning (forced, periodic) refresh counts. + /// + /// `handshake` maps elapsed seconds to what `wg show ... latest-handshakes` + /// would report at that moment. Every refresh is treated as succeeding, so + /// the counts isolate the decision logic from the backoff. + /// + /// Deliberately starts from `Checker::default()`, not `Checker::starting`: + /// the upstream harness began with `LAST_REFRESH=0`, so this reproduces its + /// cadence exactly. The startup grace period is covered separately. + fn replay(duration: i64, handshake: impl Fn(i64) -> Option) -> (usize, usize) { + let mut checker = Checker::default(); + let (mut forced, mut periodic) = (0, 0); + let mut elapsed = 0; + while elapsed < duration { + let observation = obs(elapsed, true, handshake(elapsed)); + if let Some(refresh) = checker.decide(observation) { + if refresh.force { + forced += 1; + } else { + periodic += 1; + } + checker.record(T0 + elapsed, refresh, true); + } + elapsed += POLL_INTERVAL.as_secs() as i64; + } + (forced, periodic) + } + + /// The scenario table from the upstream fix that made the forced branch + /// reachable and rate limited it (`fix(guest): let the missing-handshake + /// timer actually expire`). These counts are the contract that commit + /// established by replaying the shell checker against a mocked clock; the + /// Rust port has to reproduce them exactly or it has silently regressed. + #[test] + fn reproduces_the_upstream_forced_refresh_scenarios() { + // Healthy tunnel: handshakes keep landing, nothing is ever forced. + assert_eq!(replay(900, |t| Some(T0 + t)).0, 0, "healthy handshakes"); + + // A peer that never handshakes must still reach the forced path, once + // per HANDSHAKE_TIMEOUT: at 180, 360, 540 and 720 seconds. + assert_eq!(replay(900, |_| None).0, 4, "never handshakes"); + + // Frozen handshake with the gateway down: the timestamp never advances, + // so it ages past the timeout and is forced on the same cadence. Before + // the rate limit this was one forced refresh per poll. + assert_eq!( + replay(900, |_| Some(T0)).0, + 4, + "handshake frozen, gateway down" + ); + + // Handshakes land for the first 180s, then the peer goes silent. The + // timer runs from the last real handshake, so forcing starts at 360. + let (forced, _) = replay(900, |t| Some(T0 + if t < 180 { t } else { 180 })); + assert_eq!(forced, 3, "handshake then peer gone"); + + // No handshake, then the gateway recovers at 540s. Two forced attempts + // (180, 360) and then the periodic cadence resumes. + let (forced, periodic) = replay(900, |t| (t >= 540).then_some(T0 + t)); + assert_eq!(forced, 2, "no handshake, gateway recovers"); + assert!(periodic >= 1, "periodic refresh must resume after recovery"); + } + + /// The regression the upstream fix was about: a periodic refresh must not + /// reset the missing-handshake timer. HANDSHAKE_TIMEOUT == REFRESH_INTERVAL, + /// so clearing it on every periodic refresh re-arms the timer one poll + /// before it can expire and the forced branch becomes unreachable. + #[test] + fn periodic_refresh_does_not_re_arm_the_missing_handshake_timer() { + assert_eq!( + HANDSHAKE_TIMEOUT, REFRESH_INTERVAL, + "this regression only bites while the two intervals are equal" + ); + let mut checker = Checker::default(); + // First poll: no handshake yet, so the timer starts and we re-register. + assert_eq!(checker.decide(obs(0, true, None)), Some(PERIODIC)); + checker.record(T0, PERIODIC, true); + assert_eq!(checker.handshake_missing_since, Some(T0)); + + // Still no handshake one full timeout later. The forced branch must + // fire; if the refresh above had cleared the timer it would not. + assert_eq!( + checker.decide(obs(HANDSHAKE_TIMEOUT, true, None)), + Some(STALE_HANDSHAKE) + ); + assert_eq!( + checker.handshake_missing_since, + Some(T0), + "the timer must survive the refresh that ran at t=0" + ); + } + + #[test] + fn staleness_outranks_the_periodic_refresh() { + let mut checker = Checker::default(); + checker.record(T0, PERIODIC, true); + checker.handshake_missing_since = Some(T0); + // At exactly REFRESH_INTERVAL both deadlines are due. The forced path + // must win: only it can rebuild a tunnel whose config is unchanged. + assert_eq!( + checker.decide(obs(REFRESH_INTERVAL, true, None)), + Some(STALE_HANDSHAKE) + ); + } + + #[test] + fn forced_refresh_is_rate_limited_while_the_tunnel_stays_dead() { + let mut checker = Checker::default(); + checker.record(T0, PERIODIC, true); + checker.handshake_missing_since = Some(T0); + + let refresh = checker.decide(obs(180, true, None)).expect("forced"); + assert!(refresh.force); + checker.record(T0 + 180, refresh, true); + + // Every poll in the next window is still stale, but must stay quiet -- + // and must not fall through to a periodic refresh either. + let mut elapsed = 190; + while elapsed < 360 { + assert_eq!( + checker.decide(obs(elapsed, true, None)), + None, + "forced refresh must not repeat at t={elapsed}" + ); + elapsed += 10; + } + assert_eq!( + checker.decide(obs(360, true, None)), + Some(STALE_HANDSHAKE), + "one forced attempt per handshake timeout" + ); + } + + #[test] + fn an_observed_handshake_is_the_only_thing_that_clears_the_timer() { + let mut checker = Checker::default(); + checker.record(T0, PERIODIC, true); + + assert!(checker.decide(obs(10, true, None)).is_none()); + assert_eq!(checker.handshake_missing_since, Some(T0 + 10)); + // A handshake lands. + assert!(checker.decide(obs(20, true, Some(T0 + 20))).is_none()); + assert_eq!(checker.handshake_missing_since, None); + // Losing it again restarts the clock rather than firing immediately. + assert!(checker.decide(obs(30, true, None)).is_none()); + assert_eq!(checker.handshake_missing_since, Some(T0 + 30)); + } + + #[test] + fn parses_max_handshake_across_peers() { + let output = "aaa\t1700000000\nbbb\t1700000042\nccc\t0\n"; + assert_eq!(parse_latest_handshake(output), Some(1_700_000_042)); + } + + #[test] + fn treats_never_handshaked_peers_as_no_handshake() { + assert_eq!(parse_latest_handshake("aaa\t0\nbbb\t0\n"), None); + assert_eq!(parse_latest_handshake(""), None); + assert_eq!(parse_latest_handshake("garbage\n"), None); + } + + #[test] + fn acts_immediately_when_boot_left_no_config() { + let mut checker = Checker::starting(T0, false); + assert_eq!(checker.decide(obs(0, false, None)), Some(MISSING_CONFIG)); + } + + /// Boot registers the CVM, then this checker starts. Re-registering right + /// away would double the gateway's registration load on every fleet reboot + /// for no gain, so a config that is already on disk starts the periodic + /// clock rather than triggering an immediate refresh. + #[test] + fn does_not_re_register_a_cvm_boot_just_registered() { + let mut checker = Checker::starting(T0, true); + assert_eq!(checker.decide(obs(0, true, Some(T0))), None); + assert_eq!(checker.decide(obs(170, true, Some(T0 + 170))), None); + assert_eq!( + checker.decide(obs(180, true, Some(T0 + 180))), + Some(PERIODIC), + "the periodic clock still runs from process start" + ); + } + + /// A checker that comes up with no config must not inherit the grace period + /// above: that state means boot's registration failed and the CVM has no + /// route at all. + #[test] + fn a_missing_config_overrides_the_startup_grace_period() { + let mut checker = Checker::starting(T0, false); + assert_eq!(checker.decide(obs(0, false, None)), Some(MISSING_CONFIG)); + checker.record(T0, MISSING_CONFIG, true); + // Once it lands, the periodic clock takes over from the refresh. + assert_eq!(checker.decide(obs(10, true, Some(T0 + 10))), None); + assert_eq!( + checker.decide(obs(180, true, Some(T0 + 180))), + Some(PERIODIC) + ); + } + + #[test] + fn retries_missing_config_with_backoff_up_to_the_cap() { + let mut checker = Checker::default(); + let mut t = 0; + // First attempt fires immediately and fails. + assert_eq!(checker.decide(obs(t, false, None)), Some(MISSING_CONFIG)); + checker.record(T0 + t, MISSING_CONFIG, false); + + // 30s base delay, then 60s, then capped at MAX_RETRY_INTERVAL. + for expected_delay in [30, 60, 120, 120] { + for early in [10, expected_delay - 10] { + assert_eq!( + checker.decide(obs(t + early, false, None)), + None, + "must not retry after {early}s while waiting {expected_delay}s" + ); + } + t += expected_delay; + assert_eq!( + checker.decide(obs(t, false, None)), + Some(MISSING_CONFIG), + "must retry once {expected_delay}s have passed" + ); + checker.record(T0 + t, MISSING_CONFIG, false); + } + } + + #[test] + fn backoff_resets_after_a_success() { + let mut backoff = Backoff::default(); + assert_eq!(backoff.delay(), 0); + backoff.record(false); + assert_eq!(backoff.delay(), 30); + backoff.record(false); + assert_eq!(backoff.delay(), 60); + backoff.record(true); + assert_eq!(backoff.delay(), 0); + } + + #[test] + fn backoff_saturates_instead_of_overflowing() { + let backoff = Backoff { + consecutive_failures: u32::MAX, + }; + assert_eq!(backoff.delay(), MAX_RETRY_INTERVAL); + } + + #[test] + fn re_registers_periodically_while_healthy() { + let mut checker = Checker::default(); + checker.record(T0, PERIODIC, true); + + assert_eq!(checker.decide(obs(170, true, Some(T0 + 170))), None); + assert_eq!( + checker.decide(obs(180, true, Some(T0 + 180))), + Some(PERIODIC), + "periodic refresh must not disrupt a healthy tunnel" + ); + } + + #[test] + fn gives_a_new_interface_a_full_timeout_to_handshake() { + let mut checker = Checker::default(); + checker.record(T0, PERIODIC, true); + + // The timer runs from t=10, when the missing handshake was first + // observed, not from the refresh at t=0. + assert_eq!(checker.decide(obs(10, true, None)), None); + assert_eq!(checker.decide(obs(170, true, None)), None); + // At t=180 the timer has only reached 170, so the periodic refresh is + // what comes due. It must not disturb the timer. + assert_eq!(checker.decide(obs(180, true, None)), Some(PERIODIC)); + checker.record(T0 + 180, PERIODIC, true); + assert_eq!(checker.handshake_missing_since, Some(T0 + 10)); + // One poll later the timer finally expires and forces a refresh. + assert_eq!(checker.decide(obs(190, true, None)), Some(STALE_HANDSHAKE)); + } +} diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index 1347b8e1b..c746d736f 100644 --- a/dstack/dstack-util/src/main.rs +++ b/dstack/dstack-util/src/main.rs @@ -7,6 +7,7 @@ use clap::{Parser, Subcommand}; use dstack_attest::emit_runtime_event; use dstack_types::{KeyProvider, KeyProviderKind}; use fs_err as fs; +use gateway_checker::{cmd_gateway_checker, GatewayCheckerArgs}; use getrandom::fill as getrandom; use host_api::HostApi; use k256::schnorr::SigningKey; @@ -30,6 +31,7 @@ use utils::AppKeys; mod crypto; mod docker_compose; +mod gateway_checker; mod host_api; mod host_shared; mod parse_env_file; @@ -72,6 +74,8 @@ enum Commands { HostShared(host_shared::HostSharedArgs), /// Refresh the dstack gateway configuration GatewayRefresh(GatewayRefreshArgs), + /// Keep the dstack gateway registration fresh (long-running) + GatewayChecker(GatewayCheckerArgs), /// Notify the host about the dstack app NotifyHost(HostNotifyArgs), /// Remove orphaned containers @@ -1322,6 +1326,9 @@ async fn main() -> Result<()> { cmd_sys_setup(args).await?; } Commands::HostShared(args) => host_shared::cmd_host_shared(args)?, + Commands::GatewayChecker(args) => { + cmd_gateway_checker(args).await?; + } Commands::GatewayRefresh(args) => { cmd_gateway_refresh(args).await?; } diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index ec2f4eecc..c593976c9 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -320,7 +320,9 @@ impl HostShared { } const GATEWAY_CACHE_PATH: &str = "/run/dstack/gateway-cache.json"; -const WG_CONFIG_PATH: &str = "/etc/wireguard/dstack-wg0.conf"; +/// Name of the WireGuard interface linking this CVM to dstack-gateway. +pub const WG_INTERFACE: &str = "dstack-wg0"; +pub const WG_CONFIG_PATH: &str = "/etc/wireguard/dstack-wg0.conf"; /// Certificate validity period in seconds (10 days) const CERT_VALIDITY_SECS: u64 = 10 * 24 * 3600; const MAX_SUPPORTED_MANIFEST_VERSION: u32 = 3; @@ -555,6 +557,14 @@ impl<'a> GatewayContext<'a> { // Get or generate key store (includes WireGuard keys and client certificate) let key_store = self.get_or_generate_key_store().await?; + // Persist the key store before attempting registration. Minting it costs a + // KMS round-trip, two cert signing requests and a TDX quote, so a gateway + // outage would otherwise make every retry pay that price again and turn a + // gateway outage into a KMS load spike across the whole fleet. + if let Err(e) = key_store.save() { + warn!("failed to save gateway cache: {e:?}"); + } + if self.shared.sys_config.gateway_urls.is_empty() { bail!("Missing gateway urls"); } @@ -604,11 +614,6 @@ impl<'a> GatewayContext<'a> { )); } - // Save cache - if let Err(e) = key_store.save() { - warn!("Failed to save gateway cache: {e:?}"); - } - // Check if config has changed (skip check if force is set) if !force { let current_config = fs::read_to_string(WG_CONFIG_PATH).ok(); @@ -1931,19 +1936,61 @@ impl Stage0<'_> { } } -pub async fn cmd_gateway_refresh(args: GatewayRefreshArgs) -> Result<()> { - let host_shared_dir = args.work_dir.join(HOST_SHARED_DIR_NAME); - let shared = HostShared::load(host_shared_dir.as_path()).with_context(|| { - format!( - "Failed to load host-shared dir: {}", - host_shared_dir.display() - ) - })?; - let keys_path = shared.dir.join(APP_KEYS); - let keys: AppKeys = deserialize_json_file(&keys_path) - .with_context(|| format!("Failed to load app keys from {}", keys_path.display()))?; +/// Owns the inputs needed to (re)register this CVM with dstack-gateway. +/// +/// Loading is separated from refreshing so a long-running caller (the gateway +/// checker) can pay the parsing cost once and then refresh repeatedly. +pub struct GatewayRefresher { + shared: HostShared, + keys: AppKeys, +} + +impl GatewayRefresher { + /// Load the host-shared config and app keys from `work_dir`. + pub fn load(work_dir: &Path) -> Result { + let host_shared_dir = work_dir.join(HOST_SHARED_DIR_NAME); + let shared = HostShared::load(host_shared_dir.as_path()).with_context(|| { + format!( + "Failed to load host-shared dir: {}", + host_shared_dir.display() + ) + })?; + let keys_path = shared.dir.join(APP_KEYS); + let keys: AppKeys = deserialize_json_file(&keys_path) + .with_context(|| format!("Failed to load app keys from {}", keys_path.display()))?; + Ok(Self { shared, keys }) + } + + /// Whether this app opted into dstack-gateway at all. + pub fn gateway_enabled(&self) -> bool { + self.shared.app_compose.gateway_enabled() + } - GatewayContext::new(&shared, &keys).setup(args.force).await + /// Validate the parts of the gateway config that can never become valid by + /// waiting. These are deployment mistakes, not outages, so callers that + /// retry should give up instead of looping forever. + pub fn check_config(&self) -> Result<()> { + if self.keys.gateway_app_id.is_empty() { + bail!("Missing allowed dstack-gateway app id"); + } + if self.shared.sys_config.gateway_urls.is_empty() { + bail!("Missing gateway urls"); + } + Ok(()) + } + + /// Register with dstack-gateway and apply the returned WireGuard config. + pub async fn refresh(&self, force: bool) -> Result<()> { + GatewayContext::new(&self.shared, &self.keys) + .setup(force) + .await + } +} + +pub async fn cmd_gateway_refresh(args: GatewayRefreshArgs) -> Result<()> { + GatewayRefresher::load(&args.work_dir)? + .refresh(args.force) + .await } struct AppIdValidator { @@ -2847,9 +2894,27 @@ impl Stage1<'_> { self.vmm .notify_q("boot.progress", "setting up dstack-gateway") .await; - GatewayContext::new(&self.shared, &self.keys) + if let Err(error) = GatewayContext::new(&self.shared, &self.keys) .setup(true) - .await?; + .await + { + warn!( + "dstack-gateway registration is unavailable during boot; continuing without a route: {error:#}" + ); + // Boot no longer fails here, so a guest log line would be the only + // trace of it: the VM would report a clean boot while having no + // ingress at all. Report it to the host so the degraded state is + // visible from the VMM. The gateway checker clears this once it + // manages to register. + self.vmm + .notify_q( + "boot.error", + &format!( + "dstack-gateway registration failed, the app has no ingress route: {error:#}" + ), + ) + .await; + } self.vmm .notify_q("boot.progress", "setting up docker") .await; diff --git a/os/common/rootfs/dstack-gateway-checker.service b/os/common/rootfs/dstack-gateway-checker.service new file mode 100644 index 000000000..4f183eb8c --- /dev/null +++ b/os/common/rootfs/dstack-gateway-checker.service @@ -0,0 +1,31 @@ +[Unit] +Description=dstack Gateway Registration Checker +After=network-online.target dstack-prepare.service +Wants=network-online.target + +[Service] +Type=notify +ExecStart=/bin/dstack-util gateway-checker --work-dir /dstack +# Every recovery path in the checker only runs while its loop runs, and a loop +# that wedges leaves a healthy-looking process systemd would never restart. The +# loop pings the watchdog itself, so a hang anywhere -- including the blocking +# wg-quick/iptables shell-outs a refresh performs, which no in-process timeout +# can cancel -- gets the service killed and restarted. The timeout is generous +# because one refresh may legitimately spend minutes across KMS certificate +# requests and every configured gateway URL. +WatchdogSec=600 +# The checker exits 0 when the app never enabled dstack-gateway, because there +# is then nothing to supervise. Restart=always would respawn that exit forever. +Restart=on-failure +RestartSec=10 +# Exit code 3 means the gateway config is broken in a way retrying cannot fix +# (no gateway app id, no gateway URLs). Both are fixed for the lifetime of the +# VM, so respawning every RestartSec would just be a slower spin. Stop +# restarting but stay in `failed` state so the mistake is visible. Keep in sync +# with EXIT_MISCONFIGURED in dstack-util's gateway_checker. +RestartPreventExitStatus=3 +StandardOutput=journal +StandardError=journal+console + +[Install] +WantedBy=multi-user.target diff --git a/os/common/rootfs/wg-checker.service b/os/common/rootfs/wg-checker.service deleted file mode 100644 index 406cadc69..000000000 --- a/os/common/rootfs/wg-checker.service +++ /dev/null @@ -1,15 +0,0 @@ -[Unit] -Description=WireGuard Endpoint Checker Service -After=network-online.target dstack-prepare.service -Wants=network-online.target - -[Service] -Type=simple -ExecStart=/bin/wg-checker.sh -Restart=always -RestartSec=10 -StandardOutput=journal -StandardError=journal+console - -[Install] -WantedBy=multi-user.target diff --git a/os/common/rootfs/wg-checker.sh b/os/common/rootfs/wg-checker.sh deleted file mode 100755 index 11773088e..000000000 --- a/os/common/rootfs/wg-checker.sh +++ /dev/null @@ -1,97 +0,0 @@ -#!/bin/sh - -# SPDX-FileCopyrightText: © 2024 Phala Network -# -# SPDX-License-Identifier: Apache-2.0 - -HANDSHAKE_TIMEOUT=180 -REFRESH_INTERVAL=180 -LAST_REFRESH=0 -LAST_FORCE=0 -STALE_SINCE=0 -DSTACK_WORK_DIR=${DSTACK_WORK_DIR:-/dstack} -IFNAME=dstack-wg0 - -get_latest_handshake() { - wg show $IFNAME latest-handshakes 2>/dev/null | awk 'BEGIN { max = 0 } NF >= 2 { if ($2 > max) max = $2 } END { print max }' -} - -do_refresh() { - now=$1 - reason=$2 - force=$3 - - if ! command -v dstack-util >/dev/null 2>&1; then - printf 'dstack-util not found; cannot refresh gateway.\n' >&2 - LAST_REFRESH=$now - return - fi - - printf '%s; refreshing dstack gateway...\n' "$reason" - if [ "$force" = "1" ]; then - cmd="dstack-util gateway-refresh --work-dir $DSTACK_WORK_DIR --force" - else - cmd="dstack-util gateway-refresh --work-dir $DSTACK_WORK_DIR" - fi - if $cmd; then - printf 'dstack gateway refresh succeeded.\n' - else - printf 'dstack gateway refresh failed.\n' >&2 - fi - - LAST_REFRESH=$now -} - -check_and_refresh() { - if ! command -v wg >/dev/null 2>&1; then - return - fi - - now=$(date +%s) - - latest=$(get_latest_handshake) - if [ -z "$latest" ]; then - latest=0 - fi - - # A peer that never completed a handshake reports 0, so there is no - # timestamp to age against; fall back to when we first observed that state. - # STALE_SINCE is cleared only by an actual handshake, never by a refresh: - # since HANDSHAKE_TIMEOUT equals REFRESH_INTERVAL, clearing it on every - # periodic refresh re-arms the timer one tick before it can expire, leaving - # the forced branch unreachable for a peer that never handshakes. - if [ "$latest" -gt 0 ]; then - STALE_SINCE=0 - stale_for=$((now - latest)) - else - if [ "$STALE_SINCE" -eq 0 ]; then - STALE_SINCE=$now - fi - stale_for=$((now - STALE_SINCE)) - fi - - # Forced refresh bounces the interface and re-requests certificates, so cap - # it at one attempt per HANDSHAKE_TIMEOUT. Without this, a gateway that - # stays unreachable would be retried on every 10s tick. - if [ "$stale_for" -ge $HANDSHAKE_TIMEOUT ]; then - if [ $((now - LAST_FORCE)) -ge $HANDSHAKE_TIMEOUT ]; then - do_refresh "$now" "WireGuard handshake stale" 1 >&2 - LAST_FORCE=$now - fi - return - fi - - # Periodic refresh every REFRESH_INTERVAL seconds (not forced). - if [ "$LAST_REFRESH" -eq 0 ] || [ $((now - LAST_REFRESH)) -ge $REFRESH_INTERVAL ]; then - do_refresh "$now" "Periodic refresh" 0 - fi -} - -while true; do - if [ -f /etc/wireguard/$IFNAME.conf ]; then - check_and_refresh - else - STALE_SINCE=0 - fi - sleep 10 -done diff --git a/os/mkosi/components/dstack-rust/dstack-rust-build.sh b/os/mkosi/components/dstack-rust/dstack-rust-build.sh index 3f8d4729a..e5f5a5298 100755 --- a/os/mkosi/components/dstack-rust/dstack-rust-build.sh +++ b/os/mkosi/components/dstack-rust/dstack-rust-build.sh @@ -9,7 +9,7 @@ install -d "$DEST/usr/bin" "$DEST/usr/lib/systemd/system" \ "$DEST/etc/systemd/journald.conf.d" "$DEST/etc/systemd/resolved.conf.d" \ "$DEST/etc/systemd/system/docker.service.d" \ "$DEST/etc/systemd/system/containerd.service.d" "$DEST/etc/sysctl.d" -for s in dstack-prepare ephemeral-docker wg-checker app-compose; do +for s in dstack-prepare ephemeral-docker app-compose; do install -m0755 "$ROOT/os/common/rootfs/$s.sh" "$DEST/usr/bin/$s.sh" done install -m0644 "$ROOT/os/common/rootfs/"*.service \ diff --git a/os/mkosi/mkosi.skeleton/usr/lib/systemd/system-preset/80-dstack.preset b/os/mkosi/mkosi.skeleton/usr/lib/systemd/system-preset/80-dstack.preset index 0cdd9bffa..28938bddd 100644 --- a/os/mkosi/mkosi.skeleton/usr/lib/systemd/system-preset/80-dstack.preset +++ b/os/mkosi/mkosi.skeleton/usr/lib/systemd/system-preset/80-dstack.preset @@ -10,7 +10,7 @@ enable dstack-prepare.service enable dstack-guest-agent.socket enable dstack-guest-agent.service enable app-compose.service -enable wg-checker.service +enable dstack-gateway-checker.service enable nvidia-persistenced.service enable nvidia-fabricmanager.service enable containerd-stargz-grpc.service diff --git a/os/mkosi/tests/acceptance.sh b/os/mkosi/tests/acceptance.sh index 21f4f8007..5d26ea08b 100755 --- a/os/mkosi/tests/acceptance.sh +++ b/os/mkosi/tests/acceptance.sh @@ -66,9 +66,34 @@ done < <(sed -n 's/^require_config \(CONFIG_[A-Z0-9_]*\) \([ynm]\)$/\1 \2/p' "$a [[ $required -ge 10 ]] grep -q '0002-acpi-sandbox' "$D/components/kernel/kernel-build.sh" grep -q -- '--fuzz=0' "$D/components/kernel/kernel-build.sh" -for service in dstack-guest-agent dstack-prepare app-compose wg-checker; do +for service in dstack-guest-agent dstack-prepare app-compose dstack-gateway-checker; do grep -q "$service" "$D/mkosi.skeleton/usr/lib/systemd/system-preset/80-dstack.preset" done +# The gateway checker is a dstack-util subcommand, not a shell script. It exits 0 +# when the app did not enable dstack-gateway, so Restart must not be "always" or +# systemd respawns it every RestartSec forever on every gateway-less CVM. +gw_unit="$D/../common/rootfs/dstack-gateway-checker.service" +grep -q '^ExecStart=/bin/dstack-util gateway-checker ' "$gw_unit" || { + echo 'dstack-gateway-checker.service must run the dstack-util gateway-checker subcommand'; exit 1; } +grep -q '^Restart=on-failure$' "$gw_unit" || { + echo 'dstack-gateway-checker.service must use Restart=on-failure'; exit 1; } +# A permanent misconfiguration exits with EXIT_MISCONFIGURED. If the unit does +# not inhibit restarts for exactly that code, "fail loudly and stop" silently +# degrades into a RestartSec respawn loop, which is what this whole exit code +# exists to avoid. +checker_src="$D/../../dstack/dstack-util/src/gateway_checker.rs" +exit_code=$(sed -n 's/^const EXIT_MISCONFIGURED: i32 = \([0-9]\+\);$/\1/p' "$checker_src") +[[ -n $exit_code ]] || { echo "cannot read EXIT_MISCONFIGURED from $checker_src"; exit 1; } +grep -q "^RestartPreventExitStatus=${exit_code}\$" "$gw_unit" || { + echo "dstack-gateway-checker.service must set RestartPreventExitStatus=$exit_code"; exit 1; } +# The loop's recovery paths only run while the loop runs, so systemd has to be +# the thing that notices a wedge. WatchdogSec is useless without Type=notify. +grep -q '^Type=notify$' "$gw_unit" || { + echo 'dstack-gateway-checker.service must use Type=notify to arm the watchdog'; exit 1; } +grep -q '^WatchdogSec=' "$gw_unit" || { + echo 'dstack-gateway-checker.service must set WatchdogSec'; exit 1; } +test ! -e "$D/../common/rootfs/wg-checker.sh" +test ! -e "$D/../common/rootfs/wg-checker.service" # systemd enables any unit that matches no preset rule, so the enable list is # only meaningful with a terminal disable. Without it, every package pulled in # by Packages= would start at boot with no diff to 80-dstack.preset. diff --git a/os/yocto/layers/meta-dstack/recipes-core/dstack-guest/dstack-guest.bb b/os/yocto/layers/meta-dstack/recipes-core/dstack-guest/dstack-guest.bb index 8dfd86e7e..197bed16f 100644 --- a/os/yocto/layers/meta-dstack/recipes-core/dstack-guest/dstack-guest.bb +++ b/os/yocto/layers/meta-dstack/recipes-core/dstack-guest/dstack-guest.bb @@ -28,7 +28,7 @@ export AWS_LC_SYS_CMAKE_BUILDER = "1" # Ensure rsync-native is built before unpack runs do_unpack[depends] += "rsync-native:do_populate_sysroot" -DSTACK_SERVICES = "dstack-guest-agent.service dstack-guest-agent.socket dstack-prepare.service app-compose.service wg-checker.service" +DSTACK_SERVICES = "dstack-guest-agent.service dstack-guest-agent.socket dstack-prepare.service app-compose.service dstack-gateway-checker.service" SYSTEMD_PACKAGES = "${@bb.utils.contains('DISTRO_FEATURES','systemd','${PN}','',d)}" SYSTEMD_SERVICE:${PN} = "${@bb.utils.contains('DISTRO_FEATURES','systemd','${DSTACK_SERVICES}','',d)}" SYSTEMD_AUTO_ENABLE:${PN} = "enable" @@ -69,7 +69,6 @@ do_install() { install -m 0755 ${CARGO_BINDIR}/dstack-volume ${D}${bindir} install -m 0755 ${DSTACK_ROOTFS_FILES}/dstack-prepare.sh ${D}${bindir} install -m 0755 ${DSTACK_ROOTFS_FILES}/ephemeral-docker.sh ${D}${bindir} - install -m 0755 ${DSTACK_ROOTFS_FILES}/wg-checker.sh ${D}${bindir} install -m 0755 ${DSTACK_ROOTFS_FILES}/app-compose.sh ${D}${bindir} install -m 0644 ${DSTACK_ROOTFS_FILES}/journald.conf ${D}${sysconfdir}/systemd/journald.conf.d/dstack.conf @@ -86,7 +85,7 @@ do_install() { install -m 0644 ${DSTACK_ROOTFS_FILES}/dstack-guest-agent.service ${D}${systemd_system_unitdir} install -m 0644 ${DSTACK_ROOTFS_FILES}/dstack-prepare.service ${D}${systemd_system_unitdir} install -m 0644 ${DSTACK_ROOTFS_FILES}/app-compose.service ${D}${systemd_system_unitdir} - install -m 0644 ${DSTACK_ROOTFS_FILES}/wg-checker.service ${D}${systemd_system_unitdir} + install -m 0644 ${DSTACK_ROOTFS_FILES}/dstack-gateway-checker.service ${D}${systemd_system_unitdir} install -m 0644 ${DSTACK_ROOTFS_FILES}/dstack-guest-agent.socket ${D}${systemd_system_unitdir} install -m 0644 ${DSTACK_ROOTFS_FILES}/llmnr.conf ${D}${sysconfdir}/systemd/resolved.conf.d install -d ${D}${sysconfdir}/systemd/system/docker.service.d