hojicha_runtime/
async_handle.rs

1//! Handle for managing cancellable async operations
2
3//! Async operation handle with cancellation support
4//!
5//! This module provides `AsyncHandle<T>` for managing long-running async operations
6//! with cooperative cancellation support. This is essential for building responsive
7//! TUI applications that need to cancel operations when users navigate away or quit.
8//!
9//! # Example
10//!
11//! ```ignore
12//! use hojicha_runtime::async_handle::AsyncHandle;
13//!
14//! // Spawn a cancellable operation
15//! let handle = program.spawn_cancellable(|token| async move {
16//!     loop {
17//!         tokio::select! {
18//!             _ = token.cancelled() => {
19//!                 // Clean up and exit
20//!                 return Ok("Cancelled");
21//!             }
22//!             result = fetch_data() => {
23//!                 return Ok(result);
24//!             }
25//!         }
26//!     }
27//! });
28//!
29//! // Later, cancel if needed
30//! handle.cancel().await;
31//! ```
32
33use tokio::task::JoinHandle;
34use tokio_util::sync::CancellationToken;
35
36/// A handle to a cancellable async operation
37///
38/// `AsyncHandle` allows you to:
39/// - Cancel long-running operations cooperatively
40/// - Check if an operation is still running
41/// - Wait for completion with `.await`
42/// - Abort forcefully if needed
43///
44/// The handle automatically cancels the operation when dropped.
45pub struct AsyncHandle<T> {
46    handle: JoinHandle<T>,
47    cancel_token: CancellationToken,
48}
49
50impl<T> AsyncHandle<T> {
51    /// Create a new async handle
52    pub(crate) fn new(handle: JoinHandle<T>, cancel_token: CancellationToken) -> Self {
53        Self {
54            handle,
55            cancel_token,
56        }
57    }
58
59    /// Cancel the operation
60    ///
61    /// This sends a cancellation signal to the async task. The task must
62    /// cooperatively check for cancellation to actually stop.
63    pub fn cancel(&self) {
64        self.cancel_token.cancel();
65    }
66
67    /// Check if the operation is cancelled
68    pub fn is_cancelled(&self) -> bool {
69        self.cancel_token.is_cancelled()
70    }
71
72    /// Check if the operation is still running
73    pub fn is_running(&self) -> bool {
74        !self.handle.is_finished() && !self.cancel_token.is_cancelled()
75    }
76
77    /// Check if the operation has finished
78    pub fn is_finished(&self) -> bool {
79        self.handle.is_finished()
80    }
81
82    /// Abort the task immediately
83    ///
84    /// This is more forceful than cancel() - it immediately aborts the task
85    /// without waiting for cooperative cancellation.
86    pub fn abort(&self) {
87        self.handle.abort();
88    }
89
90    /// Get the cancellation token for cooperative cancellation
91    ///
92    /// This can be cloned and passed to child tasks for hierarchical cancellation.
93    pub fn cancellation_token(&self) -> &CancellationToken {
94        &self.cancel_token
95    }
96}
97
98impl<T> Drop for AsyncHandle<T> {
99    fn drop(&mut self) {
100        // Cancel the operation when the handle is dropped
101        self.cancel_token.cancel();
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    use std::time::Duration;
109
110    #[tokio::test]
111    async fn test_async_handle_cancel() {
112        let token = CancellationToken::new();
113        let token_clone = token.clone();
114
115        let handle = tokio::spawn(async move {
116            loop {
117                if token_clone.is_cancelled() {
118                    return "Cancelled";
119                }
120                tokio::time::sleep(Duration::from_millis(10)).await;
121            }
122        });
123
124        let async_handle = AsyncHandle::new(handle, token);
125        assert!(async_handle.is_running());
126
127        async_handle.cancel();
128        assert!(async_handle.is_cancelled());
129
130        tokio::time::sleep(Duration::from_millis(50)).await;
131        assert!(async_handle.is_finished());
132    }
133
134    #[tokio::test]
135    async fn test_async_handle_drop_cancels() {
136        let token = CancellationToken::new();
137        let token_clone = token.clone();
138        let token_check = token.clone();
139
140        let handle = tokio::spawn(async move {
141            loop {
142                if token_clone.is_cancelled() {
143                    break;
144                }
145                tokio::time::sleep(Duration::from_millis(10)).await;
146            }
147        });
148
149        {
150            let _async_handle = AsyncHandle::new(handle, token);
151            // Handle dropped here
152        }
153
154        // Should be cancelled after drop
155        assert!(token_check.is_cancelled());
156    }
157
158    #[tokio::test]
159    async fn test_async_handle_abort() {
160        let token = CancellationToken::new();
161        let token_clone = token.clone();
162
163        let handle = tokio::spawn(async move {
164            loop {
165                if token_clone.is_cancelled() {
166                    return "Cancelled";
167                }
168                tokio::time::sleep(Duration::from_millis(100)).await;
169            }
170        });
171
172        let async_handle = AsyncHandle::new(handle, token);
173
174        // Abort immediately
175        async_handle.abort();
176
177        // Should finish quickly
178        tokio::time::sleep(Duration::from_millis(10)).await;
179        assert!(async_handle.is_finished());
180    }
181}