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}