file_engine/
handle.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use tokio::task::JoinHandle;
6use tokio_stream::wrappers::UnboundedReceiverStream;
7use tokio_stream::Stream;
8use tokio_util::sync::CancellationToken;
9
10use crate::error::Result;
11use crate::progress::Progress;
12
13/// See dev-docs/design/handle-progress.md. One generic type reused across
14/// every operation (`T` = `OperationOutcome` for copy/move/compress,
15/// `SyncOutcome` for sync) rather than a handle type per operation.
16pub struct Handle<T> {
17    join_handle: JoinHandle<Result<T>>,
18    progress: UnboundedReceiverStream<Progress>,
19    cancel: CancellationToken,
20}
21
22impl<T> Handle<T> {
23    pub(crate) fn new(
24        join_handle: JoinHandle<Result<T>>,
25        progress: tokio::sync::mpsc::UnboundedReceiver<Progress>,
26        cancel: CancellationToken,
27    ) -> Self {
28        Self { join_handle, progress: UnboundedReceiverStream::new(progress), cancel }
29    }
30
31    pub fn progress(&mut self) -> &mut (impl Stream<Item = Progress> + Unpin) {
32        &mut self.progress
33    }
34
35    /// Cooperative — see dev-docs/design/batching-engine.md's cancellation
36    /// contract. Detached-on-drop: dropping the `Handle` without calling
37    /// this keeps the operation running to completion; see
38    /// dev-docs/design/handle-progress.md, "Drop behavior".
39    pub fn cancel(&self) {
40        self.cancel.cancel();
41    }
42}
43
44impl<T> Future for Handle<T> {
45    type Output = Result<T>;
46
47    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
48        let this = self.get_mut();
49        match Pin::new(&mut this.join_handle).poll(cx) {
50            Poll::Ready(Ok(result)) => Poll::Ready(result),
51            // The task panicked (dispatcher.rs's own internal
52            // `.expect()`s already treat that as unrecoverable) or was
53            // aborted (never happens — nothing calls `.abort()` on this
54            // JoinHandle). Propagate rather than silently swallowing it.
55            Poll::Ready(Err(join_err)) => std::panic::resume_unwind(join_err.into_panic()),
56            Poll::Pending => Poll::Pending,
57        }
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use std::sync::{Arc, Mutex};
64    use std::time::Duration;
65
66    use tokio::sync::mpsc;
67
68    use super::*;
69
70    #[tokio::test]
71    async fn awaiting_yields_the_wrapped_ok_value() {
72        let (_tx, rx) = mpsc::unbounded_channel();
73        let join_handle = tokio::spawn(async { Ok::<_, crate::error::Error>(42) });
74        let handle = Handle::new(join_handle, rx, CancellationToken::new());
75
76        assert_eq!(handle.await.unwrap(), 42);
77    }
78
79    #[tokio::test]
80    async fn awaiting_yields_the_wrapped_err_value() {
81        let (_tx, rx) = mpsc::unbounded_channel();
82        let join_handle = tokio::spawn(async { Err::<i32, _>(crate::error::Error::Cancelled) });
83        let handle = Handle::new(join_handle, rx, CancellationToken::new());
84
85        assert!(matches!(handle.await, Err(crate::error::Error::Cancelled)));
86    }
87
88    #[tokio::test]
89    async fn progress_stream_yields_events_in_send_order_then_ends() {
90        let (tx, rx) = mpsc::unbounded_channel();
91        let reporter = crate::progress::ProgressReporter::new(tx);
92        let join_handle = tokio::spawn(async move {
93            reporter.send(Progress::Started { bytes_total: None, entries_total: 1 });
94            reporter.send(Progress::EntryStarted {
95                entry: test_entry(),
96            });
97            // Reporter (and its clones) drop here, closing the channel.
98            Ok::<_, crate::error::Error>(())
99        });
100        let mut handle = Handle::new(join_handle, rx, CancellationToken::new());
101
102        use tokio_stream::StreamExt;
103        assert!(matches!(handle.progress().next().await, Some(Progress::Started { .. })));
104        assert!(matches!(handle.progress().next().await, Some(Progress::EntryStarted { .. })));
105        assert!(handle.progress().next().await.is_none());
106
107        handle.await.unwrap();
108    }
109
110    #[tokio::test]
111    async fn cancel_triggers_the_wrapped_cancellation_token() {
112        let (_tx, rx) = mpsc::unbounded_channel();
113        let cancel = CancellationToken::new();
114        let join_handle = tokio::spawn(async { Ok::<_, crate::error::Error>(()) });
115        let handle = Handle::new(join_handle, rx, cancel.clone());
116
117        handle.cancel();
118        assert!(cancel.is_cancelled());
119    }
120
121    #[tokio::test]
122    async fn dropping_the_handle_does_not_stop_the_wrapped_task() {
123        let (_tx, rx) = mpsc::unbounded_channel();
124        let sentinel = Arc::new(Mutex::new(0));
125        let sentinel_for_task = Arc::clone(&sentinel);
126
127        let join_handle = tokio::spawn(async move {
128            *sentinel_for_task.lock().unwrap() = 1;
129            tokio::time::sleep(Duration::from_millis(30)).await;
130            *sentinel_for_task.lock().unwrap() = 2;
131            Ok::<_, crate::error::Error>(())
132        });
133        let handle = Handle::new(join_handle, rx, CancellationToken::new());
134
135        tokio::time::sleep(Duration::from_millis(10)).await;
136        assert_eq!(*sentinel.lock().unwrap(), 1);
137        drop(handle);
138
139        tokio::time::sleep(Duration::from_millis(50)).await;
140        assert_eq!(*sentinel.lock().unwrap(), 2, "task should have run to completion despite the handle being dropped");
141    }
142
143    fn test_entry() -> crate::profiler::Entry {
144        crate::profiler::Entry {
145            path: std::path::PathBuf::from("a"),
146            relative_path: std::path::PathBuf::from("a"),
147            size: 1,
148            modified: None,
149        }
150    }
151}
file_engine/handle.rs

file_engine/
handle.rs
