diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index be5986d3..788016b5 100644 --- a/easytier/src/dns/server.rs +++ b/easytier/src/dns/server.rs @@ -12,9 +12,9 @@ use crate::proto::rpc_impl::standalone::StandAloneServer; use crate::tunnel::common::bind_socket; use crate::tunnel::packet_def::ZCPacket; use crate::tunnel::tcp::TcpTunnelListener; +use crate::utils::AsyncRuntime; use derivative::Derivative; use derive_more::{Deref, DerefMut, From, Into}; -use futures_util::StreamExt; use hickory_proto::rr::Record; use hickory_proto::serialize::binary::{BinDecodable, BinEncoder}; use hickory_proto::xfer::Protocol; @@ -117,43 +117,6 @@ impl ResponseHandler for Response { } } -struct DnsServerRuntime { - token: CancellationToken, - task: Option>, -} - -impl DnsServerRuntime { - fn start(mut server: ServerFuture) -> Self { - Self { - token: server.shutdown_token().clone(), - task: Some(tokio::spawn(async move { - server - .block_until_done() - .await - .unwrap_or_else(|e| tracing::error!("DNS server exited with error: {:?}", e)); - })), - } - } - - async fn stop(mut self) -> anyhow::Result<()> { - self.token.cancel(); - if let Some(task) = self.task.take() { - task.await?; - } - Ok(()) - } -} - -impl Drop for DnsServerRuntime { - fn drop(&mut self) { - self.token.cancel(); - if let Some(task) = self.task.take() { - task.abort(); - tracing::warn!("DNS server runtime is leaked"); - } - } -} - #[derive(Derivative)] #[derivative(Debug)] pub struct DnsServer { @@ -202,26 +165,35 @@ impl DnsServer { async fn reload_listeners( &self, listeners: impl IntoIterator, - runtime: &mut Option, + runtime: &mut Option, ) -> anyhow::Result<()> { - if let Some(old) = runtime.take() { - old.stop().await?; + if let Some(runtime) = runtime.as_ref() { + if let Some(Err(e)) = runtime.stop().await { + tracing::error!("failed to stop old DNS server runtime: {}", e); + } } - let mut new = ServerFuture::new(self.catalog.clone()); + let runtime = runtime.get_or_insert_default(); + + let mut server = ServerFuture::new(self.catalog.clone()); for listener in listeners { let addr = listener.addr; if let Err(e) = match listener.protocol { - Protocol::Udp => bind_socket(addr, None).map(|s| new.register_socket(s)), + Protocol::Udp => bind_socket(addr, None).map(|s| server.register_socket(s)), Protocol::Tcp => bind_socket(addr, None) - .map(|s| new.register_listener(s, DNS_SERVER_LISTENER_TCP_TIMEOUT)), + .map(|s| server.register_listener(s, DNS_SERVER_LISTENER_TCP_TIMEOUT)), _ => unimplemented!(), } { tracing::error!("failed to bind DNS server on {}: {:?}", addr, e); } } - runtime.replace(DnsServerRuntime::start(new)); + runtime.start(Some(server.shutdown_token().clone()), |_| async move { + server + .block_until_done() + .await + .unwrap_or_else(|e| tracing::error!("DNS server exited with error: {:?}", e)); + }); Ok(()) }