node & peer_mgr: fix rpc register/unregister

mut node

easytier/src/dns/node.rs

@@ -207,7 +207,7 @@ impl DnsNodeRuntime {
 #[derive(Debug)]
 pub struct DnsNode {
     runtime: DnsNodeRuntime,
-    task: CancellableTask<()>,
+    task: Option<CancellableTask<()>>,
 }
 
 impl DnsNode {
@@ -225,19 +225,34 @@ impl DnsNode {
             elect: Default::default(),
         };
 
-        let task = {
+        Self {
-            let runtime = runtime.clone();
+            runtime,
-            CancellableTask::spawn(|token| async move {
+            task: None,
-                runtime.elect.notify_one();
+        }
-                tokio::join!(runtime.run_election(token.clone()), runtime.run(token));
-            })
-        };
-
-        Self { runtime, task }
     }
 
-    pub async fn stop(self) -> io::Result<()> {
+    pub fn start(&mut self) {
-        self.task.stop(None).await
+        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();
     }
 }
 

easytier/src/dns/peer_mgr.rs

@@ -176,17 +176,15 @@ pub struct DnsPeerMgr(Arc<DnsPeerMgrInner>);
 
 impl DnsPeerMgr {
     pub fn new(peer_mgr: Arc<PeerManager>, 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());
-    }
 }

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() {

easytier/src/dns/tests.rs

@@ -74,14 +74,18 @@ pub fn start_dns_node(peer_mgr: Arc<PeerManager>, 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<PeerManager>) -> 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<PeerManager> {
@@ -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 mut node_a = start_dns_node_without_nic(peer_a);
-    let node_b = start_dns_node_without_nic(peer_b);
+    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 mut node_a = start_dns_node_without_nic(peer_a.clone());
-    let node_b = start_dns_node_without_nic(peer_b.clone());
+    let mut node_b = start_dns_node_without_nic(peer_b.clone());
-    let node_c = start_dns_node_without_nic(peer_c.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 mut node_a = start_dns_node_without_nic(peer_a);
-    let node_b = start_dns_node_without_nic(peer_b);
+    let mut node_b = start_dns_node_without_nic(peer_b);
-    let node_c = start_dns_node_without_nic(peer_c);
+    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 mut node_a = start_dns_node_without_nic(peer_a.clone());
-    let node_b = start_dns_node_without_nic(peer_b.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

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")]