Skip to main content

file_engine/analysis/
handle.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4use std::time::{Duration, Instant};
5
6use tokio::task::JoinHandle;
7use tokio_stream::wrappers::UnboundedReceiverStream;
8use tokio_stream::Stream;
9use tokio_util::sync::CancellationToken;
10
11use crate::error::Result;
12
13use super::progress::AnalysisProgress;
14use super::report::AnalysisReport;
15
16/// Structurally similar to `Handle<T>` (event stream + `.cancel()` +
17/// awaitable) but a distinct type, following the same precedent
18/// `WatchHandle` already sets: `Handle<T>`'s progress stream is hard-coded
19/// to `crate::progress::Progress`, which lives behind the `operations`
20/// feature — `analyze` doesn't require `operations`, and `AnalysisProgress`
21/// isn't `Progress`, so genericizing the shared `Handle<T>` would mean
22/// either coupling `analyze` to `operations` or threading a second type
23/// parameter through every existing call site for no benefit to them.
24pub struct AnalysisHandle {
25    join_handle: JoinHandle<Result<AnalysisReport>>,
26    progress: UnboundedReceiverStream<AnalysisProgress>,
27    cancel: CancellationToken,
28    started: Instant,
29    finished: Option<Instant>,
30}
31
32impl AnalysisHandle {
33    pub(crate) fn new(
34        join_handle: JoinHandle<Result<AnalysisReport>>,
35        progress: tokio::sync::mpsc::UnboundedReceiver<AnalysisProgress>,
36        cancel: CancellationToken,
37    ) -> Self {
38        Self {
39            join_handle,
40            progress: UnboundedReceiverStream::new(progress),
41            cancel,
42            started: Instant::now(),
43            finished: None,
44        }
45    }
46
47    pub fn progress(&mut self) -> &mut (impl Stream<Item = AnalysisProgress> + Unpin) {
48        &mut self.progress
49    }
50
51    /// Wall time since the analysis was spawned. Frozen once the handle
52    /// has been polled to completion — see `Handle::elapsed`.
53    pub fn elapsed(&self) -> Duration {
54        match self.finished {
55            Some(finished) => finished.saturating_duration_since(self.started),
56            None => self.started.elapsed(),
57        }
58    }
59
60    /// Cooperative. Dropping the `AnalysisHandle` without calling this
61    /// keeps the walk running to completion.
62    pub fn cancel(&self) {
63        self.cancel.cancel();
64    }
65}
66
67impl Future for AnalysisHandle {
68    type Output = Result<AnalysisReport>;
69
70    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
71        let this = self.get_mut();
72        match Pin::new(&mut this.join_handle).poll(cx) {
73            Poll::Ready(Ok(result)) => {
74                this.finished.get_or_insert_with(Instant::now);
75                Poll::Ready(result)
76            }
77            Poll::Ready(Err(join_err)) => std::panic::resume_unwind(join_err.into_panic()),
78            Poll::Pending => Poll::Pending,
79        }
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use tokio::sync::mpsc;
86
87    use super::*;
88
89    #[tokio::test]
90    async fn awaiting_yields_the_wrapped_ok_value() {
91        let (_tx, rx) = mpsc::unbounded_channel();
92        let join_handle = tokio::spawn(async {
93            Ok(AnalysisReport {
94                file_count: 1,
95                dir_count: 0,
96                total_size: 0,
97                largest_files: Vec::new(),
98                by_extension: Default::default(),
99                by_mime: Default::default(),
100                age_buckets: Default::default(),
101                errors: Vec::new(),
102                errors_total: 0,
103                #[cfg(feature = "checksum")]
104                duplicates: Vec::new(),
105                #[cfg(feature = "checksum")]
106                duplicate_groups_total: 0,
107                #[cfg(feature = "checksum")]
108                duplicate_bytes_wasted: 0,
109                duration: Duration::ZERO,
110            })
111        });
112        let handle = AnalysisHandle::new(join_handle, rx, CancellationToken::new());
113
114        assert_eq!(handle.await.unwrap().file_count, 1);
115    }
116
117    #[tokio::test]
118    async fn cancel_triggers_the_wrapped_cancellation_token() {
119        let (_tx, rx) = mpsc::unbounded_channel();
120        let cancel = CancellationToken::new();
121        let join_handle = tokio::spawn(async { Err(crate::error::Error::Cancelled) });
122        let handle = AnalysisHandle::new(join_handle, rx, cancel.clone());
123
124        handle.cancel();
125        assert!(cancel.is_cancelled());
126        assert!(matches!(handle.await, Err(crate::error::Error::Cancelled)));
127    }
128}