datum/stream/
completion.rs

1//! `StreamCompletion<T>` — the handle returned immediately at materialization —
2//! and `Cancellable`.
3//!
4//! `StreamCompletion` is `#[must_use]`: dropping it cancels the running stream
5//! (its `Drop` impl). `.wait()` blocks until the result is ready (via
6//! `block_on`); `.try_wait()` polls without blocking.
7
8use super::*;
9
10#[derive(Clone)]
11pub struct Cancellable {
12    pub(super) cancelled: Arc<AtomicBool>,
13    _keep_alive: Option<Arc<dyn Send + Sync>>,
14}
15
16impl Cancellable {
17    pub(super) fn new_with_keep_alive(keep_alive: Option<Arc<dyn Send + Sync>>) -> Self {
18        Self {
19            cancelled: Arc::new(AtomicBool::new(false)),
20            _keep_alive: keep_alive,
21        }
22    }
23
24    pub fn cancel(&self) -> bool {
25        !self.cancelled.swap(true, Ordering::SeqCst)
26    }
27
28    #[must_use]
29    pub fn is_cancelled(&self) -> bool {
30        self.cancelled.load(Ordering::SeqCst)
31    }
32}
33
34impl fmt::Debug for Cancellable {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.debug_struct("Cancellable").finish_non_exhaustive()
37    }
38}
39
40#[must_use = "dropping the StreamCompletion cancels the running stream; call .wait()/.try_wait() or keep it alive"]
41pub struct StreamCompletion<T> {
42    state: StreamCompletionState<T>,
43    cancel_on_drop: Option<StreamCancellation>,
44}
45
46enum StreamCompletionState<T> {
47    Ready(Option<StreamResult<T>>),
48    Receiver(oneshot::Receiver<StreamResult<T>>),
49}
50
51#[derive(Clone)]
52pub(crate) struct StreamCancellation {
53    cancelled: Arc<AtomicBool>,
54    worker: Arc<Mutex<Option<thread::Thread>>>,
55}
56
57pub(super) struct RegisteredStreamWorker {
58    cancellation: StreamCancellation,
59}
60
61impl StreamCancellation {
62    pub(super) fn new() -> Self {
63        Self {
64            cancelled: Arc::new(AtomicBool::new(false)),
65            worker: Arc::new(Mutex::new(None)),
66        }
67    }
68
69    pub(crate) fn for_external_completion() -> Self {
70        Self::new()
71    }
72
73    pub(crate) fn cancelled(&self) -> Arc<AtomicBool> {
74        Arc::clone(&self.cancelled)
75    }
76
77    pub(super) fn register_current_worker(&self) -> RegisteredStreamWorker {
78        *self
79            .worker
80            .lock()
81            .unwrap_or_else(|poison| poison.into_inner()) = Some(thread::current());
82        RegisteredStreamWorker {
83            cancellation: self.clone(),
84        }
85    }
86
87    fn cancel(&self) {
88        self.cancelled.store(true, Ordering::SeqCst);
89        let worker = self
90            .worker
91            .lock()
92            .unwrap_or_else(|poison| poison.into_inner())
93            .clone();
94        if let Some(worker) = worker {
95            worker.unpark();
96        }
97    }
98}
99
100impl Drop for RegisteredStreamWorker {
101    fn drop(&mut self) {
102        *self
103            .cancellation
104            .worker
105            .lock()
106            .unwrap_or_else(|poison| poison.into_inner()) = None;
107    }
108}
109
110impl<T> StreamCompletion<T> {
111    pub(crate) fn from_receiver(
112        receiver: oneshot::Receiver<StreamResult<T>>,
113        cancel_on_drop: Option<StreamCancellation>,
114    ) -> Self {
115        Self {
116            state: StreamCompletionState::Receiver(receiver),
117            cancel_on_drop,
118        }
119    }
120
121    pub(crate) fn ready(result: StreamResult<T>) -> Self {
122        Self {
123            state: StreamCompletionState::Ready(Some(result)),
124            cancel_on_drop: None,
125        }
126    }
127
128    pub fn wait(self) -> StreamResult<T> {
129        block_on(self)
130    }
131
132    #[must_use]
133    pub fn try_wait(&mut self) -> Option<StreamResult<T>> {
134        match &mut self.state {
135            StreamCompletionState::Ready(result) => result.take(),
136            StreamCompletionState::Receiver(receiver) => match receiver.try_recv() {
137                Ok(Some(result)) => Some(result),
138                Ok(None) => None,
139                Err(_) => Some(Err(StreamError::AbruptTermination)),
140            },
141        }
142    }
143}
144
145impl<T> Future for StreamCompletion<T> {
146    type Output = StreamResult<T>;
147
148    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
149        match &mut self.state {
150            StreamCompletionState::Ready(result) => {
151                Poll::Ready(result.take().unwrap_or(Err(StreamError::AbruptTermination)))
152            }
153            StreamCompletionState::Receiver(receiver) => match Pin::new(receiver).poll(cx) {
154                Poll::Ready(Ok(result)) => Poll::Ready(result),
155                Poll::Ready(Err(_)) => Poll::Ready(Err(StreamError::AbruptTermination)),
156                Poll::Pending => Poll::Pending,
157            },
158        }
159    }
160}
161
162impl<T> Unpin for StreamCompletion<T> {}
163
164impl<T> Drop for StreamCompletion<T> {
165    fn drop(&mut self) {
166        if let Some(cancellation) = &self.cancel_on_drop {
167            cancellation.cancel();
168        }
169    }
170}
171
172impl<T> fmt::Debug for StreamCompletion<T> {
173    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
174        f.debug_struct("StreamCompletion").finish_non_exhaustive()
175    }
176}
datum/stream/completion.rs

datum/stream/
completion.rs
