diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index d410aa451..c20f8cd8c 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 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; +} + // ==================== 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/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 1c6ce5705..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:#}")); } } } @@ -787,6 +786,35 @@ 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 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 = { + 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 + .into_iter() + .any(|entry| entry.is_some_and(|entry| !entry.is_deleted()))) + } + // ==================== Node Status Sync ==================== /// Set node status (stored separately from NodeData) @@ -946,6 +974,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 +993,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 +1176,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..6b9763c23 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,28 @@ 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 { + /// 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, +} + +/// 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 +189,51 @@ 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" + ); + // 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"))?; + 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 +996,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..7135c9007 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. @@ -664,6 +664,163 @@ 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 — 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("corrupt record")); + 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); + + // 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] +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(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; + 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;