From ef05a83a1e8389f16e1e87967890e5fa83bdedcd Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 30 Jul 2026 14:31:00 +0000 Subject: [PATCH 1/7] fix(guest): decouple Gateway outage from app boot --- dstack/dstack-util/src/system_setup.rs | 9 +++++++-- os/common/rootfs/wg-checker.sh | 5 +++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 6f258d34b..f55a3bcc5 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -2775,9 +2775,14 @@ 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:#}" + ); + } self.vmm .notify_q("boot.progress", "setting up docker") .await; diff --git a/os/common/rootfs/wg-checker.sh b/os/common/rootfs/wg-checker.sh index ba4149472..b81428187 100755 --- a/os/common/rootfs/wg-checker.sh +++ b/os/common/rootfs/wg-checker.sh @@ -6,6 +6,7 @@ HANDSHAKE_TIMEOUT=180 REFRESH_INTERVAL=180 +MISSING_CONFIG_RETRY_INTERVAL=30 LAST_REFRESH=0 STALE_SINCE=0 DSTACK_WORK_DIR=${DSTACK_WORK_DIR:-/dstack} @@ -80,6 +81,10 @@ while true; do check_and_refresh else STALE_SINCE=0 + now=$(date +%s) + if [ "$LAST_REFRESH" -eq 0 ] || [ $((now - LAST_REFRESH)) -ge $MISSING_CONFIG_RETRY_INTERVAL ]; then + do_refresh "$now" "WireGuard configuration missing" 1 + fi fi sleep 10 done From 5ff0768ee65f34b9e676c522807565f2b8218189 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 04:43:16 -0700 Subject: [PATCH 2/7] fix(guest): bound gateway retry cost and move wg-checker to Rust Follow-up to the boot decoupling. Three problems with retrying gateway registration after boot: 1. `GatewayKeyStore::save()` ran only after `register_cvm` succeeded, so during a gateway outage the cache in /run was never written and every retry re-minted the key store: a KMS round-trip, two cert signing requests and a TDX quote. A gateway outage is fleet-wide, so a fixed 30s retry had the whole fleet hammering the KMS. Persist the key store as soon as it is minted, and back the retry off 30s -> 60s -> 120s. 2. `setup()` returns early when the app never enabled dstack-gateway, so those CVMs never have a WireGuard config and the missing-config retry spun forever on them. The checker now reads `gateway_enabled()` directly and exits 0; the unit switches to Restart=on-failure. 3. A missing gateway app id or gateway URL is a deployment mistake, not an outage. `GatewayRefresher::check_config` separates those so the checker fails loudly instead of retrying something that can never succeed. wg-checker.sh is replaced by `dstack-util gateway-checker`. The refresh decision is a pure function of an observation, so the three trigger conditions, the backoff and the handshake timing are unit tested without a gateway, a KMS or a WireGuard interface. --- dstack/dstack-util/src/gateway_checker.rs | 441 ++++++++++++++++++ dstack/dstack-util/src/main.rs | 7 + dstack/dstack-util/src/system_setup.rs | 83 +++- os/common/rootfs/wg-checker.service | 14 +- os/common/rootfs/wg-checker.sh | 90 ---- .../dstack-rust/dstack-rust-build.sh | 2 +- os/mkosi/tests/acceptance.sh | 18 + .../recipes-core/dstack-guest/dstack-guest.bb | 1 - 8 files changed, 543 insertions(+), 113 deletions(-) create mode 100644 dstack/dstack-util/src/gateway_checker.rs delete mode 100755 os/common/rootfs/wg-checker.sh diff --git a/dstack/dstack-util/src/gateway_checker.rs b/dstack/dstack-util/src/gateway_checker.rs new file mode 100644 index 000000000..e72fb4422 --- /dev/null +++ b/dstack/dstack-util/src/gateway_checker.rs @@ -0,0 +1,441 @@ +// 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, in priority order: +//! +//! 1. **No WireGuard config.** Boot-time registration never succeeded, so the +//! CVM has no route. Retried on a backoff (see [`Backoff`]). +//! 2. **Periodic re-registration.** The gateway expires idle registrations, so +//! re-register every [`REFRESH_INTERVAL`] even when everything looks fine. +//! 3. **Stale WireGuard handshake.** The tunnel exists but the peer stopped +//! answering for [`HANDSHAKE_TIMEOUT`], which usually means the gateway +//! restarted and forgot us. +//! +//! 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 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 wg-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, + /// When the tunnel was first seen without any handshake; `None` while a + /// handshake exists or the config is absent. + handshake_missing_since: Option, + backoff: Backoff, +} + +impl Checker { + /// Decide whether to refresh now. Pure: same state plus same observation + /// always yields the same answer. + fn decide(&mut self, obs: Observation) -> Option { + // An attempt that has not cooled down yet blocks every reason below, so + // repeated failures cannot bypass the backoff by changing category. + let elapsed = match self.last_attempt { + // Nothing tried yet: act immediately. + None => return Some(Self::first_refresh(obs)), + Some(last) => obs.now.saturating_sub(last), + }; + if elapsed < self.backoff.delay() { + return None; + } + + if !obs.config_present { + self.handshake_missing_since = None; + return (elapsed >= MISSING_CONFIG_RETRY_INTERVAL).then_some(Refresh { + force: true, + reason: "WireGuard config is missing", + }); + } + + if elapsed >= REFRESH_INTERVAL { + self.handshake_missing_since = None; + return Some(Refresh { + force: false, + reason: "periodic re-registration", + }); + } + + // The tunnel is configured and recently re-registered; the only thing + // left that can be wrong is the tunnel itself going quiet. + 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), + }; + (obs.now.saturating_sub(silent_since) >= HANDSHAKE_TIMEOUT).then_some(Refresh { + force: true, + reason: "WireGuard handshake is stale", + }) + } + + /// The first poll always acts: boot may have left the CVM unregistered, and + /// re-registering an already-healthy CVM is cheap and idempotent. + fn first_refresh(obs: Observation) -> Refresh { + if obs.config_present { + Refresh { + force: false, + reason: "initial re-registration", + } + } else { + Refresh { + force: true, + reason: "WireGuard config is missing", + } + } + } + + /// Record the outcome of a refresh triggered by [`Checker::decide`]. + fn record(&mut self, now: i64, succeeded: bool) { + self.last_attempt = Some(now); + self.handshake_missing_since = None; + 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 observe(now: i64) -> Observation { + Observation { + now, + config_present: Path::new(WG_CONFIG_PATH).exists(), + latest_handshake: latest_handshake(WG_INTERFACE), + } +} + +pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { + 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"); + let mut checker = Checker::default(); + loop { + 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) => { + // now_secs() is re-read below: the refresh itself can block + // on network timeouts for a long time, and the backoff must + // count from when the attempt ended. + warn!("dstack-gateway refresh failed: {error:#}"); + false + } + }; + checker.record(now_secs(), 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, + } + } + + #[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_on_first_poll() { + let mut checker = Checker::default(); + let refresh = checker.decide(obs(0, false, None)).expect("should refresh"); + assert!(refresh.force); + + let mut checker = Checker::default(); + let refresh = checker.decide(obs(0, true, None)).expect("should refresh"); + assert!(!refresh.force, "an existing config needs no forced reapply"); + } + + #[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!(checker.decide(obs(t, false, None)).is_some()); + checker.record(T0 + t, 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!( + checker.decide(obs(t + early, false, None)).is_none(), + "must not retry after {early}s while waiting {expected_delay}s" + ); + } + t += expected_delay; + assert!( + checker.decide(obs(t, false, None)).is_some(), + "must retry once {expected_delay}s have passed" + ); + checker.record(T0 + t, 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, true); + + // Healthy tunnel, fresh handshake: quiet until REFRESH_INTERVAL. + assert!(checker.decide(obs(170, true, Some(T0 + 170))).is_none()); + let refresh = checker + .decide(obs(180, true, Some(T0 + 180))) + .expect("periodic refresh is due"); + assert!( + !refresh.force, + "periodic refresh must not disrupt the tunnel" + ); + } + + #[test] + fn forces_refresh_when_the_handshake_goes_stale() { + let mut checker = Checker::default(); + checker.record(T0, true); + + // Handshake 179s old: still within tolerance. + assert!(checker.decide(obs(170, true, Some(T0 - 9))).is_none()); + // 180s old and we have not hit the periodic interval yet. + let refresh = checker + .decide(obs(170, true, Some(T0 - 10))) + .expect("stale handshake must force a refresh"); + assert!(refresh.force); + assert_eq!(refresh.reason, "WireGuard handshake is stale"); + } + + #[test] + fn gives_a_new_interface_a_full_timeout_to_handshake() { + let mut checker = Checker::default(); + // Config appeared at t=0 but no handshake yet. + checker.record(T0, true); + + assert!(checker.decide(obs(10, true, None)).is_none()); + assert!(checker.decide(obs(179, true, None)).is_none()); + // The timeout runs from t=10, when the missing handshake was first + // observed, not from the refresh at t=0. + assert!(checker.decide(obs(189, true, None)).is_some()); + } + + #[test] + fn a_recovered_handshake_clears_the_missing_timer() { + let mut checker = Checker::default(); + checker.record(T0, true); + + assert!(checker.decide(obs(10, true, None)).is_none()); + // Handshake completes. + assert!(checker.decide(obs(20, true, Some(T0 + 20))).is_none()); + assert_eq!(checker.handshake_missing_since, None); + // Losing it again restarts the clock instead of firing immediately. + assert!(checker.decide(obs(30, true, None)).is_none()); + assert_eq!(checker.handshake_missing_since, Some(T0 + 30)); + } + + #[test] + fn backoff_outranks_the_missing_config_interval() { + let mut checker = Checker::default(); + // Three consecutive failures put the backoff at the 120s cap, which is + // longer than the 30s missing-config interval. + checker.backoff.consecutive_failures = 3; + checker.record(T0, false); + assert_eq!(checker.backoff.consecutive_failures, 4); + + assert!(checker.decide(obs(30, false, None)).is_none()); + assert!(checker.decide(obs(119, false, None)).is_none()); + assert!(checker.decide(obs(120, false, None)).is_some()); + } +} diff --git a/dstack/dstack-util/src/main.rs b/dstack/dstack-util/src/main.rs index 9bb2e5e3b..7c688289e 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; @@ -29,6 +30,7 @@ use utils::AppKeys; mod crypto; mod docker_compose; +mod gateway_checker; mod host_api; mod host_shared; mod parse_env_file; @@ -71,6 +73,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 @@ -1296,6 +1300,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 f55a3bcc5..35c64d01c 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -315,7 +315,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; @@ -550,6 +552,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"); } @@ -599,11 +609,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(); @@ -1923,19 +1928,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() + } + + /// 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 + } +} - GatewayContext::new(&shared, &keys).setup(args.force).await +pub async fn cmd_gateway_refresh(args: GatewayRefreshArgs) -> Result<()> { + GatewayRefresher::load(&args.work_dir)? + .refresh(args.force) + .await } struct AppIdValidator { diff --git a/os/common/rootfs/wg-checker.service b/os/common/rootfs/wg-checker.service index 406cadc69..a3e41f6af 100644 --- a/os/common/rootfs/wg-checker.service +++ b/os/common/rootfs/wg-checker.service @@ -1,13 +1,21 @@ [Unit] -Description=WireGuard Endpoint Checker Service +Description=dstack Gateway Registration Checker After=network-online.target dstack-prepare.service Wants=network-online.target [Service] Type=simple -ExecStart=/bin/wg-checker.sh -Restart=always +ExecStart=/bin/dstack-util gateway-checker --work-dir /dstack +# 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 diff --git a/os/common/rootfs/wg-checker.sh b/os/common/rootfs/wg-checker.sh deleted file mode 100755 index b81428187..000000000 --- a/os/common/rootfs/wg-checker.sh +++ /dev/null @@ -1,90 +0,0 @@ -#!/bin/sh - -# SPDX-FileCopyrightText: © 2024 Phala Network -# -# SPDX-License-Identifier: Apache-2.0 - -HANDSHAKE_TIMEOUT=180 -REFRESH_INTERVAL=180 -MISSING_CONFIG_RETRY_INTERVAL=30 -LAST_REFRESH=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 - STALE_SINCE=0 -} - -check_and_refresh() { - if ! command -v wg >/dev/null 2>&1; then - return - fi - - now=$(date +%s) - - # 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 - return - fi - - # Check handshake staleness (forced refresh) - latest=$(get_latest_handshake) - if [ -z "$latest" ]; then - latest=0 - fi - - if [ "$latest" -gt 0 ]; then - if [ $((now - latest)) -ge $HANDSHAKE_TIMEOUT ]; then - do_refresh "$now" "WireGuard handshake stale" 1 >&2 - fi - else - if [ "$STALE_SINCE" -eq 0 ]; then - STALE_SINCE=$now - fi - if [ $((now - STALE_SINCE)) -ge $HANDSHAKE_TIMEOUT ]; then - do_refresh "$now" "WireGuard handshake stale" 1 >&2 - fi - fi -} - -while true; do - if [ -f /etc/wireguard/$IFNAME.conf ]; then - check_and_refresh - else - STALE_SINCE=0 - now=$(date +%s) - if [ "$LAST_REFRESH" -eq 0 ] || [ $((now - LAST_REFRESH)) -ge $MISSING_CONFIG_RETRY_INTERVAL ]; then - do_refresh "$now" "WireGuard configuration missing" 1 - fi - 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 8546abd9a..a244af5e5 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/tests/acceptance.sh b/os/mkosi/tests/acceptance.sh index 21f4f8007..820c6c7e8 100755 --- a/os/mkosi/tests/acceptance.sh +++ b/os/mkosi/tests/acceptance.sh @@ -69,6 +69,24 @@ grep -q -- '--fuzz=0' "$D/components/kernel/kernel-build.sh" for service in dstack-guest-agent dstack-prepare app-compose wg-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. +wg_unit="$D/../common/rootfs/wg-checker.service" +grep -q '^ExecStart=/bin/dstack-util gateway-checker ' "$wg_unit" || { + echo 'wg-checker.service must run the dstack-util gateway-checker subcommand'; exit 1; } +grep -q '^Restart=on-failure$' "$wg_unit" || { + echo 'wg-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}\$" "$wg_unit" || { + echo "wg-checker.service must set RestartPreventExitStatus=$exit_code"; exit 1; } +test ! -e "$D/../common/rootfs/wg-checker.sh" # 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..37eb83dc5 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 @@ -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 From db83224d5301efc1ba8b408423bc66861d09ab22 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 04:54:04 -0700 Subject: [PATCH 3/7] fix(guest): report gateway outage to the host and rename the checker unit Boot no longer fails when gateway registration fails, so a guest log line was the only trace of it: the VM reported a clean boot while having no ingress route at all. Report `boot.error` to the host instead, which the VMM already surfaces via VmInfo and vmm-cli. The report is retracted once the route comes back. The checker seeds its notion of "already reported" from whether a WireGuard config exists at startup rather than assuming health, so a boot-time failure that the checker later recovers from does not leave a stale error on screen forever. Reporting only on transitions keeps the event ring buffer clean. The unit is renamed wg-checker.service -> dstack-gateway-checker.service now that it supervises gateway registration rather than poking WireGuard. The rootfs is read-only and rebuilt per image, so no upgrade path is needed; acceptance asserts the old names are gone. --- dstack/dstack-util/src/gateway_checker.rs | 29 +++++++++++++++---- dstack/dstack-util/src/system_setup.rs | 24 +++++++++++++++ ...service => dstack-gateway-checker.service} | 0 .../systemd/system-preset/80-dstack.preset | 2 +- os/mkosi/tests/acceptance.sh | 17 ++++++----- .../recipes-core/dstack-guest/dstack-guest.bb | 4 +-- 6 files changed, 60 insertions(+), 16 deletions(-) rename os/common/rootfs/{wg-checker.service => dstack-gateway-checker.service} (100%) diff --git a/dstack/dstack-util/src/gateway_checker.rs b/dstack/dstack-util/src/gateway_checker.rs index e72fb4422..d3f4f74a4 100644 --- a/dstack/dstack-util/src/gateway_checker.rs +++ b/dstack/dstack-util/src/gateway_checker.rs @@ -32,7 +32,9 @@ use anyhow::{Context, Result}; use cmd_lib::run_fun as cmd; use tracing::{error, info, warn}; -use crate::system_setup::{GatewayRefresher, WG_CONFIG_PATH, WG_INTERFACE}; +use crate::system_setup::{ + gateway_unavailable_message, GatewayRefresher, WG_CONFIG_PATH, WG_INTERFACE, +}; /// How often the loop samples the world. const POLL_INTERVAL: Duration = Duration::from_secs(10); @@ -46,7 +48,7 @@ const MISSING_CONFIG_RETRY_INTERVAL: i64 = 30; const MAX_RETRY_INTERVAL: i64 = 120; /// Exit code meaning "the gateway config is broken in a way retrying cannot -/// fix". Pinned by `RestartPreventExitStatus` in wg-checker.service, so +/// fix". Pinned by `RestartPreventExitStatus` in dstack-gateway-checker.service, so /// changing it requires changing the unit too. const EXIT_MISCONFIGURED: i32 = 3; @@ -260,6 +262,12 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { } info!("watching dstack-gateway registration"); + let vmm = refresher.host_api(); + // Seed from the observable world rather than assuming health: no WireGuard + // config here means boot-time registration failed and already reported it, + // so the first success owes the host a retraction. Assuming health instead + // would leave that boot error on screen forever after we recover. + let mut reported_degraded = !Path::new(WG_CONFIG_PATH).exists(); let mut checker = Checker::default(); loop { let now = now_secs(); @@ -268,16 +276,27 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { let succeeded = match refresher.refresh(refresh.force).await { Ok(()) => { info!("dstack-gateway refresh succeeded"); + if reported_degraded { + info!("dstack-gateway route restored; clearing the reported error"); + // Empty body resets the host's boot_error field. + vmm.notify_q("boot.error", "").await; + reported_degraded = false; + } true } Err(error) => { - // now_secs() is re-read below: the refresh itself can block - // on network timeouts for a long time, and the backoff must - // count from when the attempt ended. warn!("dstack-gateway refresh failed: {error:#}"); + if !reported_degraded { + vmm.notify_q("boot.error", &gateway_unavailable_message(&error)) + .await; + reported_degraded = true; + } 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(), succeeded); } tokio::time::sleep(POLL_INTERVAL).await; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index 35c64d01c..1fb29109b 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -1928,6 +1928,14 @@ impl Stage0<'_> { } } +/// The message reported to the host while this CVM has no gateway route. +/// +/// Boot and the gateway checker share it so the operator sees one consistent +/// string no matter which of the two noticed the outage. +pub fn gateway_unavailable_message(error: &anyhow::Error) -> String { + format!("dstack-gateway registration failed, the app has no ingress route: {error:#}") +} + /// Owns the inputs needed to (re)register this CVM with dstack-gateway. /// /// Loading is separated from refreshing so a long-running caller (the gateway @@ -1958,6 +1966,14 @@ impl GatewayRefresher { self.shared.app_compose.gateway_enabled() } + /// Client for reporting guest state back to the host. + pub fn host_api(&self) -> HostApi { + HostApi::new( + self.shared.sys_config.host_api_url.clone(), + self.shared.sys_config.collateral_urls().pccs, + ) + } + /// 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. @@ -2829,6 +2845,14 @@ impl Stage1<'_> { 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", &gateway_unavailable_message(&error)) + .await; } self.vmm .notify_q("boot.progress", "setting up docker") diff --git a/os/common/rootfs/wg-checker.service b/os/common/rootfs/dstack-gateway-checker.service similarity index 100% rename from os/common/rootfs/wg-checker.service rename to os/common/rootfs/dstack-gateway-checker.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 820c6c7e8..b638c6292 100755 --- a/os/mkosi/tests/acceptance.sh +++ b/os/mkosi/tests/acceptance.sh @@ -66,17 +66,17 @@ 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. -wg_unit="$D/../common/rootfs/wg-checker.service" -grep -q '^ExecStart=/bin/dstack-util gateway-checker ' "$wg_unit" || { - echo 'wg-checker.service must run the dstack-util gateway-checker subcommand'; exit 1; } -grep -q '^Restart=on-failure$' "$wg_unit" || { - echo 'wg-checker.service must use Restart=on-failure'; exit 1; } +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 @@ -84,9 +84,10 @@ grep -q '^Restart=on-failure$' "$wg_unit" || { 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}\$" "$wg_unit" || { - echo "wg-checker.service must set RestartPreventExitStatus=$exit_code"; exit 1; } +grep -q "^RestartPreventExitStatus=${exit_code}\$" "$gw_unit" || { + echo "dstack-gateway-checker.service must set RestartPreventExitStatus=$exit_code"; 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 37eb83dc5..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" @@ -85,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 From 034153710b903be7b0bd539306577fcd88f0a4ef Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 05:43:03 -0700 Subject: [PATCH 4/7] fix(guest): port the upstream forced-refresh fixes into the Rust checker Master fixed two bugs in wg-checker.sh after this branch was cut, and the Rust rewrite reintroduced both: 1. Staleness was checked after the periodic refresh instead of before it. 2. Every refresh cleared the missing-handshake timer. Either one alone makes the forced branch unreachable for a peer that never handshakes: HANDSHAKE_TIMEOUT equals REFRESH_INTERVAL, so the periodic refresh at last_attempt+180 reset a timer that had only reached 170. That matters because gateway setup short-circuits on an unchanged config unless force is set, so a CVM whose WireGuard config is correct but whose tunnel is dead could only ever recover through the forced path it could not reach. handshake_missing_since now means exactly one thing -- when we first observed a peer with no handshake -- and is cleared by an observed handshake and by nothing else. Staleness is evaluated first and returns either way, since a periodic refresh cannot rebuild a dead tunnel. Reaching the forced branch exposes the third thing master added and this branch lacked: nothing rate limited it, so a gateway that stays down would be bounced and re-certified on every 10s poll. last_force is tracked separately from last_attempt (a cheap periodic refresh must not re-arm the limit that guards the expensive forced one) and allows one forced attempt per HANDSHAKE_TIMEOUT. The tests now replay master's verification table over 900s of simulated time at the real poll interval, so a future refactor that breaks these semantics fails against the same numbers that commit was verified with: scenario forced healthy handshakes 0 never handshakes 4 handshake frozen, gw down 4 handshake then peer gone 3 no handshake, gw recovers 2, then periodic The old `gives_a_new_interface_a_full_timeout_to_handshake` asserted only that *some* refresh happened at t=189 and passed against the broken ordering, because the periodic refresh satisfied it. It now pins which refresh fires at which poll. --- dstack/dstack-util/src/gateway_checker.rs | 368 ++++++++++++++-------- 1 file changed, 245 insertions(+), 123 deletions(-) diff --git a/dstack/dstack-util/src/gateway_checker.rs b/dstack/dstack-util/src/gateway_checker.rs index d3f4f74a4..66cc81333 100644 --- a/dstack/dstack-util/src/gateway_checker.rs +++ b/dstack/dstack-util/src/gateway_checker.rs @@ -11,15 +11,18 @@ //! re-running the same registration whenever the observable state says it is //! needed. //! -//! There are exactly three reasons to refresh, in priority order: +//! 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. **Periodic re-registration.** The gateway expires idle registrations, so -//! re-register every [`REFRESH_INTERVAL`] even when everything looks fine. -//! 3. **Stale WireGuard handshake.** The tunnel exists but the peer stopped +//! 2. **Stale WireGuard handshake.** The tunnel exists but the peer stopped //! answering for [`HANDSHAKE_TIMEOUT`], which usually means the gateway -//! restarted and forgot us. +//! 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 @@ -121,45 +124,53 @@ struct Refresh { 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. + /// 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 { /// Decide whether to refresh now. Pure: same state plus same observation /// always yields the same answer. fn decide(&mut self, obs: Observation) -> Option { - // An attempt that has not cooled down yet blocks every reason below, so - // repeated failures cannot bypass the backoff by changing category. - let elapsed = match self.last_attempt { - // Nothing tried yet: act immediately. - None => return Some(Self::first_refresh(obs)), - Some(last) => obs.now.saturating_sub(last), - }; - if elapsed < self.backoff.delay() { - return None; - } - if !obs.config_present { + // No config means no interface, so there is no handshake to age. self.handshake_missing_since = None; - return (elapsed >= MISSING_CONFIG_RETRY_INTERVAL).then_some(Refresh { - force: true, - reason: "WireGuard config is missing", - }); - } - - if elapsed >= REFRESH_INTERVAL { - self.handshake_missing_since = None; - return Some(Refresh { - force: false, - reason: "periodic re-registration", - }); + 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); } - // The tunnel is configured and recently re-registered; the only thing - // left that can be wrong is the tunnel itself going quiet. + // 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; @@ -170,32 +181,35 @@ impl Checker { // HANDSHAKE_TIMEOUT to complete its first handshake. None => *self.handshake_missing_since.get_or_insert(obs.now), }; - (obs.now.saturating_sub(silent_since) >= HANDSHAKE_TIMEOUT).then_some(Refresh { - force: true, - reason: "WireGuard handshake is stale", - }) - } - - /// The first poll always acts: boot may have left the CVM unregistered, and - /// re-registering an already-healthy CVM is cheap and idempotent. - fn first_refresh(obs: Observation) -> Refresh { - if obs.config_present { - Refresh { - force: false, - reason: "initial re-registration", - } - } else { - Refresh { - force: true, - reason: "WireGuard config is missing", - } + 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, succeeded: bool) { + fn record(&mut self, now: i64, refresh: Refresh, succeeded: bool) { self.last_attempt = Some(now); - self.handshake_missing_since = None; + if refresh.force { + self.last_force = Some(now); + } + // handshake_missing_since is deliberately NOT cleared here; see decide(). self.backoff.record(succeeded); } } @@ -297,7 +311,7 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { // 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(), succeeded); + checker.record(now_secs(), refresh, succeeded); } tokio::time::sleep(POLL_INTERVAL).await; } @@ -317,6 +331,151 @@ mod tests { } } + /// 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. + 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"; @@ -333,12 +492,14 @@ mod tests { #[test] fn acts_immediately_on_first_poll() { let mut checker = Checker::default(); - let refresh = checker.decide(obs(0, false, None)).expect("should refresh"); - assert!(refresh.force); + assert_eq!(checker.decide(obs(0, false, None)), Some(MISSING_CONFIG)); let mut checker = Checker::default(); - let refresh = checker.decide(obs(0, true, None)).expect("should refresh"); - assert!(!refresh.force, "an existing config needs no forced reapply"); + assert_eq!( + checker.decide(obs(0, true, None)), + Some(PERIODIC), + "an existing config needs no forced reapply" + ); } #[test] @@ -346,23 +507,25 @@ mod tests { let mut checker = Checker::default(); let mut t = 0; // First attempt fires immediately and fails. - assert!(checker.decide(obs(t, false, None)).is_some()); - checker.record(T0 + t, false); + 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!( - checker.decide(obs(t + early, false, None)).is_none(), + 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!( - checker.decide(obs(t, false, None)).is_some(), + assert_eq!( + checker.decide(obs(t, false, None)), + Some(MISSING_CONFIG), "must retry once {expected_delay}s have passed" ); - checker.record(T0 + t, false); + checker.record(T0 + t, MISSING_CONFIG, false); } } @@ -389,72 +552,31 @@ mod tests { #[test] fn re_registers_periodically_while_healthy() { let mut checker = Checker::default(); - checker.record(T0, true); - - // Healthy tunnel, fresh handshake: quiet until REFRESH_INTERVAL. - assert!(checker.decide(obs(170, true, Some(T0 + 170))).is_none()); - let refresh = checker - .decide(obs(180, true, Some(T0 + 180))) - .expect("periodic refresh is due"); - assert!( - !refresh.force, - "periodic refresh must not disrupt the tunnel" - ); - } + checker.record(T0, PERIODIC, true); - #[test] - fn forces_refresh_when_the_handshake_goes_stale() { - let mut checker = Checker::default(); - checker.record(T0, true); - - // Handshake 179s old: still within tolerance. - assert!(checker.decide(obs(170, true, Some(T0 - 9))).is_none()); - // 180s old and we have not hit the periodic interval yet. - let refresh = checker - .decide(obs(170, true, Some(T0 - 10))) - .expect("stale handshake must force a refresh"); - assert!(refresh.force); - assert_eq!(refresh.reason, "WireGuard handshake is stale"); + 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(); - // Config appeared at t=0 but no handshake yet. - checker.record(T0, true); + checker.record(T0, PERIODIC, true); - assert!(checker.decide(obs(10, true, None)).is_none()); - assert!(checker.decide(obs(179, true, None)).is_none()); - // The timeout runs from t=10, when the missing handshake was first + // The timer runs from t=10, when the missing handshake was first // observed, not from the refresh at t=0. - assert!(checker.decide(obs(189, true, None)).is_some()); - } - - #[test] - fn a_recovered_handshake_clears_the_missing_timer() { - let mut checker = Checker::default(); - checker.record(T0, true); - - assert!(checker.decide(obs(10, true, None)).is_none()); - // Handshake completes. - assert!(checker.decide(obs(20, true, Some(T0 + 20))).is_none()); - assert_eq!(checker.handshake_missing_since, None); - // Losing it again restarts the clock instead of firing immediately. - assert!(checker.decide(obs(30, true, None)).is_none()); - assert_eq!(checker.handshake_missing_since, Some(T0 + 30)); - } - - #[test] - fn backoff_outranks_the_missing_config_interval() { - let mut checker = Checker::default(); - // Three consecutive failures put the backoff at the 120s cap, which is - // longer than the 30s missing-config interval. - checker.backoff.consecutive_failures = 3; - checker.record(T0, false); - assert_eq!(checker.backoff.consecutive_failures, 4); - - assert!(checker.decide(obs(30, false, None)).is_none()); - assert!(checker.decide(obs(119, false, None)).is_none()); - assert!(checker.decide(obs(120, false, None)).is_some()); + 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)); } } From d6df97e9af62ab28d0fcfbd6b59a5c95885ece1c Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 06:14:05 -0700 Subject: [PATCH 5/7] perf(guest): trim wasted work from the gateway checker loop Three findings from re-reading the poll loop end to end. 1. `observe()` forked `wg show` on every poll even with no WireGuard config, where `decide()` never looks at the handshake. That is exactly the state a gateway outage parks the CVM in, so the useless fork ran every 10s for the whole outage. Probe only when a config exists. 2. The checker re-registered immediately at startup, seconds after boot had just registered the same CVM. /etc is a volatile overlay, so a config on disk can only mean this boot wrote it; treat that as the start of the periodic clock instead. A fleet rebooting together no longer serves the gateway a second full round of registrations right after the first. With no config, boot's registration failed and fast recovery is the point, so the first poll still acts immediately. 3. `notify_q` swallows its error, so a host-API blip while reporting or retracting the gateway error flipped `reported_degraded` anyway and the state was never re-sent. A stale error could sit on the VMM for the life of the VM, or an outage could go unreported. Use `notify` and only move the flag once the host has actually taken the message. Considered and rejected: replacing the fixed 10s poll with a sleep until the next deadline. It would cut wakeups from ~8.6k/day to a few hundred, but the loop now costs one stat() per poll in the outage case, the CPU saving is on the order of 0.03%, and it would mean encoding the same deadlines in both `decide()` and a `next_wake()` that must never drift apart. Not worth the second source of truth. --- dstack/dstack-util/src/gateway_checker.rs | 110 ++++++++++++++++++---- 1 file changed, 94 insertions(+), 16 deletions(-) diff --git a/dstack/dstack-util/src/gateway_checker.rs b/dstack/dstack-util/src/gateway_checker.rs index 66cc81333..59be08b4d 100644 --- a/dstack/dstack-util/src/gateway_checker.rs +++ b/dstack/dstack-util/src/gateway_checker.rs @@ -149,6 +149,24 @@ const PERIODIC: Refresh = Refresh { }; 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 { @@ -241,11 +259,23 @@ fn parse_latest_handshake(output: &str) -> Option { .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: Path::new(WG_CONFIG_PATH).exists(), - latest_handshake: latest_handshake(WG_INTERFACE), + 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(), } } @@ -281,8 +311,9 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { // config here means boot-time registration failed and already reported it, // so the first success owes the host a retraction. Assuming health instead // would leave that boot error on screen forever after we recover. - let mut reported_degraded = !Path::new(WG_CONFIG_PATH).exists(); - let mut checker = Checker::default(); + let config_present = wg_config_present(); + let mut reported_degraded = !config_present; + let mut checker = Checker::starting(now_secs(), config_present); loop { let now = now_secs(); if let Some(refresh) = checker.decide(observe(now)) { @@ -291,19 +322,37 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { Ok(()) => { info!("dstack-gateway refresh succeeded"); if reported_degraded { - info!("dstack-gateway route restored; clearing the reported error"); - // Empty body resets the host's boot_error field. - vmm.notify_q("boot.error", "").await; - reported_degraded = false; + // Flip the flag only once the host has actually taken + // the retraction. notify_q swallows the error, so a + // host-API blip would strand a stale gateway error on + // the VMM for the life of the VM. Empty body resets + // the host's boot_error field. + match vmm.notify("boot.error", "").await { + Ok(()) => { + info!("dstack-gateway route restored; cleared the reported error"); + reported_degraded = false; + } + Err(error) => { + warn!("failed to retract the gateway error: {error:#}"); + } + } } true } Err(error) => { warn!("dstack-gateway refresh failed: {error:#}"); if !reported_degraded { - vmm.notify_q("boot.error", &gateway_unavailable_message(&error)) - .await; - reported_degraded = true; + // Same reasoning in reverse: a report the host never + // received must be retried, not marked as delivered. + match vmm + .notify("boot.error", &gateway_unavailable_message(&error)) + .await + { + Ok(()) => reported_degraded = true, + Err(error) => { + warn!("failed to report the gateway outage: {error:#}"); + } + } } false } @@ -337,6 +386,10 @@ mod tests { /// `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); @@ -490,15 +543,40 @@ mod tests { } #[test] - fn acts_immediately_on_first_poll() { - let mut checker = Checker::default(); + 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)); + } - let mut checker = Checker::default(); + /// 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(0, true, None)), + checker.decide(obs(180, true, Some(T0 + 180))), Some(PERIODIC), - "an existing config needs no forced reapply" + "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) ); } From 17d3368824223a3f1bf21c79848446da792cd095 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 07:32:11 -0700 Subject: [PATCH 6/7] revert(guest): keep gateway error reporting fire-and-forget Checking notify's result meant the loop had to carry a third state -- degraded, and the host may or may not know -- so the two flag flips grew into two nested matches and doubled the body of the refresh arm. That is a lot of machinery for a host-API blip that the next transition into or out of the degraded state reports again anyway. Back to notify_q, with the tradeoff written down where the flag is seeded. The observe() and startup-registration changes from the previous commit stand. --- dstack/dstack-util/src/gateway_checker.rs | 37 ++++++++--------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/dstack/dstack-util/src/gateway_checker.rs b/dstack/dstack-util/src/gateway_checker.rs index 59be08b4d..6737dcdd4 100644 --- a/dstack/dstack-util/src/gateway_checker.rs +++ b/dstack/dstack-util/src/gateway_checker.rs @@ -311,6 +311,11 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { // config here means boot-time registration failed and already reported it, // so the first success owes the host a retraction. Assuming health instead // would leave that boot error on screen forever after we recover. + // + // Reporting is deliberately fire-and-forget. Tracking delivery would mean + // carrying a third state ("degraded, and the host may or may not know") + // through the loop to cover a host-API blip that the next refresh cycle + // already re-reports on the way in or out of the degraded state. let config_present = wg_config_present(); let mut reported_degraded = !config_present; let mut checker = Checker::starting(now_secs(), config_present); @@ -322,37 +327,19 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { Ok(()) => { info!("dstack-gateway refresh succeeded"); if reported_degraded { - // Flip the flag only once the host has actually taken - // the retraction. notify_q swallows the error, so a - // host-API blip would strand a stale gateway error on - // the VMM for the life of the VM. Empty body resets - // the host's boot_error field. - match vmm.notify("boot.error", "").await { - Ok(()) => { - info!("dstack-gateway route restored; cleared the reported error"); - reported_degraded = false; - } - Err(error) => { - warn!("failed to retract the gateway error: {error:#}"); - } - } + info!("dstack-gateway route restored; clearing the reported error"); + // Empty body resets the host's boot_error field. + vmm.notify_q("boot.error", "").await; + reported_degraded = false; } true } Err(error) => { warn!("dstack-gateway refresh failed: {error:#}"); if !reported_degraded { - // Same reasoning in reverse: a report the host never - // received must be retried, not marked as delivered. - match vmm - .notify("boot.error", &gateway_unavailable_message(&error)) - .await - { - Ok(()) => reported_degraded = true, - Err(error) => { - warn!("failed to report the gateway outage: {error:#}"); - } - } + vmm.notify_q("boot.error", &gateway_unavailable_message(&error)) + .await; + reported_degraded = true; } false } From b5b111f8059c4e1876ca00d705619c9aa2b9b655 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 5 Aug 2026 07:42:15 -0700 Subject: [PATCH 7/7] fix(guest): make the gateway checker loop self-recovering Two changes, one dropping state and one closing the last failure mode that had no automatic recovery. Drop reported_degraded. Mirroring gateway state to the host meant the loop had to remember what the host had been told, and that memory bought nothing the loop needs: recovery does not depend on it. Boot still reports the one signal that matters -- this CVM came up without a route. The cost is that the reported error stays on the VMM after the checker recovers, until the VM restarts, which is the intended trade. Add a systemd watchdog. Every recovery path here only runs while the loop runs, and nothing below the loop could guarantee that: a refresh that never returns leaves a healthy-looking process that Restart= would never fire on, and the CVM sits without a route indefinitely. That is not hypothetical -- a refresh spends most of its time in blocking cmd! shell-outs, and wg-quick resolves peer endpoints, so a hung DNS lookup is enough. tokio::time::timeout cannot fix this: those calls block a runtime worker with no await point to cancel at. Neither can a watchdog task on another worker, which would keep pinging while this loop is wedged. The ping has to come from the loop itself, before the work rather than after, so that a refresh which never returns stops the pings. systemd then kills and restarts the service, covering a hang whatever its cause. Follows the Type=notify + WatchdogSec + sd_notify pattern already used by dstack-guest-agent. Readiness is reported before the checker decides whether it has anything to do, so the exit-0 (gateway disabled) and exit-3 (misconfigured) paths remain a started service that then stopped. WatchdogSec=600 clears the worst legitimate refresh, which can span KMS certificate requests plus every configured gateway URL at 60s each. Verified end to end against a stub notify socket: the process emits READY=1 at startup and WATCHDOG=1 on each iteration, and stays inert when the unit does not arm the watchdog, so manual runs still work. --- dstack/Cargo.lock | 1 + dstack/dstack-util/Cargo.toml | 1 + dstack/dstack-util/src/gateway_checker.rs | 83 +++++++++++++------ dstack/dstack-util/src/system_setup.rs | 23 ++--- .../rootfs/dstack-gateway-checker.service | 10 ++- os/mkosi/tests/acceptance.sh | 6 ++ 6 files changed, 79 insertions(+), 45 deletions(-) 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 index 6737dcdd4..e373c46d6 100644 --- a/dstack/dstack-util/src/gateway_checker.rs +++ b/dstack/dstack-util/src/gateway_checker.rs @@ -33,11 +33,10 @@ 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::{ - gateway_unavailable_message, GatewayRefresher, WG_CONFIG_PATH, WG_INTERFACE, -}; +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); @@ -279,7 +278,45 @@ fn observe(now: i64) -> Observation { } } +/// 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")?; @@ -306,41 +343,33 @@ pub async fn cmd_gateway_checker(args: GatewayCheckerArgs) -> Result<()> { } info!("watching dstack-gateway registration"); - let vmm = refresher.host_api(); - // Seed from the observable world rather than assuming health: no WireGuard - // config here means boot-time registration failed and already reported it, - // so the first success owes the host a retraction. Assuming health instead - // would leave that boot error on screen forever after we recover. - // - // Reporting is deliberately fire-and-forget. Tracking delivery would mean - // carrying a third state ("degraded, and the host may or may not know") - // through the loop to cover a host-API blip that the next refresh cycle - // already re-reports on the way in or out of the degraded state. - let config_present = wg_config_present(); - let mut reported_degraded = !config_present; - let mut checker = Checker::starting(now_secs(), config_present); + // 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"); - if reported_degraded { - info!("dstack-gateway route restored; clearing the reported error"); - // Empty body resets the host's boot_error field. - vmm.notify_q("boot.error", "").await; - reported_degraded = false; - } true } Err(error) => { warn!("dstack-gateway refresh failed: {error:#}"); - if !reported_degraded { - vmm.notify_q("boot.error", &gateway_unavailable_message(&error)) - .await; - reported_degraded = true; - } false } }; diff --git a/dstack/dstack-util/src/system_setup.rs b/dstack/dstack-util/src/system_setup.rs index d0e53129c..c593976c9 100644 --- a/dstack/dstack-util/src/system_setup.rs +++ b/dstack/dstack-util/src/system_setup.rs @@ -1936,14 +1936,6 @@ impl Stage0<'_> { } } -/// The message reported to the host while this CVM has no gateway route. -/// -/// Boot and the gateway checker share it so the operator sees one consistent -/// string no matter which of the two noticed the outage. -pub fn gateway_unavailable_message(error: &anyhow::Error) -> String { - format!("dstack-gateway registration failed, the app has no ingress route: {error:#}") -} - /// Owns the inputs needed to (re)register this CVM with dstack-gateway. /// /// Loading is separated from refreshing so a long-running caller (the gateway @@ -1974,14 +1966,6 @@ impl GatewayRefresher { self.shared.app_compose.gateway_enabled() } - /// Client for reporting guest state back to the host. - pub fn host_api(&self) -> HostApi { - HostApi::new( - self.shared.sys_config.host_api_url.clone(), - self.shared.sys_config.collateral_urls().pccs, - ) - } - /// 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. @@ -2923,7 +2907,12 @@ impl Stage1<'_> { // visible from the VMM. The gateway checker clears this once it // manages to register. self.vmm - .notify_q("boot.error", &gateway_unavailable_message(&error)) + .notify_q( + "boot.error", + &format!( + "dstack-gateway registration failed, the app has no ingress route: {error:#}" + ), + ) .await; } self.vmm diff --git a/os/common/rootfs/dstack-gateway-checker.service b/os/common/rootfs/dstack-gateway-checker.service index a3e41f6af..4f183eb8c 100644 --- a/os/common/rootfs/dstack-gateway-checker.service +++ b/os/common/rootfs/dstack-gateway-checker.service @@ -4,8 +4,16 @@ After=network-online.target dstack-prepare.service Wants=network-online.target [Service] -Type=simple +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 diff --git a/os/mkosi/tests/acceptance.sh b/os/mkosi/tests/acceptance.sh index b638c6292..5d26ea08b 100755 --- a/os/mkosi/tests/acceptance.sh +++ b/os/mkosi/tests/acceptance.sh @@ -86,6 +86,12 @@ exit_code=$(sed -n 's/^const EXIT_MISCONFIGURED: i32 = \([0-9]\+\);$/\1/p' "$che [[ -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