refactor: remove ScopedTask (#2125)

* replace ScopedTask with AbortOnDropHandle

easytier/src/peers/acl_filter.rs

@@ -13,7 +13,6 @@ use pnet::packet::{
     Packet as _, ip::IpNextHeaderProtocols, ipv4::Ipv4Packet, tcp::TcpPacket, udp::UdpPacket,
 };
 
-use crate::common::scoped_task::ScopedTask;
 use crate::proto::acl::{AclStats, Protocol};
 use crate::tunnel::packet_def::PacketType;
 use crate::{
@@ -21,6 +20,7 @@ use crate::{
     proto::acl::{Acl, Action, ChainType},
     tunnel::packet_def::ZCPacket,
 };
+use tokio_util::task::AbortOnDropHandle;
 
 #[derive(Debug, Eq, PartialEq, Hash)]
 struct OutboundAllowRecord {
@@ -63,7 +63,7 @@ pub struct AclFilter {
     // Track allowed outbound packets and automatically allow their corresponding inbound response
     // packets, even if they would normally be dropped by ACL rules
     outbound_allow_records: Arc<DashMap<OutboundAllowRecord, Instant>>,
-    clean_task: ScopedTask<()>,
+    clean_task: AbortOnDropHandle<()>,
 }
 
 impl Default for AclFilter {
@@ -80,14 +80,13 @@ impl AclFilter {
             acl_processor: ArcSwap::from(Arc::new(AclProcessor::new(Acl::default()))),
             acl_enabled: Arc::new(AtomicBool::new(false)),
             outbound_allow_records,
-            clean_task: tokio::spawn(async move {
+            clean_task: AbortOnDropHandle::new(tokio::spawn(async move {
                 let max_life = std::time::Duration::from_secs(30);
                 loop {
                     record_clone.retain(|_, v| v.elapsed() < max_life);
                     tokio::time::sleep(std::time::Duration::from_secs(30)).await;
                 }
-            })
-            .into(),
+            })),
         }
     }
 

easytier/src/peers/foreign_network_client.rs

@@ -1,9 +1,10 @@
 use std::sync::{Arc, Mutex};
 
 use crate::{
-    common::{PeerId, error::Error, global_ctx::ArcGlobalCtx, scoped_task::ScopedTask},
+    common::{PeerId, error::Error, global_ctx::ArcGlobalCtx},
     tunnel::packet_def::ZCPacket,
 };
+use tokio_util::task::AbortOnDropHandle;
 
 use super::{PacketRecvChan, peer_conn::PeerConn, peer_map::PeerMap, peer_rpc::PeerRpcManager};
 
@@ -13,7 +14,7 @@ pub struct ForeignNetworkClient {
     my_peer_id: PeerId,
 
     peer_map: Arc<PeerMap>,
-    task: Mutex<Option<ScopedTask<()>>>,
+    task: Mutex<Option<AbortOnDropHandle<()>>>,
 }
 
 impl ForeignNetworkClient {
@@ -82,18 +83,15 @@ impl ForeignNetworkClient {
 
     pub async fn run(&self) {
         let peer_map = Arc::downgrade(&self.peer_map);
-        *self.task.lock().unwrap() = Some(
-            tokio::spawn(async move {
-                loop {
-                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
-                    let Some(peer_map) = peer_map.upgrade() else {
-                        break;
-                    };
-                    peer_map.clean_peer_without_conn().await;
-                }
-            })
-            .into(),
-        );
+        *self.task.lock().unwrap() = Some(AbortOnDropHandle::new(tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
+                let Some(peer_map) = peer_map.upgrade() else {
+                    break;
+                };
+                peer_map.clean_peer_without_conn().await;
+            }
+        })));
     }
 
     pub fn get_peer_map(&self) -> Arc<PeerMap> {

easytier/src/peers/peer.rs

@@ -12,6 +12,7 @@ use super::{
     PacketRecvChan,
     peer_conn::{PeerConn, PeerConnId},
 };
+use crate::{common::shrink_dashmap, proto::api::instance::PeerConnInfo};
 use crate::{
     common::{
         PeerId,
@@ -21,10 +22,7 @@ use crate::{
     proto::peer_rpc::PeerIdentityType,
     tunnel::packet_def::ZCPacket,
 };
-use crate::{
-    common::{scoped_task::ScopedTask, shrink_dashmap},
-    proto::api::instance::PeerConnInfo,
-};
+use tokio_util::task::AbortOnDropHandle;
 
 type ArcPeerConn = Arc<PeerConn>;
 type ConnMap = Arc<DashMap<PeerConnId, ArcPeerConn>>;
@@ -37,14 +35,14 @@ pub struct Peer {
     packet_recv_chan: PacketRecvChan,
 
     close_event_sender: mpsc::Sender<PeerConnId>,
-    close_event_listener: ScopedTask<()>,
+    close_event_listener: AbortOnDropHandle<()>,
 
     shutdown_notifier: Arc<tokio::sync::Notify>,
 
     default_conn_id: Arc<AtomicCell<PeerConnId>>,
     peer_identity_type: Arc<AtomicCell<Option<PeerIdentityType>>>,
     peer_public_key: Arc<RwLock<Option<Vec<u8>>>>,
-    default_conn_id_clear_task: ScopedTask<()>,
+    default_conn_id_clear_task: AbortOnDropHandle<()>,
 }
 
 impl Peer {
@@ -64,7 +62,7 @@ impl Peer {
         let conns_copy = conns.clone();
         let shutdown_notifier_copy = shutdown_notifier.clone();
         let global_ctx_copy = global_ctx.clone();
-        let close_event_listener = tokio::spawn(
+        let close_event_listener = AbortOnDropHandle::new(tokio::spawn(
             async move {
                 loop {
                     select! {
@@ -103,14 +101,13 @@ impl Peer {
                 "peer_close_event_listener",
                 ?peer_node_id,
             )),
-        )
-        .into();
+        ));
 
         let default_conn_id = Arc::new(AtomicCell::new(PeerConnId::default()));
 
         let conns_copy = conns.clone();
         let default_conn_id_copy = default_conn_id.clone();
-        let default_conn_id_clear_task = ScopedTask::from(tokio::spawn(async move {
+        let default_conn_id_clear_task = AbortOnDropHandle::new(tokio::spawn(async move {
             loop {
                 tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                 if conns_copy.len() > 1 {

easytier/src/peers/peer_conn.rs

@@ -1606,7 +1606,6 @@ pub mod tests {
     use crate::common::global_ctx::GlobalCtx;
     use crate::common::global_ctx::tests::get_mock_global_ctx;
     use crate::common::new_peer_id;
-    use crate::common::scoped_task::ScopedTask;
     use crate::common::stats_manager::{LabelSet, LabelType, MetricName};
     use crate::peers::create_packet_recv_chan;
     use crate::peers::recv_packet_from_chan;
@@ -1614,6 +1613,7 @@ pub mod tests {
     use crate::tunnel::filter::PacketRecorderTunnelFilter;
     use crate::tunnel::filter::tests::DropSendTunnelFilter;
     use crate::tunnel::ring::create_ring_tunnel_pair;
+    use tokio_util::task::AbortOnDropHandle;
 
     pub fn set_secure_mode_cfg(global_ctx: &GlobalCtx, enabled: bool) {
         if !enabled {
@@ -2200,7 +2200,7 @@ pub mod tests {
         c_peer.start_recv_loop(create_packet_recv_chan().0).await;
 
         let throughput = c_peer.throughput.clone();
-        let _t = ScopedTask::from(tokio::spawn(async move {
+        let _t = AbortOnDropHandle::new(tokio::spawn(async move {
             // if not drop both, we mock some rx traffic for client peer to test pinger
             if drop_both {
                 return;

easytier/src/peers/peer_task.rs

@@ -11,8 +11,8 @@ use tokio::select;
 use tokio::sync::Notify;
 use tokio::task::JoinHandle;
 
-use crate::common::scoped_task::ScopedTask;
 use anyhow::Error;
+use tokio_util::task::AbortOnDropHandle;
 
 use super::peer_manager::PeerManager;
 
@@ -72,7 +72,7 @@ pub trait PeerTaskLauncher: Send + Sync + Clone + 'static {
 
 pub struct PeerTaskManager<Launcher: PeerTaskLauncher> {
     launcher: Launcher,
-    main_loop_task: Mutex<Option<ScopedTask<()>>>,
+    main_loop_task: Mutex<Option<AbortOnDropHandle<()>>>,
     run_signal: Arc<Notify>,
     external_signal: Option<Arc<ExternalTaskSignal>>,
     data: Launcher::Data,
@@ -105,13 +105,12 @@ where
     }
 
     pub fn start(&self) {
-        let task = tokio::spawn(Self::main_loop(
+        let task = AbortOnDropHandle::new(tokio::spawn(Self::main_loop(
             self.launcher.clone(),
             self.data.clone(),
             self.run_signal.clone(),
             self.external_signal.clone(),
-        ))
-        .into();
+        )));
         self.main_loop_task.lock().unwrap().replace(task);
     }
 
@@ -121,7 +120,7 @@ where
         signal: Arc<Notify>,
         external_signal: Option<Arc<ExternalTaskSignal>>,
     ) {
-        let peer_task_map = Arc::new(DashMap::<C, ScopedTask<Result<T, Error>>>::new());
+        let peer_task_map = Arc::new(DashMap::<C, AbortOnDropHandle<Result<T, Error>>>::new());
         let mut external_signal_version = external_signal.as_ref().map(|signal| signal.version());
 
         loop {
@@ -158,8 +157,10 @@ where
                     }
 
                     tracing::debug!(?item, "launch hole punching task");
-                    peer_task_map
-                        .insert(item.clone(), launcher.launch_task(&data, item).await.into());
+                    peer_task_map.insert(
+                        item.clone(),
+                        AbortOnDropHandle::new(launcher.launch_task(&data, item).await),
+                    );
                 }
             } else if peer_task_map.is_empty() {
                 launcher.all_task_done(&data).await;