From 4f6f8f3d934aa42fc62a3b7651266c04c0fae369 Mon Sep 17 00:00:00 2001 From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com> Date: Sat, 18 Apr 2026 14:51:04 +0200 Subject: [PATCH] use CancellableTask in DnsNode remove AsyncRuntime from DnsNode node n n clippy --- easytier/src/dns/node.rs | 214 +++++++++++++++--------------- easytier/src/dns/peer_mgr.rs | 25 +--- easytier/src/dns/tests.rs | 24 +--- easytier/src/instance/instance.rs | 37 ++---- 4 files changed, 132 insertions(+), 168 deletions(-) diff --git a/easytier/src/dns/node.rs b/easytier/src/dns/node.rs index 50368e18..ac2e91d8 100644 --- a/easytier/src/dns/node.rs +++ b/easytier/src/dns/node.rs @@ -1,5 +1,4 @@ use crate::common::global_ctx::{ArcGlobalCtx, GlobalCtxEvent}; -use crate::common::join_joinset_background; use crate::dns::config::{ DNS_NODE_RR_INTERVAL, DNS_SERVER_ELECTION_INTERVAL, DNS_SERVER_RPC_ADDR, DnsGlobalCtxExt, }; @@ -13,8 +12,8 @@ use crate::proto::dns::{DnsNodeMgrRpcClientFactory, HeartbeatRequest}; use crate::proto::rpc_impl::standalone::{StandAloneClient, StandAloneServer}; use crate::proto::rpc_types::controller::BaseController; use crate::tunnel::tcp::{TcpTunnelConnector, TcpTunnelListener}; -use crate::utils::task::AsyncRuntime; -use std::sync::{Arc, Mutex}; +use crate::utils::task::CancellableTask; +use std::sync::Arc; use tokio::sync::{Notify, broadcast}; use tokio::task::{JoinError, JoinSet}; use tokio::time::{Instant, sleep, sleep_until}; @@ -23,7 +22,7 @@ use tracing::instrument; use uuid::Uuid; #[derive(Debug, Clone)] -pub struct DnsNode { +struct DnsNodeRuntime { mgr: Arc, #[cfg(feature = "tun")] @@ -33,44 +32,13 @@ pub struct DnsNode { global_ctx: ArcGlobalCtx, elect: Arc, - runtime: AsyncRuntime, } -impl DnsNode { - pub fn new( - peer_mgr: Arc, - global_ctx: ArcGlobalCtx, - #[cfg(feature = "tun")] nic_ctx: ArcNicCtx, // TODO: REMOVE THIS - ) -> Self { - Self { - mgr: Arc::new(DnsPeerMgr::new(peer_mgr.clone(), global_ctx.clone())), - #[cfg(feature = "tun")] - nic_ctx, - peer_mgr, - global_ctx, - elect: Default::default(), - runtime: Default::default(), - } - } - - pub fn id(&self) -> Uuid { +impl DnsNodeRuntime { + fn id(&self) -> Uuid { self.global_ctx.get_id() } - pub fn start(&self) -> anyhow::Result<()> { - self.mgr.register(); - let this = self.clone(); - self.runtime.start(None, |token| async move { - tracing::info!("starting DnsNode"); - this.elect.notify_one(); - tokio::join!(this.run_election(token.clone()), this.run(token)); - }) - } - - pub async fn stop(&self) -> Result<(), JoinError> { - self.runtime.stop(None).await.unwrap_or(Ok(())) - } - #[instrument(skip_all, name = "DnsNode election loop")] async fn run_election(&self, token: CancellationToken) { loop { @@ -138,8 +106,7 @@ impl DnsNode { tokio::pin!(sleep); let mut subscriber = self.global_ctx.subscribe(); - let tasks = Arc::new(Mutex::new(JoinSet::new())); - join_joinset_background(tasks.clone(), "DnsNode".to_owned()); + let mut tasks = JoinSet::new(); loop { // Dynamic interval: slower if dirty (throttled), faster if clean (fast liveness check) @@ -175,7 +142,7 @@ impl DnsNode { Ok(GlobalCtxEvent::PeerInfoUpdated(peer_ids)) => { for peer_id in peer_ids { let mgr = self.mgr.clone(); - tasks.lock().unwrap().spawn(async move { + tasks.spawn(async move { mgr.refresh(peer_id).await; }); } @@ -203,6 +170,12 @@ impl DnsNode { self.mgr.dirty.mark(); } + + result = tasks.join_next(), if !tasks.is_empty() => { + if let Some(Err(error)) = result { + tracing::error!(?error, "refresh task panicked"); + } + } } } } @@ -239,6 +212,43 @@ impl DnsNode { } } +#[derive(Debug)] +pub struct DnsNode { + runtime: DnsNodeRuntime, + task: CancellableTask, +} + +impl DnsNode { + pub fn new( + peer_mgr: Arc, + global_ctx: ArcGlobalCtx, + #[cfg(feature = "tun")] nic_ctx: ArcNicCtx, // TODO: REMOVE THIS + ) -> Self { + let runtime = DnsNodeRuntime { + mgr: Arc::new(DnsPeerMgr::new(peer_mgr.clone(), global_ctx.clone())), + #[cfg(feature = "tun")] + nic_ctx, + peer_mgr, + global_ctx, + 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 } + } + + pub async fn stop(self) -> Result<(), JoinError> { + self.task.stop(None).await + } +} + #[cfg(all(test, feature = "tun"))] mod tests { use super::*; @@ -292,11 +302,26 @@ mod tests { } } - async fn build_test_node() -> DnsNode { + async fn build_test_runtime() -> DnsNodeRuntime { let peer_mgr = create_mock_peer_manager().await; 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 nic_ctx: ArcNicCtx = Arc::new(Mutex::new(None)); + DnsNodeRuntime { + mgr: Arc::new(DnsPeerMgr::new(peer_mgr.clone(), global_ctx.clone())), + nic_ctx, + peer_mgr, + global_ctx, + elect: Default::default(), + } + } + + async fn build_test_node() -> DnsNode { + let runtime = build_test_runtime().await; + DnsNode::new( + runtime.peer_mgr.clone(), + runtime.global_ctx.clone(), + runtime.nic_ctx.clone(), + ) } async fn start_recording_rpc_server( @@ -311,14 +336,13 @@ mod tests { .registry() .register(DnsNodeMgrRpcServer::new_arc(mgr.clone()), ""); server.serve().await?; - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; Ok((mgr, server)) } async fn occupy_dns_rpc_addr() -> StandAloneServer { let mut server = StandAloneServer::new(TcpTunnelListener::new(DNS_SERVER_RPC_ADDR.clone())); server.serve().await.unwrap(); - tokio::time::sleep(Duration::from_millis(50)).await; server } @@ -326,7 +350,7 @@ mod tests { #[serial_test::serial(dns_node_rpc_addr)] async fn heartbeat_first_send_includes_snapshot() { let (_mgr, server) = start_recording_rpc_server(false).await.unwrap(); - let node = build_test_node().await; + let node = build_test_runtime().await; let mut rpc = StandAloneClient::new(TcpTunnelConnector::new(DNS_SERVER_RPC_ADDR.clone())); let mut heartbeat = HeartbeatRequest { @@ -337,7 +361,7 @@ mod tests { node.heartbeat(&mut rpc, &mut heartbeat).await.unwrap(); drop(server); - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; assert!(heartbeat.snapshot.is_some()); assert!(!heartbeat.digest.is_empty()); @@ -347,7 +371,7 @@ mod tests { #[serial_test::serial(dns_node_rpc_addr)] async fn heartbeat_clean_send_digest_only() { let (mgr, server) = start_recording_rpc_server(false).await.unwrap(); - let node = build_test_node().await; + let node = build_test_runtime().await; let mut rpc = StandAloneClient::new(TcpTunnelConnector::new(DNS_SERVER_RPC_ADDR.clone())); let mut heartbeat = HeartbeatRequest { @@ -361,7 +385,7 @@ mod tests { let requests = mgr.recorded_requests().await; drop(server); - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; assert_eq!(requests.len(), 2); assert!(requests[0].snapshot.is_some()); @@ -373,7 +397,7 @@ mod tests { #[serial_test::serial(dns_node_rpc_addr)] async fn heartbeat_dirty_forces_full_snapshot() { let (mgr, server) = start_recording_rpc_server(false).await.unwrap(); - let node = build_test_node().await; + let node = build_test_runtime().await; let mut rpc = StandAloneClient::new(TcpTunnelConnector::new(DNS_SERVER_RPC_ADDR.clone())); let mut heartbeat = HeartbeatRequest { @@ -387,7 +411,7 @@ mod tests { let requests = mgr.recorded_requests().await; drop(server); - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; assert_eq!(requests.len(), 2); assert!(requests[0].snapshot.is_some()); @@ -398,7 +422,7 @@ mod tests { #[serial_test::serial(dns_node_rpc_addr)] async fn heartbeat_resync_triggers_second_send() { let (mgr, server) = start_recording_rpc_server(true).await.unwrap(); - let node = build_test_node().await; + let node = build_test_runtime().await; let mut rpc = StandAloneClient::new(TcpTunnelConnector::new(DNS_SERVER_RPC_ADDR.clone())); let mut heartbeat = HeartbeatRequest { @@ -410,7 +434,7 @@ mod tests { let requests = mgr.recorded_requests().await; drop(server); - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; assert_eq!(requests.len(), 2); assert!(requests[0].snapshot.is_some()); @@ -420,7 +444,8 @@ mod tests { #[tokio::test] #[serial_test::serial(dns_node_rpc_addr)] async fn run_marks_dirty_on_dhcp_event() { - let node = build_test_node().await; + let node = build_test_runtime().await; + let _ = node.mgr.dirty.reset(); assert!(!node.mgr.dirty.peek()); @@ -431,7 +456,7 @@ mod tests { async move { node.run(token).await } }); - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; node.global_ctx .issue_event(GlobalCtxEvent::DhcpIpv4Changed(None, None)); @@ -447,7 +472,8 @@ mod tests { #[tokio::test] #[serial_test::serial(dns_node_rpc_addr)] async fn run_marks_dirty_on_config_patched_event() { - let node = build_test_node().await; + let node = build_test_runtime().await; + let _ = node.mgr.dirty.reset(); assert!(!node.mgr.dirty.peek()); @@ -458,7 +484,7 @@ mod tests { async move { node.run(token).await } }); - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; node.global_ctx .issue_event(GlobalCtxEvent::ConfigPatched(InstanceConfigPatch::default())); @@ -474,7 +500,8 @@ mod tests { #[tokio::test] #[serial_test::serial(dns_node_rpc_addr)] async fn run_peer_info_updated_non_self_does_not_mark_dirty() { - let node = build_test_node().await; + let node = build_test_runtime().await; + let _ = node.mgr.dirty.reset(); assert!(!node.mgr.dirty.peek()); @@ -485,10 +512,10 @@ mod tests { async move { node.run(token).await } }); - tokio::time::sleep(Duration::from_millis(50)).await; + sleep(Duration::from_millis(50)).await; node.global_ctx .issue_event(GlobalCtxEvent::PeerInfoUpdated(vec![u32::MAX])); - tokio::time::sleep(Duration::from_millis(200)).await; + sleep(Duration::from_millis(200)).await; assert!(!node.mgr.dirty.peek()); @@ -502,7 +529,8 @@ mod tests { #[tokio::test] #[serial_test::serial(dns_node_rpc_addr)] async fn run_heartbeat_error_notifies_election() { - let node = build_test_node().await; + let node = build_test_runtime().await; + let _ = node.mgr.dirty.reset(); let token = CancellationToken::new(); @@ -526,7 +554,8 @@ mod tests { #[tokio::test] async fn id_matches_global_ctx_id() { - let node = build_test_node().await; + let node = build_test_runtime().await; + assert_eq!(node.id(), node.global_ctx.get_id()); } @@ -540,28 +569,19 @@ mod tests { #[serial_test::serial(dns_node_rpc_addr)] async fn election_wins_and_sets_dns_server() { let node = build_test_node().await; - let token = CancellationToken::new(); - let handle = tokio::spawn({ - let node = node.clone(); - let token = token.clone(); - async move { node.run_election(token).await } - }); + let runtime = node.runtime.clone(); - node.elect.notify_one(); wait_for_condition( - async || node.global_ctx.dns_server().is_some(), + async || runtime.global_ctx.dns_server().is_some(), Duration::from_secs(2), ) .await; - token.cancel(); - tokio::time::timeout(Duration::from_secs(3), handle) - .await - .unwrap() - .unwrap(); + let global_ctx = runtime.global_ctx.clone(); + node.stop().await.unwrap(); wait_for_condition( - async || node.global_ctx.dns_server().is_none(), + async || global_ctx.dns_server().is_none(), Duration::from_secs(2), ) .await; @@ -572,22 +592,12 @@ mod tests { async fn election_loses_when_rpc_addr_is_occupied() { let holder = occupy_dns_rpc_addr().await; let node = build_test_node().await; - let token = CancellationToken::new(); - let handle = tokio::spawn({ - let node = node.clone(); - let token = token.clone(); - async move { node.run_election(token).await } - }); + let runtime = node.runtime.clone(); - node.elect.notify_one(); - tokio::time::sleep(Duration::from_millis(300)).await; - assert!(node.global_ctx.dns_server().is_none()); + sleep(Duration::from_millis(300)).await; + assert!(runtime.global_ctx.dns_server().is_none()); - token.cancel(); - tokio::time::timeout(Duration::from_secs(3), handle) - .await - .unwrap() - .unwrap(); + node.stop().await.unwrap(); drop(holder); } @@ -596,31 +606,19 @@ mod tests { async fn election_retries_after_losing_then_wins() { let holder = occupy_dns_rpc_addr().await; let node = build_test_node().await; - let token = CancellationToken::new(); - let handle = tokio::spawn({ - let node = node.clone(); - let token = token.clone(); - async move { node.run_election(token).await } - }); + let runtime = node.runtime.clone(); - node.elect.notify_one(); - tokio::time::sleep(Duration::from_millis(300)).await; - assert!(node.global_ctx.dns_server().is_none()); + sleep(Duration::from_millis(300)).await; + assert!(runtime.global_ctx.dns_server().is_none()); drop(holder); - tokio::time::sleep(Duration::from_millis(100)).await; - - node.elect.notify_one(); + sleep(DNS_NODE_RR_INTERVAL).await; wait_for_condition( - async || node.global_ctx.dns_server().is_some(), + async || runtime.global_ctx.dns_server().is_some(), Duration::from_secs(2), ) .await; - token.cancel(); - tokio::time::timeout(Duration::from_secs(3), handle) - .await - .unwrap() - .unwrap(); + node.stop().await.unwrap(); } } diff --git a/easytier/src/dns/peer_mgr.rs b/easytier/src/dns/peer_mgr.rs index 19146dc5..5a59394e 100644 --- a/easytier/src/dns/peer_mgr.rs +++ b/easytier/src/dns/peer_mgr.rs @@ -144,15 +144,17 @@ pub struct DnsPeerMgr(Arc); impl DnsPeerMgr { pub fn new(peer_mgr: Arc, global_ctx: ArcGlobalCtx) -> Self { - Self(Arc::new(DnsPeerMgrInner { + let this = Self(Arc::new(DnsPeerMgrInner { peers: Cache::builder().time_to_idle(DNS_PEER_TTI).build(), dirty: Default::default(), peer_mgr, global_ctx, - })) + })); + this.register(); + this } - pub fn register(&self) { + fn register(&self) { self.peer_mgr .get_peer_rpc_mgr() .rpc_server() @@ -163,7 +165,7 @@ impl DnsPeerMgr { ); } - pub fn unregister(&self) -> Option<()> { + fn unregister(&self) -> Option<()> { self.peer_mgr .get_peer_rpc_mgr() .rpc_server() @@ -812,24 +814,9 @@ 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 unregister_without_register_returns_none() { - let peer_mgr = create_peer_manager_with_zone( - "unregister-peer", - "unregister-zone.test", - Ipv4Addr::new(10, 1, 0, 2), - ) - .await; - let mgr = DnsPeerMgr::new(peer_mgr, get_mock_global_ctx()); - - assert!(mgr.unregister().is_none()); - } - #[tokio::test] async fn drop_triggers_unregister() { let peer_mgr = create_peer_manager_with_zone( diff --git a/easytier/src/dns/tests.rs b/easytier/src/dns/tests.rs index efe655e2..86e36253 100644 --- a/easytier/src/dns/tests.rs +++ b/easytier/src/dns/tests.rs @@ -74,18 +74,14 @@ 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)))); - let dns_node = DnsNode::new(peer_mgr, global_ctx, nic_ctx); - dns_node.start().expect("failed to start dns node"); - dns_node + DnsNode::new(peer_mgr, global_ctx, nic_ctx) } 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)); - let dns_node = DnsNode::new(peer_mgr, global_ctx, nic_ctx); - dns_node.start().expect("failed to start dns node"); - dns_node + DnsNode::new(peer_mgr, global_ctx, nic_ctx) } pub async fn prepare_env_from_config_str(config_str: &str) -> Arc { @@ -489,9 +485,7 @@ 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()); - dns_a.register(); - + let _dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); let dns_b = DnsPeerMgr::new(peer_b.clone(), peer_b.get_global_ctx()); dns_b.refresh(peer_a.my_peer_id()).await; @@ -534,9 +528,7 @@ disabled = true .await .expect("route should appear"); - let dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); - dns_a.register(); - + let _dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); let dns_b = DnsPeerMgr::new(peer_b.clone(), peer_b.get_global_ctx()); dns_b.refresh(peer_a.my_peer_id()).await; @@ -873,9 +865,7 @@ 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()); - dns_c.register(); - + let _dns_c = DnsPeerMgr::new(peer_c.clone(), peer_c.get_global_ctx()); let dns_a = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); dns_a.refresh(peer_c.my_peer_id()).await; @@ -906,9 +896,7 @@ 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()); - dns_b.register(); - + let _dns_b = DnsPeerMgr::new(peer_b.clone(), peer_b.get_global_ctx()); let dns_a_online = DnsPeerMgr::new(peer_a.clone(), peer_a.get_global_ctx()); dns_a_online.refresh(peer_b.my_peer_id()).await; assert!( diff --git a/easytier/src/instance/instance.rs b/easytier/src/instance/instance.rs index 9e334355..9dcb0b53 100644 --- a/easytier/src/instance/instance.rs +++ b/easytier/src/instance/instance.rs @@ -472,7 +472,7 @@ pub struct Instance { #[cfg(feature = "tun")] nic_ctx: ArcNicCtx, #[cfg(feature = "magic-dns")] - dns: Arc, + dns: Option, peer_packet_receiver: Arc>, peer_manager: Arc, @@ -558,14 +558,6 @@ impl Instance { #[cfg(feature = "tun")] let nic_ctx = Arc::new(Mutex::new(None)); - #[cfg(feature = "magic-dns")] - let dns = Arc::new(DnsNode::new( - peer_manager.clone(), - global_ctx.clone(), - #[cfg(feature = "tun")] - nic_ctx.clone(), - )); - Instance { inst_name: global_ctx.inst_name.clone(), id, @@ -575,7 +567,7 @@ impl Instance { #[cfg(feature = "tun")] nic_ctx, #[cfg(feature = "magic-dns")] - dns, + dns: None, peer_manager, listener_manager, @@ -848,7 +840,14 @@ impl Instance { } #[cfg(feature = "magic-dns")] - self.dns.start()?; + { + self.dns = Some(DnsNode::new( + self.get_peer_manager(), + self.get_global_ctx(), + #[cfg(feature = "tun")] + self.get_nic_ctx(), + )); + } if self.global_ctx.config.get_dhcp() { self.check_dhcp_ip_conflict(); @@ -1396,9 +1395,9 @@ impl Instance { pub async fn clear_resources(&mut self) { self.peer_manager.clear_resources().await; #[cfg(feature = "magic-dns")] - self.dns.stop().await.unwrap_or_else(|e| { - tracing::error!("failed to stop dns, err: {:?}", e); - }); + if let Some(node) = self.dns.take() { + let _ = node.stop().await; + } #[cfg(feature = "tun")] let _ = self.nic_ctx.lock().await.take(); } @@ -1410,18 +1409,10 @@ impl Drop for Instance { let my_peer_id = self.peer_manager.my_peer_id(); let pm = Arc::downgrade(&self.peer_manager); #[cfg(feature = "magic-dns")] - let dns = self.dns.clone(); + let _ = self.dns.take(); // force abort #[cfg(feature = "tun")] let nic_ctx = self.nic_ctx.clone(); tokio::spawn(async move { - // TODO: change this - #[cfg(feature = "magic-dns")] - { - dns.stop().await.unwrap_or_else(|e| { - tracing::error!("failed to stop dns, err: {:?}", e); - }); - drop(dns); - } #[cfg(feature = "tun")] nic_ctx.lock().await.take(); if let Some(pm) = pm.upgrade() {