file_engine/analysis/
handle.rs1use 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
16pub 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 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 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}