tokio_process_tools/
collector.rs

1use crate::output_stream::Next;
2use std::fmt::Debug;
3use std::future::Future;
4use std::pin::Pin;
5use thiserror::Error;
6use tokio::sync::oneshot::Sender;
7use tokio::task::JoinHandle;
8
9/// Errors that can occur when collecting stream data.
10#[derive(Debug, Error)]
11pub enum CollectorError {
12    /// The collector task could not be joined/terminated.
13    #[error("Failed to join/terminate the collector task over stream '{stream_name}': {source}")]
14    TaskJoin {
15        /// The name of the stream this collector operates on.
16        stream_name: &'static str,
17
18        /// The source error.
19        #[source]
20        source: tokio::task::JoinError,
21    },
22}
23
24/// A trait for types that can act as sinks for collected stream data.
25///
26/// This is automatically implemented for any type that is `Debug + Send + Sync + 'static`.
27pub trait Sink: Debug + Send + Sync + 'static {}
28
29impl<T> Sink for T where T: Debug + Send + Sync + 'static {}
30
31// NOTE: We use Pin<Box> here to force usage of Higher-Rank Trait Bounds (HRTBs).
32// The returned futures will most-likely capture the `&mut T`and are therefore poised
33// by its lifetime. Without the trait-object usage, this would not work.
34pub type AsyncCollectFn<'a> = Pin<Box<dyn Future<Output = Next> + Send + 'a>>;
35
36/// A collector for stream data, inspecting it chunk by chunk but also providing mutable access
37/// to a sink in which the data can be stored.
38///
39/// See the `collect_*` functions on `BroadcastOutputStream` and `SingleOutputStream`.
40///
41/// For proper cleanup, call
42/// - `wait()`, which waits for the collection task to complete.
43/// - `cancel()`, which sends a termination signal and then waits for the collection task to complete.
44///
45/// If not cleaned up, the termination signal will be sent when dropping this collector,
46/// but the task will be aborted (forceful, not waiting for its regular completion).
47pub struct Collector<S: Sink> {
48    /// The name of the stream this collector operates on.
49    pub(crate) stream_name: &'static str,
50
51    pub(crate) task: Option<JoinHandle<S>>,
52    pub(crate) task_termination_sender: Option<Sender<()>>,
53}
54
55impl<S: Sink> Collector<S> {
56    /// Checks if this task has finished.
57    pub fn is_finished(&self) -> bool {
58        self.task.as_ref().map(|t| t.is_finished()).unwrap_or(true)
59    }
60
61    /// Wait for the collector to terminate naturally.
62    ///
63    /// A collector will automatically terminate when either:
64    ///
65    /// 1. The underlying write-side of the stream is dropped.
66    /// 2. The underlying stream is closed (by receiving an EOF / final read of 0 bytes).
67    /// 3. The first `Next::Break` is observed.
68    ///
69    /// If none of these may occur in your case, this could/will hang forever!
70    pub async fn wait(mut self) -> Result<S, CollectorError> {
71        // Take the `task_termination_sender`. Let's make sure nobody can ever interfere with us
72        // waiting here. DO NOT drop it, or the task will terminate (at least if it also takes the
73        // receive-error as a signal to terminate)!
74        let tts = self.task_termination_sender.take();
75
76        let sink = self
77            .task
78            .take()
79            .expect("`task` to be present.")
80            .await
81            .map_err(|err| CollectorError::TaskJoin {
82                stream_name: self.stream_name,
83                source: err,
84            });
85
86        // Drop the termination sender, we don't need it. Task is now terminated.
87        drop(tts);
88
89        sink
90    }
91
92    /// Sends a cancellation event to the collector, letting it shut down.
93    pub async fn cancel(mut self) -> Result<S, CollectorError> {
94        // We ignore any potential error here.
95        // Sending may fail if the task is already terminated (for example, by reaching EOF),
96        // which in turn dropped the receiver end!
97        let _res = self
98            .task_termination_sender
99            .take()
100            .expect("`task_termination_sender` to be present.")
101            .send(());
102
103        self.wait().await
104    }
105}
106
107impl<S: Sink> Drop for Collector<S> {
108    fn drop(&mut self) {
109        if let Some(task_termination_sender) = self.task_termination_sender.take() {
110            // We ignore any potential error here.
111            // Sending may fail if the task is already terminated (for example, by reaching EOF),
112            // which in turn dropped the receiver end!
113            let _res = task_termination_sender.send(());
114        }
115        if let Some(task) = self.task.take() {
116            task.abort();
117        }
118    }
119}
tokio_process_tools/collector.rs

tokio_process_tools/
collector.rs
