From a60f101bda7dbd99ae8ace74160da78d97f532c2 Mon Sep 17 00:00:00 2001 From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com> Date: Fri, 17 Apr 2026 17:45:43 +0200 Subject: [PATCH] merge --- easytier/src/dns/node.rs | 20 ++-- easytier/src/dns/node_mgr.rs | 12 +-- easytier/src/dns/peer_mgr.rs | 104 +++++++++++-------- easytier/src/dns/server.rs | 26 ++--- easytier/src/dns/tests.rs | 10 +- easytier/src/dns/utils/addr.rs | 2 +- easytier/src/dns/zone.rs | 4 +- easytier/src/instance/instance.rs | 2 +- easytier/src/instance/proxy_cidrs_monitor.rs | 22 ++-- easytier/src/peers/peer_ospf_route.rs | 28 ++--- easytier/src/proto/dns.rs | 4 +- easytier/src/tunnel/common.rs | 52 +--------- easytier/src/utils/task.rs | 19 ++-- 13 files changed, 137 insertions(+), 168 deletions(-) diff --git a/easytier/src/dns/node.rs b/easytier/src/dns/node.rs index fe91a034..50368e18 100644 --- a/easytier/src/dns/node.rs +++ b/easytier/src/dns/node.rs @@ -1,23 +1,23 @@ use crate::common::global_ctx::{ArcGlobalCtx, GlobalCtxEvent}; use crate::common::join_joinset_background; use crate::dns::config::{ - DnsGlobalCtxExt, DNS_NODE_RR_INTERVAL, DNS_SERVER_ELECTION_INTERVAL, DNS_SERVER_RPC_ADDR, + DNS_NODE_RR_INTERVAL, DNS_SERVER_ELECTION_INTERVAL, DNS_SERVER_RPC_ADDR, DnsGlobalCtxExt, }; use crate::dns::peer_mgr::DnsPeerMgr; use crate::dns::server::DnsServer; #[cfg(feature = "tun")] use crate::instance::instance::ArcNicCtx; -use crate::peers::peer_manager::PeerManager; use crate::peers::NicPacketFilter; +use crate::peers::peer_manager::PeerManager; 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::AsyncRuntime; +use crate::utils::task::AsyncRuntime; use std::sync::{Arc, Mutex}; -use tokio::sync::{broadcast, Notify}; -use tokio::task::JoinSet; -use tokio::time::{sleep, sleep_until, Instant}; +use tokio::sync::{Notify, broadcast}; +use tokio::task::{JoinError, JoinSet}; +use tokio::time::{Instant, sleep, sleep_until}; use tokio_util::sync::CancellationToken; use tracing::instrument; use uuid::Uuid; @@ -57,18 +57,18 @@ impl DnsNode { self.global_ctx.get_id() } - pub fn start(&self) { + 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) -> anyhow::Result<()> { - self.runtime.stop().await.unwrap_or(Ok(())) + pub async fn stop(&self) -> Result<(), JoinError> { + self.runtime.stop(None).await.unwrap_or(Ok(())) } #[instrument(skip_all, name = "DnsNode election loop")] diff --git a/easytier/src/dns/node_mgr.rs b/easytier/src/dns/node_mgr.rs index 2d9d62e5..aebe5c8e 100644 --- a/easytier/src/dns/node_mgr.rs +++ b/easytier/src/dns/node_mgr.rs @@ -6,7 +6,7 @@ use crate::proto::dns::DnsNodeMgrRpc; use crate::proto::dns::{DnsSnapshot, HeartbeatRequest, HeartbeatResponse}; use crate::proto::rpc_types; use crate::proto::rpc_types::controller::BaseController; -use crate::utils::DeterministicDigest; +use crate::proto::utils::TransientDigest; use anyhow::Error; use hickory_server::authority::Catalog; use itertools::Itertools; @@ -17,7 +17,7 @@ use uuid::Uuid; #[derive(Debug, Clone, Default)] struct DnsNodeInfo { - digest: Vec, + digest: [u8; 32], zones: ZoneGroup, addresses: HashSet, listeners: HashSet, @@ -29,7 +29,7 @@ impl TryFrom<&DnsSnapshot> for DnsNodeInfo { fn try_from(value: &DnsSnapshot) -> Result { Ok(Self { digest: value.digest(), - zones: (&value.zones).try_into()?, + zones: value.zones.as_slice().try_into()?, addresses: value .addresses .iter() @@ -156,7 +156,7 @@ impl DnsNodeMgrRpc for DnsNodeMgr { self.nodes .get(&id) .await - .is_none_or(|info| info.digest != input.digest) + .is_none_or(|info| input.digest != info.digest) }; Ok(HeartbeatResponse { resync }) @@ -173,9 +173,9 @@ mod tests { }; use crate::dns::utils::response::ResponseHandle; use hickory_proto::op::{Message, ResponseCode}; - use hickory_proto::rr::{rdata, RData, RecordType}; + use hickory_proto::rr::{RData, RecordType, rdata}; use std::net::Ipv4Addr; - use tokio::time::{sleep, Duration}; + use tokio::time::{Duration, sleep}; fn heartbeat_digest_only(id: Uuid, digest: Vec) -> HeartbeatRequest { HeartbeatRequest { diff --git a/easytier/src/dns/peer_mgr.rs b/easytier/src/dns/peer_mgr.rs index 6798cce0..46130e44 100644 --- a/easytier/src/dns/peer_mgr.rs +++ b/easytier/src/dns/peer_mgr.rs @@ -1,6 +1,6 @@ -use crate::common::global_ctx::ArcGlobalCtx; use crate::common::PeerId; -use crate::dns::config::{DnsExportConfig, DnsGlobalCtxExt, DNS_PEER_TTI}; +use crate::common::global_ctx::ArcGlobalCtx; +use crate::dns::config::{DNS_PEER_TTI, DnsExportConfig, DnsGlobalCtxExt}; use crate::dns::utils::dirty::DirtyFlag; use crate::dns::zone::ZoneGroup; use crate::peer_center::instance::PeerCenterPeerManagerTrait; @@ -12,7 +12,7 @@ use crate::proto::dns::{ }; use crate::proto::rpc_types; use crate::proto::rpc_types::controller::BaseController; -use crate::utils::DeterministicDigest; +use crate::proto::utils::TransientDigest; use anyhow::Context; use moka::future::Cache; use std::ops::Deref; @@ -21,7 +21,7 @@ use tracing::instrument; #[derive(Debug, Clone)] struct DnsPeerInfo { - digest: Vec, + digest: [u8; 32], zones: Vec, } @@ -29,7 +29,7 @@ impl TryFrom for DnsPeerInfo { type Error = anyhow::Error; fn try_from(value: DnsExportConfig) -> Result { - let _ = ZoneGroup::try_from(&value.zones)?; + let _ = ZoneGroup::try_from(value.zones.as_slice())?; Ok(Self { digest: value.digest(), zones: value.zones, @@ -86,7 +86,7 @@ impl DnsPeerMgrInner { .peers .get(&peer_id) .await - .is_some_and(|info| info.digest == route.dns) + .is_some_and(|info| route.dns == info.digest) { return; } @@ -201,7 +201,7 @@ mod tests { use crate::proto::dns::GetExportConfigRequest; use std::collections::HashSet; use std::net::Ipv4Addr; - use tokio::time::{sleep, Duration}; + use tokio::time::{Duration, sleep}; use uuid::Uuid; async fn create_peer_manager_with_zone( @@ -279,14 +279,18 @@ mod tests { .await; let snapshot = mgr.snapshot(); - assert!(snapshot - .zones - .iter() - .any(|z| z.origin.contains("peer-cache.test"))); - assert!(snapshot - .zones - .iter() - .any(|z| z.origin.contains("local-custom.test"))); + assert!( + snapshot + .zones + .iter() + .any(|z| z.origin.contains("peer-cache.test")) + ); + assert!( + snapshot + .zones + .iter() + .any(|z| z.origin.contains("local-custom.test")) + ); } #[tokio::test] @@ -527,10 +531,12 @@ mod tests { assert!(local_dns.dirty.peek()); let snapshot = local_dns.snapshot(); - assert!(snapshot - .zones - .iter() - .any(|z| z.origin.contains("remote-export.test"))); + assert!( + snapshot + .zones + .iter() + .any(|z| z.origin.contains("remote-export.test")) + ); } #[tokio::test] @@ -564,14 +570,18 @@ mod tests { local_dns.refresh(peer_a.my_peer_id()).await; let snapshot = local_dns.snapshot(); - assert!(snapshot - .zones - .iter() - .any(|z| z.origin.contains("remote-a.test"))); - assert!(!snapshot - .zones - .iter() - .any(|z| z.origin.contains("remote-b.test"))); + assert!( + snapshot + .zones + .iter() + .any(|z| z.origin.contains("remote-a.test")) + ); + assert!( + !snapshot + .zones + .iter() + .any(|z| z.origin.contains("remote-b.test")) + ); } #[tokio::test] @@ -719,10 +729,12 @@ mod tests { .get(&unchanged_id) .await .expect("unchanged peer cache should stay"); - assert!(unchanged_cache - .zones - .iter() - .any(|z| z.origin.contains("cached-unchanged.test"))); + assert!( + unchanged_cache + .zones + .iter() + .any(|z| z.origin.contains("cached-unchanged.test")) + ); } #[tokio::test] @@ -748,14 +760,18 @@ mod tests { .await; let before = mgr.snapshot(); - assert!(before - .zones - .iter() - .any(|z| z.origin.contains("cached-expire.test"))); - assert!(before - .zones - .iter() - .any(|z| z.origin.contains("local-tti.test"))); + assert!( + before + .zones + .iter() + .any(|z| z.origin.contains("cached-expire.test")) + ); + assert!( + before + .zones + .iter() + .any(|z| z.origin.contains("local-tti.test")) + ); let deadline = tokio::time::Instant::now() + DNS_PEER_TTI + Duration::from_secs(3); loop { @@ -765,10 +781,12 @@ mod tests { .iter() .any(|z| z.origin.contains("cached-expire.test")); if expired { - assert!(now_snapshot - .zones - .iter() - .any(|z| z.origin.contains("local-tti.test"))); + assert!( + now_snapshot + .zones + .iter() + .any(|z| z.origin.contains("local-tti.test")) + ); break; } diff --git a/easytier/src/dns/server.rs b/easytier/src/dns/server.rs index 900e148f..e5b9ee77 100644 --- a/easytier/src/dns/server.rs +++ b/easytier/src/dns/server.rs @@ -4,37 +4,37 @@ use crate::dns::system; use crate::dns::utils::addr::NameServerAddr; use crate::dns::utils::response::ResponseHandle; use crate::peer_center::instance::PeerCenterPeerManagerTrait; -use crate::peers::peer_manager::PeerManager; use crate::peers::NicPacketFilter; +use crate::peers::peer_manager::PeerManager; use crate::proto::dns::DnsNodeMgrRpcServer; 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 crate::utils::task::AsyncRuntime; use derivative::Derivative; use hickory_proto::serialize::binary::BinDecodable; use hickory_proto::xfer::Protocol; use hickory_server::authority::MessageRequest; use hickory_server::{ + ServerFuture, authority::Catalog, server::{Request, RequestHandler, ResponseHandler, ResponseInfo}, - ServerFuture, }; use parking_lot::RwLock; use pnet::packet::icmp::{IcmpTypes, MutableIcmpPacket}; use pnet::packet::ip::IpNextHeaderProtocols; use pnet::packet::ipv4::{Ipv4Packet, MutableIpv4Packet}; use pnet::packet::udp::{MutableUdpPacket, UdpPacket}; -use pnet::packet::{icmp, ipv4, udp, MutablePacket, Packet}; +use pnet::packet::{MutablePacket, Packet, icmp, ipv4, udp}; use std::collections::HashSet; use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4}; use std::{sync::Arc, time::Duration}; use tokio_util::sync::CancellationToken; -use tracing::{instrument, Instrument}; +use tracing::{Instrument, instrument}; #[cfg(feature = "tun")] use crate::instance::instance::{ArcNicCtx, NicCtx}; +use crate::tunnel::common::bind; #[derive(Clone)] struct DynamicCatalog { @@ -182,7 +182,7 @@ impl DnsServer { tracing::info!(?listeners, "reloading"); if let Some(runtime) = runtime.as_ref() { - if let Some(Err(error)) = runtime.stop().await { + if let Some(Err(error)) = runtime.stop(None).await { tracing::error!(?error, "failed to stop old DNS server runtime"); } } @@ -194,9 +194,11 @@ impl DnsServer { let addr = listener.addr; tracing::info!(?addr, "binding listener"); if let Err(error) = match listener.protocol { - Protocol::Udp => bind_socket(addr, None).map(|s| server.register_socket(s)), - Protocol::Tcp => bind_socket(addr, None) + Protocol::Tcp => bind() + .addr(addr) + .call() .map(|s| server.register_listener(s, DNS_SERVER_LISTENER_TCP_TIMEOUT)), + Protocol::Udp => bind().addr(addr).call().map(|s| server.register_socket(s)), _ => unimplemented!(), } { tracing::error!(?addr, ?error, "failed to bind listener"); @@ -294,7 +296,7 @@ impl DnsServer { } if let Some(runtime) = runtime.take() { - let _ = runtime.stop().await; + let _ = runtime.stop(None).await; } } } @@ -485,7 +487,7 @@ mod tests { use crate::proto::rpc_types::controller::BaseController; use hickory_client::client::{Client, ClientHandle}; use hickory_proto::op::{Message, MessageType, OpCode, Query}; - use hickory_proto::rr::{rdata, DNSClass, Name, RData, Record, RecordType}; + use hickory_proto::rr::{DNSClass, Name, RData, Record, RecordType, rdata}; use hickory_proto::runtime::TokioRuntimeProvider; use hickory_proto::serialize::binary::BinEncodable; use hickory_proto::udp::UdpClientStream; @@ -496,7 +498,7 @@ mod tests { use pnet::packet::ip::IpNextHeaderProtocols; use pnet::packet::ipv4::{Ipv4Packet, MutableIpv4Packet}; use pnet::packet::udp::{MutableUdpPacket, UdpPacket}; - use pnet::packet::{icmp, ipv4, udp, MutablePacket, Packet}; + use pnet::packet::{MutablePacket, Packet, icmp, ipv4, udp}; use std::net::{Ipv4Addr, SocketAddr}; use std::str::FromStr; use std::time::Duration; diff --git a/easytier/src/dns/tests.rs b/easytier/src/dns/tests.rs index a3317dd4..7357cf3b 100644 --- a/easytier/src/dns/tests.rs +++ b/easytier/src/dns/tests.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use crate::common::config::TomlConfigLoader; -use crate::common::global_ctx::tests::get_mock_global_ctx; use crate::common::global_ctx::GlobalCtx; +use crate::common::global_ctx::tests::get_mock_global_ctx; use crate::connector::udp_hole_punch::tests::replace_stun_info_collector; use crate::dns::config::zone::ZoneConfigInner; use crate::dns::node::DnsNode; @@ -616,9 +616,11 @@ async fn config_patch_reloads_listener_binding() { check_dns_record_at(old_addr, "listener-patch.mesh-test.", "10.144.150.11").await; let mut dns = peer.get_global_ctx().config.get_dns(); - dns.listeners = vec![format!("udp://127.0.0.1:{listener_new}") - .parse() - .expect("invalid listener")] + dns.listeners = vec![ + format!("udp://127.0.0.1:{listener_new}") + .parse() + .expect("invalid listener"), + ] .into(); peer.get_global_ctx().config.set_dns(Some(dns)); peer.get_global_ctx() diff --git a/easytier/src/dns/utils/addr.rs b/easytier/src/dns/utils/addr.rs index 443aba26..2db3d05a 100644 --- a/easytier/src/dns/utils/addr.rs +++ b/easytier/src/dns/utils/addr.rs @@ -1,7 +1,7 @@ use crate::dns::config::DNS_SUPPORTED_PROTOCOLS; use crate::proto; use crate::proto::utils::RepeatedMessageModel; -use anyhow::{anyhow, Error}; +use anyhow::{Error, anyhow}; use hickory_proto::xfer::Protocol; use hickory_resolver::config::{NameServerConfig, NameServerConfigGroup}; use serde_with::{DeserializeFromStr, SerializeDisplay}; diff --git a/easytier/src/dns/zone.rs b/easytier/src/dns/zone.rs index c6b87f04..2fcafe0c 100644 --- a/easytier/src/dns/zone.rs +++ b/easytier/src/dns/zone.rs @@ -167,9 +167,9 @@ mod tests { use crate::proto::common::Url; use crate::proto::dns::ZoneData; use hickory_proto::op::{Message, ResponseCode}; - use hickory_proto::rr::{rdata, RData, Record, RecordType, RrsetRecords}; - use hickory_server::authority::Catalog; + use hickory_proto::rr::{RData, Record, RecordType, RrsetRecords, rdata}; use hickory_server::ServerFuture; + use hickory_server::authority::Catalog; use std::net::{Ipv4Addr, SocketAddr}; use std::str::FromStr; use tokio::net::UdpSocket; diff --git a/easytier/src/instance/instance.rs b/easytier/src/instance/instance.rs index a32f1e04..9e334355 100644 --- a/easytier/src/instance/instance.rs +++ b/easytier/src/instance/instance.rs @@ -848,7 +848,7 @@ impl Instance { } #[cfg(feature = "magic-dns")] - self.dns.start(); + self.dns.start()?; if self.global_ctx.config.get_dhcp() { self.check_dhcp_ip_conflict(); diff --git a/easytier/src/instance/proxy_cidrs_monitor.rs b/easytier/src/instance/proxy_cidrs_monitor.rs index fe7e7878..796e6093 100644 --- a/easytier/src/instance/proxy_cidrs_monitor.rs +++ b/easytier/src/instance/proxy_cidrs_monitor.rs @@ -44,20 +44,20 @@ impl ProxyCidrsMonitor { proxy_cidrs.insert(vpn_cfg.client_cidr); } + #[cfg(feature = "magic-dns")] + { + use crate::dns::config::DnsGlobalCtxExt; + if let Some(dns) = global_ctx.dns_server() { + proxy_cidrs.extend(dns.addresses().into_iter().filter_map(|a| match a.ip() { + IpAddr::V4(ip) => Some(cidr::Ipv4Cidr::new_host(ip)), + _ => None, + })) + } + } + proxy_cidrs }; - #[cfg(feature = "magic-dns")] - { - use crate::dns::config::DnsGlobalCtxExt; - if let Some(dns) = global_ctx.dns_server() { - proxy_cidrs.extend(dns.addresses().into_iter().filter_map(|a| match a.ip() { - IpAddr::V4(ip) => Some(cidr::Ipv4Cidr::new_host(ip)), - _ => None, - })) - } - } - // Calculate diff if cur_proxy_cidrs == &proxy_cidrs { return (proxy_cidrs, Vec::new(), Vec::new()); diff --git a/easytier/src/peers/peer_ospf_route.rs b/easytier/src/peers/peer_ospf_route.rs index cf1d3d47..122e57a3 100644 --- a/easytier/src/peers/peer_ospf_route.rs +++ b/easytier/src/peers/peer_ospf_route.rs @@ -30,17 +30,8 @@ use tokio::{ task::{JoinHandle, JoinSet}, }; -use super::{ - graph_algo::dijkstra_with_first_hop, - peer_rpc::PeerRpcManager, - route_trait::{ - DefaultRouteCostCalculator, ForeignNetworkRouteInfoMap, NextHopPolicy, RouteCostCalculator, - RouteCostCalculatorInterface, - }, - PeerPacketFilter, -}; use crate::common::config::ConfigLoader; -use crate::utils::DeterministicDigest; +use crate::proto::utils::TransientDigest; use crate::{ common::{ PeerId, @@ -228,14 +219,15 @@ impl RoutePeerInfo { .and_then(|cfg| cfg.public_key().ok()) .map(|pk| pk.as_bytes().to_vec()) .unwrap_or_default(); - cfg_if! { - if #[cfg(feature = "magic-dns")] { + + let dns = cfg_select! { + feature = "magic-dns" => {{ use crate::dns::config::DnsGlobalCtxExt; - let dns = global_ctx.dns_export_config().digest(); - } else { - let dns = Default::default(); - } - } + global_ctx.dns_export_config().digest() + }} + _ => Default::default(), + }; + Self { peer_id: my_peer_id, inst_id: Some(global_ctx.get_id().into()), @@ -268,7 +260,7 @@ impl RoutePeerInfo { ipv6_addr: global_ctx.get_ipv6().map(|x| x.into()), groups: global_ctx.get_acl_groups(my_peer_id), - dns, + dns: dns.into(), noise_static_pubkey, diff --git a/easytier/src/proto/dns.rs b/easytier/src/proto/dns.rs index d1526dd8..253cc3a9 100644 --- a/easytier/src/proto/dns.rs +++ b/easytier/src/proto/dns.rs @@ -1,11 +1,11 @@ -use crate::utils::DeterministicDigest; +use crate::proto::utils::TransientDigest; use std::fmt::Display; include!(concat!(env!("OUT_DIR"), "/dns.rs")); impl HeartbeatRequest { pub fn update(&mut self, snapshot: DnsSnapshot) { - self.digest = snapshot.digest(); + self.digest = snapshot.digest().into(); self.snapshot = Some(snapshot); } } diff --git a/easytier/src/tunnel/common.rs b/easytier/src/tunnel/common.rs index 2543b56e..f73e53e7 100644 --- a/easytier/src/tunnel/common.rs +++ b/easytier/src/tunnel/common.rs @@ -17,7 +17,7 @@ use super::{ buf::BufList, packet_def::{TCP_TUNNEL_HEADER_SIZE, TCPTunnelHeader, ZCPacketType}, }; -use crate::common::error::Error; +use crate::common::netns::NetNS; use crate::tunnel::packet_def::{PEER_MANAGER_HEADER_SIZE, ZCPacket}; use bytes::{Buf, BufMut, Bytes, BytesMut}; use tokio::net::{TcpListener, TcpSocket, UdpSocket}; @@ -535,56 +535,6 @@ pub fn bind( B::finalize(socket) } -pub trait Bindable: Sized { - const TY: socket2::Type; - const PROTOCOL: Option; - - fn finalize(socket: socket2::Socket) -> Result; -} - -impl Bindable for TcpListener { - const TY: socket2::Type = socket2::Type::STREAM; - const PROTOCOL: Option = Some(socket2::Protocol::TCP); - - fn finalize(socket: socket2::Socket) -> Result { - let tcp_socket = TcpSocket::from_std_stream(socket.into()); - - if let Err(e) = tcp_socket.set_nodelay(true) { - tracing::warn!(?e, "set_nodelay fail in listen"); - } - - Ok(tcp_socket.listen(1024)?) - } -} - -impl Bindable for UdpSocket { - const TY: socket2::Type = socket2::Type::DGRAM; - const PROTOCOL: Option = Some(socket2::Protocol::UDP); - - fn finalize(socket: socket2::Socket) -> Result { - Ok(UdpSocket::from_std(socket.into())?) - } -} - -pub fn bind_socket(addr: SocketAddr, net_ns: Option) -> Result { - let _g = net_ns.map(|n| n.guard()); - - let socket2_socket = - socket2::Socket::new(socket2::Domain::for_address(addr), B::TY, B::PROTOCOL)?; - - setup_sokcet2(&socket2_socket, &addr)?; - - B::finalize(socket2_socket) -} - -pub fn bind_tcp_socket(addr: SocketAddr, net_ns: NetNS) -> Result { - bind_socket(addr, Some(net_ns)) -} - -pub fn bind_udp_socket(addr: SocketAddr, net_ns: NetNS) -> Result { - bind_socket(addr, Some(net_ns)) -} - pub fn reserve_buf(buf: &mut BytesMut, min_size: usize, max_size: usize) { if buf.capacity() < min_size { buf.reserve(max_size); diff --git a/easytier/src/utils/task.rs b/easytier/src/utils/task.rs index db25df64..cf436373 100644 --- a/easytier/src/utils/task.rs +++ b/easytier/src/utils/task.rs @@ -79,7 +79,7 @@ impl AsyncRuntime { Ok(()) } - pub async fn stop(&self, timeout: Duration) -> Option> { + pub async fn stop(&self, timeout: Option) -> Option> { let state = { let mut state = self.state.lock(); match &*state { @@ -104,12 +104,17 @@ impl AsyncRuntime { }; token.cancel(); - let result = if let Ok(result) = tokio::time::timeout(timeout, &mut task).await { - result - } else { - task.abort(); - tracing::warn!("task stop timeout after {:?}, aborted", timeout); - task.await + let result = match timeout { + Some(duration) => { + if let Ok(result) = tokio::time::timeout(duration, &mut task).await { + result + } else { + task.abort(); + tracing::warn!("task stop timeout after {:?}, aborted", duration); + task.await + } + } + None => task.await, }; {