node & peer_mgr: fix rpc register/unregister

mut node
This commit is contained in:
Luna Yao
2026-04-28 15:52:31 +02:00
parent 8cf9d493b8
commit 46d1c2aef3
5 changed files with 71 additions and 75 deletions
+27 -12
View File
@@ -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 = {
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();
}
}
+6 -39
View File
@@ -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());
}
}
+1 -1
View File
@@ -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() {
+32 -20
View File
@@ -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 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
+5 -3
View File
@@ -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")]