diff --git a/easytier/src/dns/node.rs b/easytier/src/dns/node.rs index 0b8b5e7b..8895e349 100644 --- a/easytier/src/dns/node.rs +++ b/easytier/src/dns/node.rs @@ -207,7 +207,7 @@ impl DnsNodeRuntime { #[derive(Debug)] pub struct DnsNode { runtime: DnsNodeRuntime, - task: CancellableTask<()>, + task: Option>, } impl DnsNode { @@ -225,19 +225,34 @@ impl DnsNode { elect: Default::default(), }; - let task = { - let runtime = runtime.clone(); - CancellableTask::spawn(|token| async move { - runtime.elect.notify_one(); - tokio::join!(runtime.run_election(token.clone()), runtime.run(token)); - }) - }; - - Self { runtime, task } + Self { + runtime, + task: None, + } } - pub async fn stop(self) -> io::Result<()> { - self.task.stop(None).await + pub fn start(&mut self) { + let runtime = self.runtime.clone(); + self.task + .replace(CancellableTask::spawn(|token| async move { + runtime.elect.notify_one(); + tokio::join!(runtime.run_election(token.clone()), runtime.run(token)); + })); + self.runtime.mgr.register(); + } + + pub async fn stop(&mut self) -> io::Result<()> { + self.runtime.mgr.unregister(); + let Some(task) = self.task.take() else { + return Ok(()); + }; + task.stop(None).await + } +} + +impl Drop for DnsNode { + fn drop(&mut self) { + self.runtime.mgr.unregister(); } } diff --git a/easytier/src/dns/peer_mgr.rs b/easytier/src/dns/peer_mgr.rs index b03eddfb..971e0511 100644 --- a/easytier/src/dns/peer_mgr.rs +++ b/easytier/src/dns/peer_mgr.rs @@ -176,17 +176,15 @@ pub struct DnsPeerMgr(Arc); impl DnsPeerMgr { pub fn new(peer_mgr: Arc, global_ctx: ArcGlobalCtx) -> Self { - let this = Self(Arc::new(DnsPeerMgrInner { + Self(Arc::new(DnsPeerMgrInner { peers: Cache::builder().time_to_idle(DNS_PEER_TTI).build(), dirty: Default::default(), peer_mgr, global_ctx, - })); - this.register(); - this + })) } - fn register(&self) { + pub fn register(&self) { self.peer_mgr .get_peer_rpc_mgr() .rpc_server() @@ -197,7 +195,7 @@ impl DnsPeerMgr { ); } - fn unregister(&self) -> Option<()> { + pub fn unregister(&self) -> Option<()> { self.peer_mgr .get_peer_rpc_mgr() .rpc_server() @@ -237,12 +235,6 @@ impl Deref for DnsPeerMgr { } } -impl Drop for DnsPeerMgr { - fn drop(&mut self) { - self.unregister(); - } -} - #[cfg(test)] mod tests { use super::*; @@ -658,6 +650,7 @@ mod tests { .await; let local_dns = DnsPeerMgr::new(local.clone(), local.get_global_ctx()); + local_dns.register(); let keep_dns = DnsPeerMgr::new(keep_peer.clone(), keep_peer.get_global_ctx()); keep_dns.register(); @@ -864,33 +857,7 @@ mod tests { ) .await; let mgr = DnsPeerMgr::new(peer_mgr.clone(), peer_mgr.get_global_ctx()); + mgr.register(); assert!(mgr.unregister().is_some()); } - - #[tokio::test] - async fn drop_triggers_unregister() { - let peer_mgr = create_peer_manager_with_zone( - "drop-peer", - "drop-zone.test", - Ipv4Addr::new(10, 1, 0, 3), - ) - .await; - let global_ctx = peer_mgr.get_global_ctx(); - - let mgr = DnsPeerMgr::new(peer_mgr.clone(), global_ctx.clone()); - let inner = mgr.0.clone(); - mgr.register(); - - drop(mgr); - - let res = peer_mgr - .get_peer_rpc_mgr() - .rpc_server() - .registry() - .unregister( - DnsPeerMgrRpcServer::new_arc(inner), - &global_ctx.get_network_name(), - ); - assert!(res.is_none()); - } } diff --git a/easytier/src/dns/system/windows.rs b/easytier/src/dns/system/windows.rs index 6c10d5b3..048f7864 100644 --- a/easytier/src/dns/system/windows.rs +++ b/easytier/src/dns/system/windows.rs @@ -190,7 +190,7 @@ mod tests { // prepare_env does not run full Instance::run, so start the monitor explicitly in test. let _monitor = ProxyCidrsMonitor::new(peer_mgr.clone(), peer_mgr.get_global_ctx()).start(); - let dns_node = start_dns_node(peer_mgr, virtual_nic); + let mut dns_node = start_dns_node(peer_mgr, virtual_nic); println!("dev_name: {}", tun_name); let fake_ip = match DNS_DEFAULT_ADDRESSES[0].addr.ip() { diff --git a/easytier/src/dns/tests.rs b/easytier/src/dns/tests.rs index fad724f6..e07c2d2e 100644 --- a/easytier/src/dns/tests.rs +++ b/easytier/src/dns/tests.rs @@ -74,14 +74,18 @@ pub fn start_dns_node(peer_mgr: Arc, virtual_nic: NicCtx) -> DnsNod let global_ctx = peer_mgr.get_global_ctx(); let nic_ctx: ArcNicCtx = Arc::new(tokio::sync::Mutex::new(Some(Box::new(virtual_nic)))); - DnsNode::new(peer_mgr, global_ctx, nic_ctx) + let mut node = DnsNode::new(peer_mgr, global_ctx, nic_ctx); + node.start(); + node } pub fn start_dns_node_without_nic(peer_mgr: Arc) -> DnsNode { let global_ctx = peer_mgr.get_global_ctx(); let nic_ctx: ArcNicCtx = Arc::new(tokio::sync::Mutex::new(None)); - DnsNode::new(peer_mgr, global_ctx, nic_ctx) + let mut node = DnsNode::new(peer_mgr, global_ctx, nic_ctx); + node.start(); + node } pub async fn prepare_env_from_config_str(config_str: &str) -> Arc { @@ -405,7 +409,7 @@ whitelist = ["*"] ); let peer_mgr = prepare_env_from_config_str(&config).await; - let dns_node = start_dns_node_without_nic(peer_mgr); + let mut dns_node = start_dns_node_without_nic(peer_mgr); let server_addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_port); check_dns_record_at(server_addr, "alpha.et-test.", "10.144.144.11").await; @@ -443,8 +447,8 @@ whitelist = ["*"] .await .expect("route should appear"); - let node_a = start_dns_node_without_nic(peer_a); - let node_b = start_dns_node_without_nic(peer_b); + let mut node_a = start_dns_node_without_nic(peer_a); + let mut node_b = start_dns_node_without_nic(peer_b); let addr_a = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_a); let addr_b = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_b); @@ -486,8 +490,10 @@ records = ["secret IN A 10.99.0.9"] // Export behavior is determined by whether `[dns.zone.export]` exists. // Verify from peer-sync view to avoid host-wide DNS-server election side effects. - let _dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); + let dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); + dns_a.register(); let dns_b = DnsPeerMgr::new(peer_b.clone(), peer_b.get_global_ctx()); + dns_b.register(); dns_b .refresh(peer_a.my_peer_id(), Default::default(), Default::default()) .await @@ -532,8 +538,10 @@ disabled = true .await .expect("route should appear"); - let _dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); + let dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); + dns_a.register(); let dns_b = DnsPeerMgr::new(peer_b.clone(), peer_b.get_global_ctx()); + dns_b.register(); dns_b .refresh(peer_a.my_peer_id(), Default::default(), Default::default()) .await @@ -567,7 +575,7 @@ records = ["api IN A 10.80.0.1"] ); let peer = prepare_env_from_config_str(&config).await; - let dns_node = start_dns_node_without_nic(peer.clone()); + let mut dns_node = start_dns_node_without_nic(peer.clone()); let server_addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener); check_dns_record_at(server_addr, "api.patch.mesh-test.", "10.80.0.1").await; @@ -605,7 +613,7 @@ async fn config_patch_reloads_listener_binding() { ); let peer = prepare_env_from_config_str(&config).await; - let dns_node = start_dns_node_without_nic(peer.clone()); + let mut dns_node = start_dns_node_without_nic(peer.clone()); let old_addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_old); let new_addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_new); @@ -662,9 +670,9 @@ records = ["svc IN A 10.77.7.7"] .await .expect("route a-c should appear via b"); - let node_a = start_dns_node_without_nic(peer_a.clone()); - let node_b = start_dns_node_without_nic(peer_b.clone()); - let node_c = start_dns_node_without_nic(peer_c.clone()); + let mut node_a = start_dns_node_without_nic(peer_a.clone()); + let mut node_b = start_dns_node_without_nic(peer_b.clone()); + let mut node_c = start_dns_node_without_nic(peer_c.clone()); let addr_a = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_a); check_dns_record_at(addr_a, "node-c7.mesh7-test.", "10.144.151.13").await; @@ -768,9 +776,9 @@ records = ["api IN A 10.66.1.8"] .await .expect("route a-c should appear via b"); - let node_a = start_dns_node_without_nic(peer_a); - let node_b = start_dns_node_without_nic(peer_b); - let node_c = start_dns_node_without_nic(peer_c); + let mut node_a = start_dns_node_without_nic(peer_a); + let mut node_b = start_dns_node_without_nic(peer_b); + let mut node_c = start_dns_node_without_nic(peer_c); let addr_a = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_a); let addr_b = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_b); @@ -818,8 +826,8 @@ records = ["svc IN A 10.66.2.8"] .await .expect("route a-b should appear"); - let node_a = start_dns_node_without_nic(peer_a.clone()); - let node_b = start_dns_node_without_nic(peer_b.clone()); + let mut node_a = start_dns_node_without_nic(peer_a.clone()); + let mut node_b = start_dns_node_without_nic(peer_b.clone()); let addr_a = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), listener_a); check_dns_record_missing_at(addr_a, "node-c4.mesh4-test.").await; @@ -832,7 +840,7 @@ records = ["svc IN A 10.66.2.8"] .await .expect("route a-c should appear after c joins"); - let node_c = start_dns_node_without_nic(peer_c); + let mut node_c = start_dns_node_without_nic(peer_c); check_dns_record_at(addr_a, "node-c4.mesh4-test.", "10.144.146.13").await; check_dns_record_at(addr_a, "svc.joined.mesh4-test.", "10.66.2.8").await; @@ -872,8 +880,10 @@ records = ["secret IN A 10.66.3.8"] .await .expect("route a-c should appear via b"); - let _dns_c = DnsPeerMgr::new(peer_c.clone(), peer_c.get_global_ctx()); + let dns_c = DnsPeerMgr::new(peer_c.clone(), peer_c.get_global_ctx()); + dns_c.register(); let dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); + dns_a.register(); dns_a .refresh(peer_c.my_peer_id(), Default::default(), Default::default()) .await @@ -906,8 +916,10 @@ async fn config_string_two_nodes_peer_dns_offline_then_rejoin() { .await .expect("route should appear"); - let _dns_b = DnsPeerMgr::new(peer_b.clone(), peer_b.get_global_ctx()); + let dns_b = DnsPeerMgr::new(peer_b.clone(), peer_b.get_global_ctx()); + dns_b.register(); let dns_a_online = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); + dns_a_online.register(); dns_a_online .refresh(peer_b.my_peer_id(), Default::default(), Default::default()) .await diff --git a/easytier/src/instance/instance.rs b/easytier/src/instance/instance.rs index 59a98646..9f8a96f9 100644 --- a/easytier/src/instance/instance.rs +++ b/easytier/src/instance/instance.rs @@ -931,12 +931,14 @@ impl Instance { #[cfg(feature = "magic-dns")] { - self.dns = Some(DnsNode::new( + let mut node = DnsNode::new( self.get_peer_manager(), self.get_global_ctx(), #[cfg(feature = "tun")] self.get_nic_ctx(), - )); + ); + node.start(); + self.dns = Some(node); } if self.global_ctx.config.get_dhcp() { @@ -1485,7 +1487,7 @@ impl Instance { pub async fn clear_resources(&mut self) { self.peer_manager.clear_resources().await; #[cfg(feature = "magic-dns")] - if let Some(node) = self.dns.take() { + if let Some(mut node) = self.dns.take() { let _ = node.stop().await; } #[cfg(feature = "tun")]