fix(gateway): every connection forks a process when wg show fails, and Admin.Status writes to the KV store under the routing lock

dstack/gateway/src/admin_service.rs

@@ -47,46 +47,60 @@ pub struct AdminRpcHandler {
 
 impl AdminRpcHandler {
     pub(crate) async fn status(self) -> Result<StatusResponse> {
+        tokio::task::spawn_blocking(move || {
+            self.state.refresh_state()?;
+            self.status_snapshot()
+        })
+        .await
+        .context("status task failed")?
+    }
+
+    fn status_snapshot(self) -> Result<StatusResponse> {
         let (base_domain, _port) = self
             .state
             .kv_store()
             .get_best_zt_domain()
             .unwrap_or_default();
-        let mut state = self.state.lock();
-        state.refresh_state()?;
-        let hosts = state
-            .state
-            .instances
-            .values()
-            .map(|instance| {
-                // Get global latest_handshake from KvStore (max across all nodes)
-                let latest_handshake = state
-                    .get_instance_latest_handshake(&instance.id)
-                    .unwrap_or(0);
-                HostInfo {
+        // KV reads can wait behind a sync merge, so do them after releasing
+        // the routing lock.
+        let mut hosts = {
+            let state = self.state.lock();
+            state
+                .state
+                .instances
+                .values()
+                .map(|instance| HostInfo {
                     instance_id: instance.id.clone(),
                     ip: instance.ip.to_string(),
                     app_id: instance.app_id.clone(),
                     base_domain: base_domain.clone(),
-                    latest_handshake,
+                    latest_handshake: 0,
                     num_connections: instance.num_connections(),
                     ready: Some(instance.is_ready()),
                     health: instance.health().as_str().to_string(),
-                }
-            })
-            .collect::<Vec<_>>();
+                })
+                .collect::<Vec<_>>()
+        };
+        for host in &mut hosts {
+            host.latest_handshake = self
+                .state
+                .kv_store()
+                .get_instance_latest_handshake(&host.instance_id)
+                .unwrap_or(0);
+        }
+        let config = &self.state.config;
         Ok(StatusResponse {
-            id: state.config.sync.node_id,
-            url: state.config.sync.my_url.clone(),
-            uuid: state.config.uuid(),
-            bootnode_url: state.config.sync.bootnode.clone(),
-            nodes: state.get_all_nodes(),
+            id: config.sync.node_id,
+            url: config.sync.my_url.clone(),
+            uuid: config.uuid(),
+            bootnode_url: config.sync.bootnode.clone(),
+            nodes: self.state.get_all_nodes(),
             hosts,
             num_connections: NUM_CONNECTIONS.load(Ordering::Relaxed),
             // Reads the post-probe config, so this is what the data path is
             // running rather than what the file asked for.
-            accel: Some(accel_status(&state.config.proxy)),
-            health_gating: state.config.proxy.health_check.enabled,
+            accel: Some(accel_status(&config.proxy)),
+            health_gating: config.proxy.health_check.enabled,
         })
     }
 }

dstack/gateway/src/main_service.rs

@@ -304,6 +304,47 @@ impl ProxyInner {
         self.handshake_cache.latest(stale_timeout)
     }
 
+    pub(crate) fn get_all_nodes(&self) -> Vec<GatewayNodeInfo> {
+        gateway_nodes(&self.kv_store, false)
+    }
+
+    /// Publish this node's WireGuard observations to the KV store.
+    ///
+    /// The KV writes wait on the store's write lock, which a sync merge holds,
+    /// so they must not run under the routing lock.
+    pub(crate) fn refresh_state(&self) -> Result<()> {
+        let handshakes = self.latest_handshakes(None)?;
+        let instance_ids: Vec<(String, u64)> = {
+            let state = self.lock();
+            state
+                .state
+                .instances
+                .iter()
+                .filter_map(|(id, info)| {
+                    let (timestamp, _) = handshakes.get(&info.public_key)?;
+                    Some((id.clone(), *timestamp))
+                })
+                .collect()
+        };
+
+        for (instance_id, timestamp) in &instance_ids {
+            if let Err(err) = self
+                .kv_store
+                .sync_instance_handshake(instance_id, *timestamp)
+            {
+                debug!("failed to sync instance handshake: {err:?}");
+            }
+        }
+
+        if let Err(err) = self
+            .kv_store
+            .sync_node_last_seen(self.config.sync.node_id, now_secs())
+        {
+            debug!("failed to sync node last_seen: {err:?}");
+        }
+        Ok(())
+    }
+
     pub async fn new(
         options: ProxyOptions,
         port_policy_tx: UnboundedSender<String>,
@@ -963,6 +1004,9 @@ fn start_recycle_thread(proxy: Proxy) {
     }
     std::thread::spawn(move || loop {
         std::thread::sleep(proxy.config.recycle.interval);
+        if let Err(err) = proxy.refresh_state() {
+            warn!("failed to refresh state: {err:?}");
+        }
         if let Err(err) = proxy.lock().recycle() {
             error!("failed to run recycle: {err:?}");
         };
@@ -2514,12 +2558,9 @@ impl ProxyState {
         Ok(())
     }
 
+    /// Drop instances the cluster has stopped seeing. The caller runs
+    /// `ProxyInner::refresh_state` first, off this lock.
     fn recycle(&mut self) -> Result<()> {
-        // Refresh state: sync local handshakes to KvStore, update local last_seen from global
-        if let Err(err) = self.refresh_state() {
-            warn!("failed to refresh state: {err:?}");
-        }
-
         // Note: Gateway nodes are not removed from KvStore, only marked offline/retired
 
         // Recycle stale CVM instances based on global last_seen (max across all nodes)
@@ -2581,87 +2622,49 @@ impl ProxyState {
         Ok(())
     }
 
-    pub(crate) fn refresh_state(&mut self) -> Result<()> {
-        // Get local WG handshakes and sync to KvStore
-        let handshakes = self.latest_handshakes(None)?;
-
-        // Build a map from public_key to instance_id for lookup
-        let pk_to_id: BTreeMap<&str, &str> = self
-            .state
-            .instances
-            .iter()
-            .map(|(id, info)| (info.public_key.as_str(), id.as_str()))
-            .collect();
-
-        // Sync local handshake observations to KvStore
-        for (pk, (ts, _)) in &handshakes {
-            if let Some(&instance_id) = pk_to_id.get(pk.as_str()) {
-                if let Err(err) = self.kv_store.sync_instance_handshake(instance_id, *ts) {
-                    debug!("failed to sync instance handshake: {err:?}");
-                }
-            }
-        }
-
-        // Update this node's last_seen in KvStore
-        let now = now_secs();
-        if let Err(err) = self
-            .kv_store
-            .sync_node_last_seen(self.config.sync.node_id, now)
-        {
-            debug!("failed to sync node last_seen: {err:?}");
-        }
-        Ok(())
-    }
-
     /// Sync connection count for an instance to KvStore
     pub(crate) fn sync_connections(&self, instance_id: &str, count: u64) {
         if let Err(err) = self.kv_store.sync_connections(instance_id, count) {
             debug!("Failed to sync connections: {err:?}");
         }
     }
 
-    /// Get latest handshake for an instance from KvStore (max across all nodes)
-    pub(crate) fn get_instance_latest_handshake(&self, instance_id: &str) -> Option<u64> {
-        self.kv_store.get_instance_latest_handshake(instance_id)
-    }
-
-    /// Get all nodes from KvStore (for admin API - includes all nodes)
-    pub(crate) fn get_all_nodes(&self) -> Vec<GatewayNodeInfo> {
-        self.get_all_nodes_filtered(false)
-    }
-
     /// Get nodes for CVM registration (excludes nodes with status "down")
     pub(crate) fn get_active_nodes(&self) -> Vec<GatewayNodeInfo> {
         self.get_all_nodes_filtered(true)
     }
 
     /// Get all nodes from KvStore with optional filtering
     fn get_all_nodes_filtered(&self, exclude_down: bool) -> Vec<GatewayNodeInfo> {
-        let node_statuses = if exclude_down {
-            self.kv_store.load_all_node_statuses()
-        } else {
-            Default::default()
-        };
-
-        self.kv_store
-            .load_all_nodes()
-            .into_iter()
-            // Shared with the metrics sampler so the gauge and the routing
-            // table cannot disagree about what "active" means.
-            .filter(|(id, _)| !exclude_down || KvStore::node_is_active(node_statuses.get(id)))
-            .map(|(id, node)| GatewayNodeInfo {
-                id,
-                uuid: node.uuid,
-                wg_public_key: node.wg_public_key,
-                wg_ip: node.wg_ip,
-                wg_endpoint: node.wg_endpoint,
-                url: node.url,
-                last_seen: self.kv_store.get_node_latest_last_seen(id).unwrap_or(0),
-            })
-            .collect()
+        gateway_nodes(&self.kv_store, exclude_down)
     }
 }
 
+fn gateway_nodes(kv_store: &KvStore, exclude_down: bool) -> Vec<GatewayNodeInfo> {
+    let node_statuses = if exclude_down {
+        kv_store.load_all_node_statuses()
+    } else {
+        Default::default()
+    };
+
+    kv_store
+        .load_all_nodes()
+        .into_iter()
+        // Shared with the metrics sampler so the gauge and the routing
+        // table cannot disagree about what "active" means.
+        .filter(|(id, _)| !exclude_down || KvStore::node_is_active(node_statuses.get(id)))
+        .map(|(id, node)| GatewayNodeInfo {
+            id,
+            uuid: node.uuid,
+            wg_public_key: node.wg_public_key,
+            wg_ip: node.wg_ip,
+            wg_endpoint: node.wg_endpoint,
+            url: node.url,
+            last_seen: kv_store.get_node_latest_last_seen(id).unwrap_or(0),
+        })
+        .collect()
+}
+
 pub struct RpcHandler {
     remote_app_id: Option<Vec<u8>>,
     remote_app_info: Option<AppInfo>,

dstack/gateway/src/main_service/handshakes.rs

@@ -51,37 +51,13 @@ impl LatestHandshakesCache {
         self.cell.set(timestamps);
     }
 
+    /// The cached snapshot, stale if need be; never runs `wg show`. Callers
+    /// hold the routing lock, and a host where `wg show` always fails would
+    /// otherwise fork once per connection.
     pub(crate) fn latest(&self, stale_timeout: Option<Duration>) -> Result<HandshakesWithAge> {
-        // Admin/public status paths call this synchronously. On fixture hosts the
-        // first successful `wg show` may not have completed yet (or the interface
-        // may be absent), so a hard Empty error collapses many Admin.* RPCs with
-        // "cached cell is empty". Prefer:
-        // 1) fresh TTL value
-        // 2) stale last-known value
-        // 3) one synchronous producer refresh
-        // 4) empty map so callers can still report registered hosts/meta
-        let timestamps = match self.cell.get() {
+        let timestamps = match self.cell.get_allow_stale() {
             Ok(snapshot) => snapshot.into_value(),
-            Err(cached_cell::GetError::Expired { .. }) | Err(cached_cell::GetError::Empty) => {
-                match self.cell.get_allow_stale() {
-                    Ok(snapshot) => snapshot.into_value(),
-                    Err(_) => {
-                        let interface = self.interface.clone();
-                        match fetch_latest_handshake_timestamps(&interface) {
-                            Ok(value) => {
-                                self.cell.set(value.clone());
-                                std::sync::Arc::new(value)
-                            }
-                            Err(err) => {
-                                warn!(
-                                    "WireGuard latest-handshakes unavailable; returning empty map: {err}"
-                                );
-                                std::sync::Arc::new(BTreeMap::new())
-                            }
-                        }
-                    }
-                }
-            }
+            Err(_) => Arc::new(BTreeMap::new()),
         };
         add_elapsed_time(timestamps.as_ref(), stale_timeout)
     }
@@ -155,6 +131,33 @@ fn add_elapsed_time(
 mod tests {
     use super::*;
 
+    #[test]
+    fn a_cold_cache_does_not_shell_out_on_the_routing_path() {
+        const CONNECTIONS: usize = 500;
+        let cache = LatestHandshakesCache::new(
+            "dstack-no-such-iface0".to_string(),
+            Duration::from_secs(30),
+        );
+
+        let started = std::time::Instant::now();
+        for _ in 0..CONNECTIONS {
+            assert!(
+                cache
+                    .latest(None)
+                    .expect("a cold cache still answers")
+                    .is_empty(),
+                "a host with no WireGuard data knows of no fresh handshake"
+            );
+        }
+        let elapsed = started.elapsed();
+
+        assert!(
+            elapsed < Duration::from_millis(100),
+            "{CONNECTIONS} cold reads took {elapsed:?}: the routing path is spawning \
+             a process per connection"
+        );
+    }
+
     #[test]
     fn parses_latest_handshake_timestamps() {
         let handshakes = parse_latest_handshake_timestamps(

dstack/gateway/src/main_service/tests.rs

@@ -2679,3 +2679,32 @@ fn tombstone_collection_triggers_on_write_count_boundaries_not_on_time() {
     // until it crosses the next boundary again.
     assert!(!tombstone_collection_due(Some(w(100, 0)), w(90, 0), 100));
 }
+
+/// A sync merge holding the KV store must not stall routing via `refresh_state`.
+#[tokio::test]
+async fn refreshing_state_leaves_the_routing_lock_free_while_it_writes_to_the_kv_store() {
+    use std::sync::atomic::AtomicBool;
+
+    let state = create_test_state().await;
+
+    // Stand in for a sync merge holding the store.
+    let _store_held = state.kv_store().ephemeral().write();
+
+    let finished = std::sync::Arc::new(AtomicBool::new(false));
+    let proxy = state.proxy.clone();
+    let done = finished.clone();
+    std::thread::spawn(move || {
+        let _ = proxy.refresh_state();
+        done.store(true, Ordering::SeqCst);
+    });
+    std::thread::sleep(std::time::Duration::from_millis(200));
+
+    assert!(
+        !finished.load(Ordering::SeqCst),
+        "test is vacuous: the refresh never reached a KV write"
+    );
+    assert!(
+        state.proxy.state.try_lock().is_ok(),
+        "the routing lock is held while the refresh waits on the KV store"
+    );
+}