From e7fc018e54aacbd92b709f10a6ea7a9b0010eea1 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 12 Aug 2026 03:14:49 -0700 Subject: [PATCH 1/3] feat(gateway): add rejected-record listing and node removal recovery APIs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three operator recovery additions in the same family as Admin.RemoveCvm: - Admin.ListRejectedInstances reports the instance records this node currently refuses to import, with the reason and whether the instance still holds local data-plane state. Rejections were previously logged only on transitions, so finding what to remove meant grepping old logs or restarting the gateway to re-log them. - Admin.RemoveNode tombstones a decommissioned gateway node's replicated records (info, status, sync address) and drops it from the sync peer set immediately. The __peer_addr tombstone doubles as the cluster-wide removal signal: every gateway watches the prefix and prunes its own peer set when the deletion replicates, so no restart is needed. An address that was never written does not count as removed, because bootstrap can add a peer before its address record has synced in. - Admin.DeleteZtDomain now works on a corrupt config record. It used to gate deletion on get_zt_domain_config, which cannot tell missing from unreadable, so a corrupt record was permanently stuck — the same trap RemoveCvm was added to fix for instance records. Both removal RPCs are idempotent and report record_existed so a mistyped ID is visible instead of silently succeeding. --- dstack/gateway/rpc/proto/gateway_rpc.proto | 43 ++++++++ dstack/gateway/src/admin_service.rs | 52 +++++++-- dstack/gateway/src/kv/mod.rs | 81 ++++++++++++++ dstack/gateway/src/main_service.rs | 80 +++++++++++++- dstack/gateway/src/main_service/tests.rs | 117 +++++++++++++++++++++ 5 files changed, 364 insertions(+), 9 deletions(-) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index d410aa451..7ea0bcb22 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -437,6 +437,14 @@ service Admin { // Remove a CVM from WaveKV and the local data plane. This is an idempotent // operator recovery action and also works when the stored record is unreadable. rpc RemoveCvm(RemoveCvmRequest) returns (RemoveCvmResponse) {} + // List the instance records this node currently refuses to import, with the + // reason for each. Pairs with RemoveCvm: list the bad records, then remove + // the ones that should not survive. + rpc ListRejectedInstances(google.protobuf.Empty) returns (ListRejectedInstancesResponse) {} + // Remove a decommissioned gateway node from WaveKV and this node's sync peer + // set. Idempotent operator recovery action; other gateways prune the node + // from their own peer sets when the removal replicates to them. + rpc RemoveNode(RemoveNodeRequest) returns (RemoveNodeResponse) {} // ==================== DNS Credential Management ==================== // List all DNS credentials @@ -517,6 +525,41 @@ message RemoveCvmResponse { bool removed_locally = 2; } +// One instance record this node refuses to import. +message RejectedInstanceInfo { + string instance_id = 1; + // Why the record is refused. + string reason = 2; + // "unusable": the record fails validation or its bytes no longer decode. + // "lost_conflict": the record lost an IP or key conflict to an older + // registration. + string rejection = 3; + // Whether the instance still holds state in this node's data plane. An + // unusable record keeps whatever the data plane already had, so removing + // an active instance also drops its routing. + bool active_locally = 4; +} + +message ListRejectedInstancesResponse { + repeated RejectedInstanceInfo rejected = 1; +} + +// Emergency operator request to remove a decommissioned gateway node. +message RemoveNodeRequest { + uint32 node_id = 1; +} + +// Outcome of RemoveNode. Both fields are false when the node was never known +// (or the removal already completed), so a mistyped node_id is visible to +// the operator instead of silently reporting success. +message RemoveNodeResponse { + // Whether a live node record existed in WaveKV before the tombstone was + // written. + bool record_existed = 1; + // Whether the node was still in this gateway's sync peer set. + bool removed_from_peer_set = 2; +} + // ==================== DNS Credential Messages ==================== // DNS credential information diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index d5c1f6710..0596e5b18 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -15,9 +15,10 @@ use dstack_gateway_rpc::{ GetInstanceHandshakesResponse, GetInstancePortPolicyRequest, GetInstancePortPolicyResponse, GetMetaResponse, GetNodeStatusesResponse, GetZtDomainRequest, GlobalConnectionsStats, HandshakeEntry, HostInfo, LastSeenEntry, ListCertAttestationsRequest, - ListCertAttestationsResponse, ListDnsCredentialsResponse, ListZtDomainsResponse, - NodeStatusEntry, PeerSyncStatus as ProtoPeerSyncStatus, PortAttrs as RpcPortAttrs, - PortPolicy as RpcPortPolicy, RemoveCvmRequest, RemoveCvmResponse, RenewCertResponse, + ListCertAttestationsResponse, ListDnsCredentialsResponse, ListRejectedInstancesResponse, + ListZtDomainsResponse, NodeStatusEntry, PeerSyncStatus as ProtoPeerSyncStatus, + PortAttrs as RpcPortAttrs, PortPolicy as RpcPortPolicy, RejectedInstanceInfo, RemoveCvmRequest, + RemoveCvmResponse, RemoveNodeRequest, RemoveNodeResponse, RenewCertResponse, RenewZtDomainCertRequest, RenewZtDomainCertResponse, RotateAcmeCredentialsResponse, SetCertbotConfigRequest, SetDefaultDnsCredentialRequest, SetInstancePortPolicyRequest, SetNodeStatusRequest, SetNodeUrlRequest, StatusResponse, StoreSyncStatus, @@ -30,8 +31,8 @@ use wavekv::node::NodeStatus as WaveKvNodeStatus; use crate::{ kv::{ - DnsCredential, DnsProvider, GlobalCertbotConfig, NodeStatus, PortFlags, PortPolicy, - ZtDomainConfig, + import::Rejection, DnsCredential, DnsProvider, GlobalCertbotConfig, NodeStatus, PortFlags, + PortPolicy, ZtDomainConfig, }, main_service::Proxy, models::PortPolicyView, @@ -325,6 +326,37 @@ impl AdminRpc for AdminRpcHandler { }) } + async fn list_rejected_instances(self) -> Result { + let rejected = self + .state + .rejected_instances() + .into_iter() + .map(|report| RejectedInstanceInfo { + instance_id: report.rejected.instance_id, + reason: format!("{:#}", report.rejected.reason), + rejection: match report.rejected.rejection { + Rejection::Unusable => "unusable".to_string(), + Rejection::LostConflict => "lost_conflict".to_string(), + }, + active_locally: report.active_locally, + }) + .collect(); + Ok(ListRejectedInstancesResponse { rejected }) + } + + async fn remove_node(self, request: RemoveNodeRequest) -> Result { + let removal = self.state.remove_node(request.node_id)?; + warn!( + "admin removed node {} from WaveKV and the sync peer set \ + (record existed: {}, was a sync peer: {})", + request.node_id, removal.record_existed, removal.removed_from_peer_set + ); + Ok(RemoveNodeResponse { + record_existed: removal.record_existed, + removed_from_peer_set: removal.removed_from_peer_set, + }) + } + // ==================== DNS Credential Management ==================== async fn list_dns_credentials(self) -> Result { @@ -556,9 +588,13 @@ impl AdminRpc for AdminRpcHandler { let kv_store = self.state.kv_store(); let domain = normalize_zt_domain(&request.domain)?; - kv_store - .get_zt_domain_config(&domain) - .context("ZT-Domain config not found")?; + // A corrupt config must still be deletable, so check for the record + // itself: get_zt_domain_config cannot tell missing from unreadable, + // and refusing would leave a corrupt record permanently stuck. + ensure!( + kv_store.zt_domain_config_exists(&domain), + "ZT-Domain config not found" + ); // Delete config (cert data, acme, attestations are kept for historical purposes) kv_store.delete_zt_domain_config(&domain)?; diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 1c6ce5705..6fed6e0b1 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -787,6 +787,26 @@ impl KvStore { .collect() } + /// Remove a gateway node from replicated state. + /// + /// Writes tombstones for the node's info, status, and sync address, and + /// drops this node's own last_seen observation of it. The `__peer_addr` + /// tombstone doubles as the cluster-wide removal signal: every gateway + /// prunes its sync peer set when it observes the deletion (see + /// [`Self::prune_removed_peers`]). + /// + /// Returns whether a live node record existed before the tombstone was + /// written. + pub fn sync_remove_node(&self, node_id: NodeId) -> Result { + let previous = self.persistent.write().delete(keys::node_info(node_id))?; + self.persistent.write().delete(keys::node_status(node_id))?; + self.persistent.write().delete(keys::peer_addr(node_id))?; + self.ephemeral + .write() + .delete(keys::last_seen_node(node_id, self.my_node_id))?; + Ok(previous.is_some_and(|entry| !entry.is_deleted())) + } + // ==================== Node Status Sync ==================== /// Set node status (stored separately from NodeData) @@ -946,6 +966,11 @@ impl KvStore { self.persistent.watch_prefix(keys::NODE_PREFIX) } + /// Watch for changes to replicated peer sync addresses + pub fn watch_peer_addrs(&self) -> watch::Receiver<()> { + self.persistent.watch_prefix(keys::PEER_ADDR_PREFIX) + } + // ==================== Persistence ==================== pub fn persist_if_dirty(&self) -> Result { @@ -960,6 +985,49 @@ impl KvStore { Ok(()) } + /// Drop a node from the sync peer set of both stores. + /// + /// Returns whether the persistent store still had it as a peer. + pub fn remove_peer(&self, peer_id: NodeId) -> Result { + let removed = self.persistent.write().remove_peer(peer_id)?; + self.ephemeral.write().remove_peer(peer_id)?; + Ok(removed) + } + + /// Drop peers whose sync address has been explicitly deleted. + /// + /// A tombstoned `__peer_addr/{id}` record is the replicated signal that + /// an operator removed the node (see [`Self::sync_remove_node`]). An + /// address that was never written does not count: bootstrap can add a + /// peer before its address record has synced in, and such a peer must + /// not be dropped for being early. + pub fn prune_removed_peers(&self) { + let peer_ids: Vec = self + .persistent + .read() + .status() + .peers + .iter() + .map(|peer| peer.id) + .collect(); + for peer_id in peer_ids { + // `get` filters tombstones out, so the deletion signal is only + // visible through the tombstone-inclusive accessor. + let tombstoned = self + .persistent + .read() + .get_including_tombstones(&keys::peer_addr(peer_id)) + .is_some_and(|entry| entry.is_deleted()); + if !tombstoned { + continue; + } + warn!("dropping removed node {peer_id} from the sync peer set"); + if let Err(err) = self.remove_peer(peer_id) { + warn!("failed to remove peer {peer_id}: {err:#}"); + } + } + } + // ==================== Peer Address (in DB) ==================== /// Register a node's sync URL in DB and add to peer list for sync @@ -1100,6 +1168,19 @@ impl KvStore { .decode(&keys::zt_domain_config(domain)) } + /// Whether any record — readable or not — exists for the domain's config. + /// + /// [`Self::get_zt_domain_config`] cannot distinguish a missing record + /// from a corrupt one; deletion must, or a corrupt record could never be + /// removed. + pub fn zt_domain_config_exists(&self, domain: &str) -> bool { + // `get` already excludes tombstones, so Some means a live record. + self.persistent + .read() + .get(&keys::zt_domain_config(domain)) + .is_some() + } + /// Save ZT-Domain configuration pub fn save_zt_domain_config(&self, config: &ZtDomainConfig) -> Result<()> { self.persistent diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 15f1e48af..cbed26c74 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -10,7 +10,7 @@ use std::{ time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; -use anyhow::{bail, Context, Result}; +use anyhow::{bail, ensure, Context, Result}; use auth_client::AuthClient; use crate::distributed_certbot::DistributedCertBot; @@ -34,6 +34,7 @@ use tokio::sync::{ }; use tokio_rustls::TlsAcceptor; use tracing::{debug, error, info, warn}; +use wavekv::types::NodeId; use crate::{ cert_store::{CertResolver, CertStoreBuilder}, @@ -140,6 +141,27 @@ pub struct CvmRemoval { pub removed_locally: bool, } +/// Outcome of an operator-initiated gateway node removal. +/// +/// Both fields are false when the node was never known (or the removal +/// already completed), which lets the operator distinguish a mistyped +/// node_id from an actual removal. +pub struct NodeRemoval { + /// A live node record existed in WaveKV. + pub record_existed: bool, + /// The node was still in this gateway's sync peer set. + pub removed_from_peer_set: bool, +} + +/// A refused instance record plus its local data-plane footprint. +pub struct RejectedInstanceReport { + pub rejected: import::RejectedInstance, + /// Whether the instance still holds state in this node's data plane. An + /// unusable record keeps whatever the data plane already had, so removing + /// an active instance also drops its routing. + pub active_locally: bool, +} + impl Proxy { /// Remove one CVM by explicit operator request. /// @@ -166,6 +188,47 @@ impl Proxy { }) } + /// Instance records this node currently refuses to import. + /// + /// Recomputed from the store on every call rather than read from the + /// cached rejection log, so the answer is current even right after a + /// restart and does not depend on when the last reload ran. + pub fn rejected_instances(&self) -> Vec { + let rejected = + import::accept_instances(&self.config.wg, self.kv_store.load_all_instances()).rejected; + let state = self.lock(); + rejected + .into_iter() + .map(|rejected| RejectedInstanceReport { + active_locally: state.state.instances.contains_key(&rejected.instance_id), + rejected, + }) + .collect() + } + + /// Remove a decommissioned gateway node by explicit operator request. + /// + /// Tombstones the node's replicated records and drops it from this + /// gateway's sync peer set immediately; other gateways prune their own + /// sets when the `__peer_addr` tombstone reaches them. A node removed by + /// mistake rejoins when it restarts (startup re-registers its records), + /// or via `SetNodeUrl` from any live gateway. + pub fn remove_node(&self, node_id: NodeId) -> Result { + ensure!( + node_id != self.config.sync.node_id, + "a node cannot remove itself" + ); + let record_existed = self + .kv_store + .sync_remove_node(node_id) + .with_context(|| format!("failed to delete node {node_id} from WaveKV"))?; + let removed_from_peer_set = self.kv_store.remove_peer(node_id)?; + Ok(NodeRemoval { + record_existed, + removed_from_peer_set, + }) + } + pub async fn new(options: ProxyOptions) -> Result { let (port_policy_tx, port_policy_rx) = unbounded_channel(); let inner = ProxyInner::new(options, port_policy_tx).await?; @@ -928,6 +991,21 @@ fn start_wavekv_watch_task(proxy: Proxy) -> Result<()> { } }); + // Watch for peer address deletions and prune the sync peer set, so a + // node removed by an operator on any gateway stops being a sync target + // here without a restart. + let mut rx = kv_store.watch_peer_addrs(); + let kv_for_peers = kv_store.clone(); + kv_for_peers.prune_removed_peers(); + tokio::spawn(async move { + loop { + if rx.changed().await.is_err() { + break; + } + kv_for_peers.prune_removed_peers(); + } + }); + // Start periodic persistence task let persist_interval = proxy.config.sync.persist_interval; if !persist_interval.is_zero() { diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index a87d18394..02a983401 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -664,6 +664,123 @@ async fn removing_an_unknown_cvm_reports_that_nothing_existed() { assert!(!removal.removed_locally); } +#[tokio::test] +async fn rejected_instance_records_are_visible_to_the_operator() { + let state = create_test_state().await; + sync_from_peer(&state, "peer-instance", "10.0.0.40", &test_pubkey("good")); + reload_instances_from_kv_store(&state.proxy, &state.kv_store).unwrap(); + assert!(state.proxy.rejected_instances().is_empty()); + + state + .kv_store + .persistent() + .write() + .put( + crate::kv::keys::inst("peer-instance"), + b"not-messagepack".to_vec(), + ) + .unwrap(); + + // The operator can see what is wrong — and that removal would also drop + // the instance's live routing — without grepping logs. + let rejected = state.proxy.rejected_instances(); + assert_eq!(rejected.len(), 1); + assert_eq!(rejected[0].rejected.instance_id, "peer-instance"); + assert!(format!("{:#}", rejected[0].rejected.reason).contains("does not decode")); + assert!(rejected[0].active_locally); + + // Once removed, the record no longer shows up as rejected. + state.proxy.remove_cvm("peer-instance").unwrap(); + assert!(state.proxy.rejected_instances().is_empty()); +} + +#[tokio::test] +async fn an_operator_can_remove_a_decommissioned_node() { + let state = create_test_state().await; + let kv = &state.kv_store; + kv.register_peer_url(7, "https://gw7.example.com:9202") + .unwrap(); + kv.sync_node( + 7, + &crate::kv::NodeData { + uuid: b"gw7-uuid".to_vec(), + url: "https://gw7.example.com:9202".to_string(), + wg_public_key: String::new(), + wg_endpoint: String::new(), + wg_ip: String::new(), + }, + ) + .unwrap(); + + let removal = state.proxy.remove_node(7).unwrap(); + assert!(removal.record_existed); + assert!(removal.removed_from_peer_set); + assert!(kv.get_peer_url(7).is_none()); + assert!(!kv.load_all_nodes().contains_key(&7)); + let peers = kv.persistent().read().status().peers; + assert!(!peers.iter().any(|peer| peer.id == 7)); + + // The recovery operation is safe to retry after a timeout or lost reply, + // and the retry tells the operator there was nothing left to remove. + let retry = state.proxy.remove_node(7).unwrap(); + assert!(!retry.record_existed); + assert!(!retry.removed_from_peer_set); +} + +#[tokio::test] +async fn a_node_cannot_remove_itself() { + let state = create_test_state().await; + let my_id = state.proxy.config.sync.node_id; + assert!(state.proxy.remove_node(my_id).is_err()); +} + +#[tokio::test] +async fn peers_prune_nodes_removed_on_another_gateway() { + let state = create_test_state().await; + let kv = &state.kv_store; + + // Simulate observing a removal performed elsewhere: the __peer_addr + // tombstone arrives via replication, not through this node's admin API. + kv.register_peer_url(9, "https://gw9.example.com:9202") + .unwrap(); + kv.persistent() + .write() + .delete(crate::kv::keys::peer_addr(9)) + .unwrap(); + kv.prune_removed_peers(); + let peers = kv.persistent().read().status().peers; + assert!(!peers.iter().any(|peer| peer.id == 9)); + + // A peer whose address record was never written is not dropped: + // bootstrap can add a peer before its address has synced in. + kv.add_peer(11).unwrap(); + kv.prune_removed_peers(); + let peers = kv.persistent().read().status().peers; + assert!(peers.iter().any(|peer| peer.id == 11)); +} + +#[tokio::test] +async fn a_zt_domain_with_a_corrupt_config_can_still_be_deleted() { + let state = create_test_state().await; + let kv = &state.kv_store; + + kv.persistent() + .write() + .put( + crate::kv::keys::zt_domain_config("bad.example.com"), + b"not-messagepack".to_vec(), + ) + .unwrap(); + + // The corrupt record is invisible to reads, but must still be deletable — + // otherwise it would be permanently stuck. + assert!(kv.get_zt_domain_config("bad.example.com").is_none()); + assert!(kv.zt_domain_config_exists("bad.example.com")); + + kv.delete_zt_domain_config("bad.example.com").unwrap(); + assert!(!kv.zt_domain_config_exists("bad.example.com")); +} + #[tokio::test] async fn an_instance_that_lost_an_ip_conflict_stops_being_routable() { let state = create_test_state().await; From 8521ce9212c018e75551ab172f71bdb1110c8031 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 12 Aug 2026 03:58:04 -0700 Subject: [PATCH 2/3] fix(gateway): address review findings on recovery APIs - RemoveNode now reports record_existed when any of the node's persistent records (info, status, or sync address) was live, so a node registered via SetNodeUrl but never booted is still visible to the operator. The three deletes also run under one write handle. - Loading instances no longer logs each undecodable record at error! level. That log line fired on every reload and every rejected-record listing, bypassing the transition-only reporting the reload path already has. The decode error now travels in LoadedInstances instead, which also gives ListRejectedInstances the actual error to show rather than a generic "record does not decode". --- dstack/gateway/rpc/proto/gateway_rpc.proto | 4 +-- dstack/gateway/src/kv/import.rs | 11 +++--- dstack/gateway/src/kv/mod.rs | 40 +++++++++++++--------- dstack/gateway/src/main_service.rs | 3 +- dstack/gateway/src/main_service/tests.rs | 15 +++++--- 5 files changed, 45 insertions(+), 28 deletions(-) diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index 7ea0bcb22..c20f8cd8c 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -553,8 +553,8 @@ message RemoveNodeRequest { // (or the removal already completed), so a mistyped node_id is visible to // the operator instead of silently reporting success. message RemoveNodeResponse { - // Whether a live node record existed in WaveKV before the tombstone was - // written. + // Whether any of the node's records (info, status, or sync address) was + // live in WaveKV before the tombstones were written. bool record_existed = 1; // Whether the node was still in this gateway's sync peer set. bool removed_from_peer_set = 2; diff --git a/dstack/gateway/src/kv/import.rs b/dstack/gateway/src/kv/import.rs index 64d0ae410..da187c9fd 100644 --- a/dstack/gateway/src/kv/import.rs +++ b/dstack/gateway/src/kv/import.rs @@ -200,9 +200,9 @@ fn accept_instances_at(wg: &WgConfig, loaded: LoadedInstances, now: u64) -> Acce let mut instances = BTreeMap::new(); let mut rejected: Vec = undecodable .into_iter() - .map(|instance_id| RejectedInstance { + .map(|(instance_id, reason)| RejectedInstance { instance_id, - reason: anyhow::anyhow!("record does not decode"), + reason: anyhow::anyhow!(reason), rejection: Rejection::Unusable, }) .collect(); @@ -250,7 +250,6 @@ fn accept_instances_at(wg: &WgConfig, loaded: LoadedInstances, now: u64) -> Acce mod tests { use super::*; use ipnet::Ipv4Net; - use std::collections::BTreeSet; /// Wall clock the tests validate against; every fixture `reg_time` below is /// well under it unless the test is about the future-timestamp horizon. @@ -292,7 +291,7 @@ mod tests { .into_iter() .map(|(id, data)| (id.to_string(), data)) .collect(), - undecodable: BTreeSet::new(), + undecodable: BTreeMap::new(), } } @@ -503,7 +502,9 @@ mod tests { decoded: [("good".to_string(), instance("10.0.0.20", &key(1), 100))] .into_iter() .collect(), - undecodable: ["corrupt".to_string()].into_iter().collect(), + undecodable: [("corrupt".to_string(), "does not decode".to_string())] + .into_iter() + .collect(), }, NOW, ); diff --git a/dstack/gateway/src/kv/mod.rs b/dstack/gateway/src/kv/mod.rs index 6fed6e0b1..d186e46d7 100644 --- a/dstack/gateway/src/kv/mod.rs +++ b/dstack/gateway/src/kv/mod.rs @@ -36,12 +36,7 @@ pub use https_client::{AppIdValidator, HttpsClientConfig}; pub use sync_service::{fetch_peers_from_bootnode, WaveKvSyncService}; use tracing::{error, warn}; -use std::{ - collections::{BTreeMap, BTreeSet}, - net::Ipv4Addr, - path::Path, - time::Duration, -}; +use std::{collections::BTreeMap, net::Ipv4Addr, path::Path, time::Duration}; use anyhow::{Context, Result}; @@ -109,8 +104,11 @@ pub struct InstanceData { pub struct LoadedInstances { /// Records that decoded successfully, keyed by instance ID. pub decoded: BTreeMap, - /// Instance IDs whose stored bytes are present but no longer decode. - pub undecodable: BTreeSet, + /// Instance IDs whose stored bytes are present but no longer decode, + /// mapped to the decode error. Loading does not log these: the reload + /// path reports them on transitions, and read-only listings must stay + /// quiet no matter how often an operator runs them. + pub undecodable: BTreeMap, } /// Gateway node status (stored separately for independent updates) @@ -758,8 +756,9 @@ impl KvStore { loaded.decoded.insert(instance_id.into(), data); } Err(err) => { - error!("{err:#}"); - loaded.undecodable.insert(instance_id.into()); + loaded + .undecodable + .insert(instance_id.into(), format!("{err:#}")); } } } @@ -795,16 +794,25 @@ impl KvStore { /// prunes its sync peer set when it observes the deletion (see /// [`Self::prune_removed_peers`]). /// - /// Returns whether a live node record existed before the tombstone was - /// written. + /// Returns whether any of the node's persistent records was live before + /// the tombstones were written, so a node known only by its sync address + /// (registered via `SetNodeUrl` but never booted) still reports as + /// existing. pub fn sync_remove_node(&self, node_id: NodeId) -> Result { - let previous = self.persistent.write().delete(keys::node_info(node_id))?; - self.persistent.write().delete(keys::node_status(node_id))?; - self.persistent.write().delete(keys::peer_addr(node_id))?; + let previous = { + let mut persistent = self.persistent.write(); + [ + persistent.delete(keys::node_info(node_id))?, + persistent.delete(keys::node_status(node_id))?, + persistent.delete(keys::peer_addr(node_id))?, + ] + }; self.ephemeral .write() .delete(keys::last_seen_node(node_id, self.my_node_id))?; - Ok(previous.is_some_and(|entry| !entry.is_deleted())) + Ok(previous + .into_iter() + .any(|entry| entry.is_some_and(|entry| !entry.is_deleted()))) } // ==================== Node Status Sync ==================== diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index cbed26c74..2f020f063 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -147,7 +147,8 @@ pub struct CvmRemoval { /// already completed), which lets the operator distinguish a mistyped /// node_id from an actual removal. pub struct NodeRemoval { - /// A live node record existed in WaveKV. + /// Any of the node's records (info, status, or sync address) was live + /// in WaveKV. pub record_existed: bool, /// The node was still in this gateway's sync peer set. pub removed_from_peer_set: bool, diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index 02a983401..9d03cad9e 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -645,7 +645,7 @@ async fn an_operator_can_remove_a_cvm_whose_kv_record_is_unreadable() { assert!(!state.lock().state.instances.contains_key("peer-instance")); let loaded = state.kv_store.load_all_instances(); assert!(!loaded.decoded.contains_key("peer-instance")); - assert!(!loaded.undecodable.contains("peer-instance")); + assert!(!loaded.undecodable.contains_key("peer-instance")); // The recovery operation is safe to retry after a timeout or lost reply, // and the retry tells the operator there was nothing left to remove. @@ -681,12 +681,12 @@ async fn rejected_instance_records_are_visible_to_the_operator() { ) .unwrap(); - // The operator can see what is wrong — and that removal would also drop - // the instance's live routing — without grepping logs. + // The operator can see what is wrong — with the actual decode error, and + // whether removal would also drop live routing — without grepping logs. let rejected = state.proxy.rejected_instances(); assert_eq!(rejected.len(), 1); assert_eq!(rejected[0].rejected.instance_id, "peer-instance"); - assert!(format!("{:#}", rejected[0].rejected.reason).contains("does not decode")); + assert!(format!("{:#}", rejected[0].rejected.reason).contains("corrupt record")); assert!(rejected[0].active_locally); // Once removed, the record no longer shows up as rejected. @@ -725,6 +725,13 @@ async fn an_operator_can_remove_a_decommissioned_node() { let retry = state.proxy.remove_node(7).unwrap(); assert!(!retry.record_existed); assert!(!retry.removed_from_peer_set); + + // A node known only by its sync address (registered via SetNodeUrl but + // never booted) still reports record_existed. + kv.register_peer_url(8, "https://gw8.example.com:9202") + .unwrap(); + let removal = state.proxy.remove_node(8).unwrap(); + assert!(removal.record_existed); } #[tokio::test] From 451e2d8240013990ede20d789d9c13b4163f4eb5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 12 Aug 2026 04:13:47 -0700 Subject: [PATCH 3/3] fix(gateway): capture peer membership before publishing the node tombstone RemoveNode wrote the __peer_addr tombstone first and asked the peer set second. The tombstone wakes the peer-address watcher, so on a multi-threaded runtime prune_removed_peers() could drop the peer before the RPC path's own remove_peer() ran, and removed_from_peer_set would report false for a peer that was present when the request began. Drop the peer before publishing the tombstone, so the reported membership no longer depends on scheduling. Add a regression test that races remove_node against a live watcher on a multi-threaded runtime. --- dstack/gateway/src/main_service.rs | 6 ++++- dstack/gateway/src/main_service/tests.rs | 33 ++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/dstack/gateway/src/main_service.rs b/dstack/gateway/src/main_service.rs index 2f020f063..6b9763c23 100644 --- a/dstack/gateway/src/main_service.rs +++ b/dstack/gateway/src/main_service.rs @@ -219,11 +219,15 @@ impl Proxy { node_id != self.config.sync.node_id, "a node cannot remove itself" ); + // Drop the peer before publishing the tombstone: the __peer_addr + // deletion wakes the peer-address watcher, whose prune would + // otherwise race this call and make the reported membership depend + // on scheduling. + let removed_from_peer_set = self.kv_store.remove_peer(node_id)?; let record_existed = self .kv_store .sync_remove_node(node_id) .with_context(|| format!("failed to delete node {node_id} from WaveKV"))?; - let removed_from_peer_set = self.kv_store.remove_peer(node_id)?; Ok(NodeRemoval { record_existed, removed_from_peer_set, diff --git a/dstack/gateway/src/main_service/tests.rs b/dstack/gateway/src/main_service/tests.rs index 9d03cad9e..7135c9007 100644 --- a/dstack/gateway/src/main_service/tests.rs +++ b/dstack/gateway/src/main_service/tests.rs @@ -766,6 +766,39 @@ async fn peers_prune_nodes_removed_on_another_gateway() { assert!(peers.iter().any(|peer| peer.id == 11)); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn remove_node_reports_peer_membership_despite_a_racing_watcher() { + let state = create_test_state().await; + let kv = state.kv_store.clone(); + + // The production watch task prunes the peer set as soon as the + // __peer_addr tombstone lands. remove_node must capture membership + // before publishing the tombstone, or the answer it returns would + // depend on which of the two gets there first. + let mut rx = kv.watch_peer_addrs(); + let kv_for_watch = kv.clone(); + let watcher = tokio::spawn(async move { + while rx.changed().await.is_ok() { + kv_for_watch.prune_removed_peers(); + } + }); + + for node_id in 100..120 { + kv.register_peer_url(node_id, "https://gw.example.com:9202") + .unwrap(); + let removal = state.proxy.remove_node(node_id).unwrap(); + assert!(removal.record_existed); + assert!( + removal.removed_from_peer_set, + "node {node_id}: membership must be captured before the tombstone publishes" + ); + let peers = kv.persistent().read().status().peers; + assert!(!peers.iter().any(|peer| peer.id == node_id)); + } + + watcher.abort(); +} + #[tokio::test] async fn a_zt_domain_with_a_corrupt_config_can_still_be_deleted() { let state = create_test_state().await;