merge

easytier/src/dns/node.rs

@@ -13,6 +13,7 @@ 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::CancellableTask;
+use std::io;
 use std::sync::Arc;
 use tokio::sync::{Notify, broadcast};
 use tokio::task::{JoinError, JoinSet};
@@ -215,7 +216,7 @@ impl DnsNodeRuntime {
 #[derive(Debug)]
 pub struct DnsNode {
     runtime: DnsNodeRuntime,
-    task: CancellableTask,
+    task: CancellableTask<()>,
 }
 
 impl DnsNode {
@@ -244,7 +245,7 @@ impl DnsNode {
         Self { runtime, task }
     }
 
-    pub async fn stop(self) -> Result<(), JoinError> {
+    pub async fn stop(self) -> io::Result<()> {
         self.task.stop(None).await
     }
 }

easytier/src/dns/server.rs

@@ -169,7 +169,7 @@ impl DnsServer {
     async fn reload_listeners(
         &self,
         listeners: impl IntoIterator<Item = NameServerAddr>,
-        runtime: &mut Option<CancellableTask>,
+        runtime: &mut Option<CancellableTask<()>>,
     ) -> anyhow::Result<()> {
         let listeners = listeners.into_iter().collect();
 
@@ -204,8 +204,8 @@ impl DnsServer {
             }
         }
 
-        *runtime = Some(CancellableTask::with_token(
+        let token = server.shutdown_token().clone();
-            server.shutdown_token().clone(),
+        let handle = tokio::spawn(
             async move {
                 server
                     .block_until_done()
@@ -213,7 +213,9 @@ impl DnsServer {
                     .unwrap_or_else(|e| tracing::error!("DNS server exited with error: {:?}", e));
             }
             .instrument(tracing::info_span!("DNS server backend runtime")),
-        ));
+        );
+
+        *runtime = Some(CancellableTask::with_handle(token, handle));
 
         *self.listeners.write() = listeners;
 

easytier/src/instance/instance.rs

@@ -7,13 +7,6 @@ use std::sync::{Arc, Weak};
 #[cfg(feature = "tun")]
 use std::time::Duration;
 
-use anyhow::Context;
-use cidr::{IpCidr, Ipv4Inet};
-use futures::FutureExt;
-use tokio::sync::{Mutex, Notify};
-#[cfg(feature = "tun")]
-use tokio::{sync::oneshot, task::JoinSet};
-
 use crate::common::PeerId;
 use crate::common::acl_processor::AclRuleBuilder;
 use crate::common::config::ConfigLoader;
@@ -59,6 +52,13 @@ use crate::proto::rpc_types::controller::BaseController;
 use crate::rpc_service::InstanceRpcService;
 use crate::utils::weak_upgrade;
 use crate::vpn_portal::{self, VpnPortal};
+use anyhow::Context;
+use cidr::{IpCidr, Ipv4Inet};
+use futures::FutureExt;
+use tokio::sync::{Mutex, Notify};
+#[cfg(feature = "tun")]
+use tokio::{sync::oneshot, task::JoinSet};
+use tokio_util::task::AbortOnDropHandle;
 
 use super::listeners::ListenerManager;
 use super::public_ipv6_provider::{

easytier/src/utils/dns.rs

@@ -190,9 +190,8 @@ pub async fn socket_addrs(
 
 #[cfg(test)]
 mod tests {
-    use crate::defer;
-
     use super::*;
+    use guarden::defer;
 
     #[test]
     fn parse_matrix_cases() {