Merge pull request #948 from Dstack-TEE/codex/fix-guest-gateway-outage-boot

fix(guest): decouple Gateway outage from app boot

Author: kvinwang

dstack/Cargo.lock

@@ -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

dstack/dstack-util/Cargo.toml

@@ -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?;
         }

dstack/dstack-util/src/gateway_checker.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<Self> {
+        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;

dstack/dstack-util/src/main.rs

@@ -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

dstack/dstack-util/src/system_setup.rs

@@ -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 \

os/common/rootfs/dstack-gateway-checker.service

@@ -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