Skip to main content

drop_handle/
lib.rs

1//! `DropHandle` is a handle that aborts the task when dropped.
2//!
3//! The task will only be aborted when the last `DropHandle` is dropped, so you can clone it to keep the task alive.
4//! This is useful for tasks that should be automatically cleaned up when they are no longer needed, without having to manually call `abort()`.
5//!
6//! Example usage:
7//! ```
8//! use drop_handle::DropHandle;
9//! use tokio::time::{sleep, Duration};
10//!
11//! #[tokio::main]
12//! async fn main() {
13//!     let drop_handle: DropHandle = tokio::spawn(async {
14//!         loop {
15//!             println!("Task is running...");
16//!             sleep(Duration::from_secs(1)).await;
17//!         }
18//!     })
19//!     .into();
20//!     // The task will be automatically aborted when `drop_handle` goes out of scope.
21//! }
22//! ```
23
24use std::{ops::Deref, sync::Arc};
25use tokio::task::{AbortHandle, JoinHandle};
26use tracing::{debug, trace};
27
28#[cfg(test)]
29mod tests;
30
31/// A handle that aborts the task when dropped.
32///
33/// The task will only be aborted when the last `DropHandle` is dropped, so you can clone it to keep the task alive.
34/// This is useful for tasks that should be automatically cleaned up when they are no longer needed, without having to manually call `abort()`.
35///
36/// Example usage:
37/// ```
38/// use drop_handle::DropHandle;
39/// use tokio::time::{sleep, Duration};
40///
41/// #[tokio::main]
42/// async fn main() {
43///     let drop_handle: DropHandle = tokio::spawn(async {
44///         loop {
45///             println!("Task is running...");
46///             sleep(Duration::from_secs(1)).await;
47///         }
48///     })
49///     .into();
50///     // The task will be automatically aborted when `drop_handle` goes out of scope.
51/// }
52/// ```
53#[derive(Clone, Debug)]
54pub struct DropHandle(Arc<AbortHandle>);
55
56impl Deref for DropHandle {
57    type Target = AbortHandle;
58
59    fn deref(&self) -> &AbortHandle {
60        &self.0
61    }
62}
63
64impl From<AbortHandle> for DropHandle {
65    fn from(value: AbortHandle) -> Self {
66        debug!("create DropHandle for task {:?}", value.id());
67        Self(Arc::new(value))
68    }
69}
70
71impl<T> From<JoinHandle<T>> for DropHandle {
72    fn from(value: JoinHandle<T>) -> Self {
73        value.abort_handle().into()
74    }
75}
76
77/// When the last `DropHandle` is dropped, the task will be aborted.
78impl Drop for DropHandle {
79    fn drop(&mut self) {
80        let drop_counter = Arc::strong_count(&self.0);
81        trace!("DropHandle counter: {}", drop_counter);
82        if drop_counter <= 1 {
83            debug!("drop DropHandle: abort task {:?}", self.0.id());
84            self.abort();
85        }
86    }
87}