use CancellableTask in DnsNode

remove AsyncRuntime from DnsNode

node

n

n

clippy
This commit is contained in:
Luna Yao
2026-04-18 14:51:04 +02:00
parent d5995ea1cb
commit 4f6f8f3d93
4 changed files with 132 additions and 168 deletions
+106 -108
View File
@@ -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<DnsPeerMgr>,
#[cfg(feature = "tun")]
@@ -33,44 +32,13 @@ pub struct DnsNode {
global_ctx: ArcGlobalCtx,
elect: Arc<Notify>,
runtime: AsyncRuntime,
}
impl DnsNode {
pub fn new(
peer_mgr: Arc<PeerManager>,
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<PeerManager>,
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<TcpTunnelListener> {
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();
}
}
+6 -19
View File
@@ -144,15 +144,17 @@ pub struct DnsPeerMgr(Arc<DnsPeerMgrInner>);
impl DnsPeerMgr {
pub fn new(peer_mgr: Arc<PeerManager>, 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(
+6 -18
View File
@@ -74,18 +74,14 @@ 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))));
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<PeerManager>) -> 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<PeerManager> {
@@ -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!(
+14 -23
View File
@@ -472,7 +472,7 @@ pub struct Instance {
#[cfg(feature = "tun")]
nic_ctx: ArcNicCtx,
#[cfg(feature = "magic-dns")]
dns: Arc<DnsNode>,
dns: Option<DnsNode>,
peer_packet_receiver: Arc<Mutex<PacketRecvChanReceiver>>,
peer_manager: Arc<PeerManager>,
@@ -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() {