vtcode_commons/
task_guard.rs

1//! Owned tokio task handle with abort-on-drop semantics.
2//!
3//! Consolidates the ad-hoc `BackgroundTaskGuard`, `SignalHandlerGuard`, and
4//! `ProgressUpdateGuard` shapes previously reimplemented per crate: each held
5//! an `Option<JoinHandle<()>>` and aborted it on drop. (`TimeoutWarningGuard`
6//! is intentionally separate: it drives cooperative cancellation through a
7//! `CancellationToken` with an explicit async `cancel()`, not abort-on-drop.)
8//! Use this type for new owned background tasks instead of adding another
9//! local guard.
10//!
11//! # Usage
12//!
13//! ```no_run
14//! use vtcode_commons::TaskGuard;
15//!
16//! let guard = TaskGuard::with_label(tokio::spawn(async {}), "file-palette");
17//! // Dropping `guard` aborts the task; `guard.disarm()` releases the handle.
18//! ```
19//!
20//! Detached use remains possible via `disarm()` plus an explicit
21//! documented-detached comment at the call site (see "Task Extent, Error
22//! Propagation, and Cancel-Safety" in `docs/guides/async-architecture.md`).
23
24use tokio::task::JoinHandle;
25
26/// RAII owner for a spawned tokio task.
27///
28/// Dropping the guard aborts the task. Await or release the handle with
29/// [`Self::disarm`] when the task must outlive the current scope.
30#[must_use = "TaskGuard aborts the task when dropped; bind it to a local"]
31pub struct TaskGuard {
32    handle: Option<JoinHandle<()>>,
33    label: &'static str,
34}
35
36impl TaskGuard {
37    /// Own a spawned task with a default label.
38    pub fn new(handle: JoinHandle<()>) -> Self {
39        Self::with_label(handle, "task")
40    }
41
42    /// Own a spawned task with a static label used for debugging.
43    pub fn with_label(handle: JoinHandle<()>, label: &'static str) -> Self {
44        Self { handle: Some(handle), label }
45    }
46
47    /// Release the handle without aborting, for documented-detached handoff.
48    ///
49    /// Returns `None` only if the guard was already disarmed.
50    pub fn disarm(&mut self) -> Option<JoinHandle<()>> {
51        self.handle.take()
52    }
53
54    /// Whether the owned task has finished.
55    ///
56    /// Returns `false` once the guard is disarmed (no handle held), since the
57    /// task is no longer owned by this guard.
58    pub fn is_finished(&self) -> bool {
59        self.handle.as_ref().is_some_and(JoinHandle::is_finished)
60    }
61
62    /// Static label identifying the owned task.
63    pub fn label(&self) -> &'static str {
64        self.label
65    }
66}
67
68impl std::fmt::Debug for TaskGuard {
69    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        formatter
71            .debug_struct("TaskGuard")
72            .field("label", &self.label)
73            .field("finished", &self.is_finished())
74            .finish_non_exhaustive()
75    }
76}
77
78impl Drop for TaskGuard {
79    fn drop(&mut self) {
80        if let Some(handle) = self.handle.take() {
81            handle.abort();
82        }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use std::sync::Arc;
90    use std::sync::atomic::{AtomicBool, Ordering};
91    use std::time::Duration;
92
93    #[tokio::test]
94    async fn abort_on_drop_prevents_completion() {
95        let completed = Arc::new(AtomicBool::new(false));
96        let flag = Arc::clone(&completed);
97        let guard = TaskGuard::with_label(
98            tokio::spawn(async move {
99                tokio::time::sleep(Duration::from_millis(200)).await;
100                flag.store(true, Ordering::SeqCst);
101            }),
102            "abort-on-drop",
103        );
104        drop(guard);
105        tokio::time::sleep(Duration::from_millis(300)).await;
106        assert!(!completed.load(Ordering::SeqCst));
107    }
108
109    #[tokio::test]
110    async fn disarm_lets_task_complete() {
111        let completed = Arc::new(AtomicBool::new(false));
112        let flag = Arc::clone(&completed);
113        let mut guard = TaskGuard::with_label(
114            tokio::spawn(async move {
115                flag.store(true, Ordering::SeqCst);
116            }),
117            "disarm",
118        );
119        let handle = guard.disarm().expect("fresh guard holds a handle");
120        assert!(guard.disarm().is_none(), "second disarm releases nothing");
121        assert!(!guard.is_finished(), "disarmed guard owns nothing");
122        handle.await.expect("disarmed task runs to completion");
123        assert!(completed.load(Ordering::SeqCst));
124    }
125
126    #[tokio::test]
127    async fn is_finished_distinguishes_completed_and_pending_tasks() {
128        let completed = TaskGuard::new(tokio::spawn(async {}));
129        // Yield until the empty task completes (bounded wait).
130        for _ in 0..100 {
131            if completed.is_finished() {
132                break;
133            }
134            tokio::task::yield_now().await;
135        }
136        assert!(completed.is_finished());
137
138        let pending = TaskGuard::new(tokio::spawn(async {
139            tokio::time::sleep(Duration::from_secs(60)).await;
140        }));
141        assert!(!pending.is_finished());
142    }
143
144    #[tokio::test]
145    async fn label_defaults_and_debug_render() {
146        let guard = TaskGuard::new(tokio::spawn(async {}));
147        assert_eq!(guard.label(), "task");
148        let rendered = format!("{:?}", guard);
149        assert!(rendered.contains("TaskGuard"));
150        assert!(rendered.contains("task"));
151    }
152}
vtcode_commons/task_guard.rs

vtcode_commons/
task_guard.rs
