refactor: remove ScopedTask (#2125)

* replace ScopedTask with AbortOnDropHandle

easytier/src/common/mod.rs

@@ -25,7 +25,6 @@ pub mod log;
 pub mod netns;
 pub mod network;
 pub mod os_info;
-pub mod scoped_task;
 pub mod stats_manager;
 pub mod stun;
 pub mod stun_codec_ext;

easytier/src/common/scoped_task.rs

@@ -1,119 +0,0 @@
-//! This crate provides a wrapper type of Tokio's JoinHandle: `ScopedTask`, which aborts the task when it's dropped.
-//! `ScopedTask` can still be awaited to join the child-task, and abort-on-drop will still trigger while it is being awaited.
-//!
-//! For example, if task A spawned task B but is doing something else, and task B is waiting for task C to join,
-//! aborting A will also abort both B and C.
-
-use derive_more::{Deref, DerefMut, From};
-use std::future::Future;
-use std::pin::Pin;
-use std::task::{Context, Poll};
-use tokio::task::JoinHandle;
-
-#[derive(Debug, From, Deref, DerefMut)]
-pub struct ScopedTask<T>(JoinHandle<T>);
-
-impl<T> Drop for ScopedTask<T> {
-    fn drop(&mut self) {
-        self.abort()
-    }
-}
-
-impl<T> Future for ScopedTask<T> {
-    type Output = <JoinHandle<T> as Future>::Output;
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        Pin::new(&mut self.0).poll(cx)
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::ScopedTask;
-    use futures_util::future::pending;
-    use std::sync::{Arc, RwLock};
-    use tokio::task::yield_now;
-
-    struct Sentry(Arc<RwLock<bool>>);
-    impl Drop for Sentry {
-        fn drop(&mut self) {
-            *self.0.write().unwrap() = true
-        }
-    }
-
-    #[tokio::test]
-    async fn drop_while_not_waiting_for_join() {
-        let dropped = Arc::new(RwLock::new(false));
-        let sentry = Sentry(dropped.clone());
-        let task = ScopedTask::from(tokio::spawn(async move {
-            let _sentry = sentry;
-            pending::<()>().await
-        }));
-        yield_now().await;
-        assert!(!*dropped.read().unwrap());
-        drop(task);
-        yield_now().await;
-        assert!(*dropped.read().unwrap());
-    }
-
-    #[tokio::test]
-    async fn drop_while_waiting_for_join() {
-        let dropped = Arc::new(RwLock::new(false));
-        let sentry = Sentry(dropped.clone());
-        let handle = tokio::spawn(async move {
-            ScopedTask::from(tokio::spawn(async move {
-                let _sentry = sentry;
-                pending::<()>().await
-            }))
-            .await
-            .unwrap()
-        });
-        yield_now().await;
-        assert!(!*dropped.read().unwrap());
-        handle.abort();
-        yield_now().await;
-        assert!(*dropped.read().unwrap());
-    }
-
-    #[tokio::test]
-    async fn no_drop_only_join() {
-        assert_eq!(
-            ScopedTask::from(tokio::spawn(async {
-                yield_now().await;
-                5
-            }))
-            .await
-            .unwrap(),
-            5
-        )
-    }
-
-    #[tokio::test]
-    async fn manually_abort_before_drop() {
-        let dropped = Arc::new(RwLock::new(false));
-        let sentry = Sentry(dropped.clone());
-        let task = ScopedTask::from(tokio::spawn(async move {
-            let _sentry = sentry;
-            pending::<()>().await
-        }));
-        yield_now().await;
-        assert!(!*dropped.read().unwrap());
-        task.abort();
-        yield_now().await;
-        assert!(*dropped.read().unwrap());
-    }
-
-    #[tokio::test]
-    async fn manually_abort_then_join() {
-        let dropped = Arc::new(RwLock::new(false));
-        let sentry = Sentry(dropped.clone());
-        let task = ScopedTask::from(tokio::spawn(async move {
-            let _sentry = sentry;
-            pending::<()>().await
-        }));
-        yield_now().await;
-        assert!(!*dropped.read().unwrap());
-        task.abort();
-        yield_now().await;
-        assert!(task.await.is_err());
-    }
-}

easytier/src/common/stats_manager.rs

@@ -5,8 +5,7 @@ use std::fmt;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 use tokio::time::interval;
-
-use crate::common::scoped_task::ScopedTask;
+use tokio_util::task::AbortOnDropHandle;
 
 /// Predefined metric names for type safety
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
@@ -578,7 +577,7 @@ impl MetricSnapshot {
 /// StatsManager manages global statistics with high performance counters
 pub struct StatsManager {
     counters: Arc<DashMap<MetricKey, Arc<MetricData>>>,
-    cleanup_task: ScopedTask<()>,
+    cleanup_task: AbortOnDropHandle<()>,
 }
 
 impl StatsManager {
@@ -611,7 +610,7 @@ impl StatsManager {
 
         Self {
             counters,
-            cleanup_task: cleanup_task.into(),
+            cleanup_task: AbortOnDropHandle::new(cleanup_task),
         }
     }
 

easytier/src/common/stun.rs

@@ -1343,11 +1343,9 @@ impl StunInfoCollectorTrait for MockStunInfoCollector {
 
 #[cfg(test)]
 mod tests {
-    use crate::{
-        common::scoped_task::ScopedTask,
-        tunnel::{TunnelListener, udp::UdpTunnelListener},
-    };
+    use crate::tunnel::{TunnelListener, udp::UdpTunnelListener};
     use tokio::time::{sleep, timeout};
+    use tokio_util::task::AbortOnDropHandle;
 
     use super::*;
 
@@ -1441,7 +1439,7 @@ mod tests {
         use stun_codec::rfc5389::attributes::XorMappedAddress;
         use tokio::net::TcpListener;
 
-        async fn spawn_tcp_stun_server() -> (SocketAddr, ScopedTask<()>) {
+        async fn spawn_tcp_stun_server() -> (SocketAddr, AbortOnDropHandle<()>) {
             let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
             let server_addr = listener.local_addr().unwrap();
 
@@ -1465,7 +1463,7 @@ mod tests {
                 stream.write_all(rsp_buf.as_slice()).await.unwrap();
             });
 
-            (server_addr, task.into())
+            (server_addr, AbortOnDropHandle::new(task))
         }
 
         let (server1, _t1) = spawn_tcp_stun_server().await;
@@ -1504,7 +1502,7 @@ mod tests {
         let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
         let server_addr = listener.local_addr().unwrap();
 
-        let _t = ScopedTask::from(tokio::spawn(async move {
+        let _t = AbortOnDropHandle::new(tokio::spawn(async move {
             for _ in 0..8 {
                 let Ok((mut stream, peer_addr)) = listener.accept().await else {
                     break;

easytier/src/common/token_bucket.rs

@@ -5,8 +5,8 @@ use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 use tokio::sync::Notify;
 use tokio::time;
+use tokio_util::task::AbortOnDropHandle;
 
-use crate::common::scoped_task::ScopedTask;
 use crate::proto::common::LimiterConfig;
 
 /// Token Bucket rate limiter using atomic operations
@@ -14,7 +14,7 @@ pub struct TokenBucket {
     available_tokens: AtomicU64, // Current token count (atomic)
     last_refill_time: AtomicU64, // Last refill time as micros since epoch
     config: BucketConfig,        // Immutable configuration
-    refill_task: Mutex<Option<ScopedTask<()>>>, // Background refill task
+    refill_task: Mutex<Option<AbortOnDropHandle<()>>>, // Background refill task
     start_time: Instant,         // Bucket creation time
 
     refill_notifier: Arc<Notify>,
@@ -91,7 +91,7 @@ impl TokenBucket {
             .refill_task
             .lock()
             .unwrap()
-            .replace(refill_task.into());
+            .replace(AbortOnDropHandle::new(refill_task));
         arc_self
     }
 
@@ -172,7 +172,7 @@ impl TokenBucket {
 pub struct TokenBucketManager {
     buckets: Arc<DashMap<String, Arc<TokenBucket>>>,
 
-    retain_task: ScopedTask<()>,
+    retain_task: AbortOnDropHandle<()>,
 }
 
 impl Default for TokenBucketManager {
@@ -205,7 +205,7 @@ impl TokenBucketManager {
 
         Self {
             buckets,
-            retain_task: retain_task.into(),
+            retain_task: AbortOnDropHandle::new(retain_task),
         }
     }