From 08a1640e3042d2f77cce57db0788925c1ed5f383 Mon Sep 17 00:00:00 2001
From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com>
Date: Tue, 31 Mar 2026 03:37:23 +0200
Subject: [PATCH] server: lifecycle

---
 easytier/src/dns/node.rs   | 6 +++---
 easytier/src/dns/server.rs | 6 +++++-
 easytier/src/utils.rs      | 6 +++---
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/easytier/src/dns/node.rs b/easytier/src/dns/node.rs
index c9cb3344..77d04f09 100644
--- a/easytier/src/dns/node.rs
+++ b/easytier/src/dns/node.rs
@@ -11,6 +11,7 @@ use crate::proto::dns::{DnsNodeMgrRpcClientFactory, DnsPeerMgrRpcServer, Heartbe
 use crate::proto::rpc_impl::standalone::{StandAloneClient, StandAloneServer};
 use crate::proto::rpc_types::controller::BaseController;
 use crate::tunnel::tcp::{TcpTunnelConnector, TcpTunnelListener};
+use crate::utils::AsyncRuntime;
 use derivative::Derivative;
 use futures::task::SpawnExt;
 use std::sync::Arc;
@@ -20,7 +21,6 @@ use tokio::task::{JoinHandle, JoinSet};
 use tokio::time::{sleep, sleep_until, Instant};
 use tokio_util::sync::CancellationToken;
 use uuid::Uuid;
-use crate::utils::AsyncRuntime;
 
 #[derive(Debug)]
 struct DnsNodeRuntime {
@@ -74,7 +74,7 @@ impl DnsNode {
 
     pub fn start(&self) {
         let this = self.clone();
-        self.runtime.start(|token| async move {
+        self.runtime.start(None, |token| async move {
             tokio::join!(this.run_election(token.clone()), this.run_node(token));
         });
     }
@@ -117,7 +117,7 @@ impl DnsNode {
             tokio::join!(
                 self.peer_mgr
                     .add_nic_packet_process_pipeline(Box::new(server.clone())),
-                server.run()
+                server.run(token.child_token())
             );
 
             self.global_ctx.set_dns(None);
diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs
index f5ea197b..be5986d3 100644
--- a/easytier/src/dns/server.rs
+++ b/easytier/src/dns/server.rs
@@ -270,7 +270,7 @@ impl DnsServer {
         Ok(())
     }
 
-    pub async fn run(&self) {
+    pub async fn run(&self, token: CancellationToken) {
         let dirty = &self.mgr.dirty;
         let mut runtime = None;
 
@@ -314,6 +314,10 @@ impl DnsServer {
         };
 
         tokio::select!(
+            _ = token.cancelled() => {
+                tracing::info!("DnsServer received shutdown signal, exiting server loop");
+            }
+
             _ = reload_catalog => {},
             _ = reload_addresses => {},
             _ = reload_listeners => {},
diff --git a/easytier/src/utils.rs b/easytier/src/utils.rs
index 6bb34d3d..c1bf09e8 100644
--- a/easytier/src/utils.rs
+++ b/easytier/src/utils.rs
@@ -195,7 +195,7 @@ impl<T: Send + 'static> AsyncRuntime<T> {
         self.inner.lock().as_ref().map(|r| r.token.clone())
     }
 
-    pub fn start<F, Fut>(&self, factory: F)
+    pub fn start<F, Fut>(&self, token: Option<CancellationToken>, factory: F)
     where
         F: FnOnce(CancellationToken) -> Fut,
         Fut: Future<Output = T> + Send + 'static,
@@ -207,8 +207,8 @@ impl<T: Send + 'static> AsyncRuntime<T> {
                 return;
             }
         }
-
-        let token = CancellationToken::new();
+        
+        let token = token.unwrap_or_default();
         runtime.replace(AsyncRuntimeInner {
             task: tokio::spawn(factory(token.clone())).into(),
             token,