kube_runtime/utils/
mod.rs

1//! Helpers for manipulating built-in streams
2
3mod backoff_reset_timer;
4pub(crate) mod delayed_init;
5mod event_decode;
6mod event_modify;
7mod predicate;
8mod reflect;
9mod stream_backoff;
10mod watch_ext;
11
12pub use backoff_reset_timer::{Backoff, ResetTimerBackoff};
13pub use event_decode::EventDecode;
14pub use event_modify::EventModify;
15pub use predicate::{Config as PredicateConfig, Predicate, PredicateFilter, predicates};
16pub use reflect::Reflect;
17pub use stream_backoff::StreamBackoff;
18pub use watch_ext::WatchStreamExt;
19
20use futures::{
21    FutureExt, Stream, StreamExt, TryStream, TryStreamExt,
22    stream::{self, Peekable},
23};
24use pin_project::pin_project;
25use std::{
26    fmt::Debug,
27    pin::{Pin, pin},
28    sync::{Arc, Mutex},
29    task::Poll,
30};
31use stream::IntoStream;
32use tokio::{runtime::Handle, task::JoinHandle};
33
34/// Allows splitting a `Stream` into several streams that each emit a disjoint subset of the input stream's items,
35/// like a streaming variant of pattern matching.
36///
37/// NOTE: The cases MUST be reunited into the same final stream (using `futures::stream::select` or similar),
38/// since cases for rejected items will *not* register wakeup correctly, and may otherwise lose items and/or deadlock.
39///
40/// NOTE: The whole set of cases will deadlock if there is ever an item that no live case wants to consume.
41#[pin_project]
42pub(crate) struct SplitCase<S: Stream, Case> {
43    // Future-unaware `Mutex` is OK because it's only taken inside single poll()s
44    inner: Arc<Mutex<Peekable<S>>>,
45    /// Tests whether an item from the stream should be consumed
46    ///
47    /// NOTE: This MUST be total over all `SplitCase`s, otherwise the input stream
48    /// will get stuck deadlocked because no candidate tries to consume the item.
49    should_consume_item: fn(&S::Item) -> bool,
50    /// Narrows the type of the consumed type, using the same precondition as `should_consume_item`.
51    ///
52    /// NOTE: This MUST return `Some` if `should_consume_item` returns `true`, since we can't put
53    /// an item back into the input stream once consumed.
54    try_extract_item_case: fn(S::Item) -> Option<Case>,
55}
56
57impl<S, Case> Stream for SplitCase<S, Case>
58where
59    S: Stream + Unpin,
60    S::Item: Debug,
61{
62    type Item = Case;
63
64    fn poll_next(
65        self: std::pin::Pin<&mut Self>,
66        cx: &mut std::task::Context<'_>,
67    ) -> std::task::Poll<Option<Self::Item>> {
68        let this = self.project();
69        let inner = this.inner.lock().unwrap();
70        let mut inner = Pin::new(inner);
71        let inner_peek = pin!(inner.as_mut().peek());
72        match inner_peek.poll(cx) {
73            Poll::Ready(Some(x_ref)) => {
74                if (this.should_consume_item)(x_ref) {
75                    let item = inner.as_mut().poll_next(cx);
76                    match item {
77                        Poll::Ready(Some(x)) => Poll::Ready(Some((this.try_extract_item_case)(x).expect(
78                            "`try_extract_item_case` returned `None` despite `should_consume_item` returning `true`",
79                        ))),
80                        res => panic!(
81                    "Peekable::poll_next() returned {res:?} when Peekable::peek() returned Ready(Some(_))"
82                ),
83                    }
84                } else {
85                    // Handled by another SplitCase instead
86                    Poll::Pending
87                }
88            }
89            Poll::Ready(None) => Poll::Ready(None),
90            Poll::Pending => Poll::Pending,
91        }
92    }
93}
94
95/// Splits a `TryStream` into separate `Ok` and `Error` streams.
96///
97/// Note: This will deadlock if one branch outlives the other
98#[allow(clippy::type_complexity)]
99fn trystream_split_result<S>(
100    stream: S,
101) -> (
102    SplitCase<IntoStream<S>, S::Ok>,
103    SplitCase<IntoStream<S>, S::Error>,
104)
105where
106    S: TryStream + Unpin,
107    S::Ok: Debug,
108    S::Error: Debug,
109{
110    let stream = Arc::new(Mutex::new(stream.into_stream().peekable()));
111    (
112        SplitCase {
113            inner: stream.clone(),
114            should_consume_item: Result::is_ok,
115            try_extract_item_case: Result::ok,
116        },
117        SplitCase {
118            inner: stream,
119            should_consume_item: Result::is_err,
120            try_extract_item_case: Result::err,
121        },
122    )
123}
124
125/// Forwards Ok elements via a stream built from `make_via_stream`, while passing errors through unmodified
126pub(crate) fn trystream_try_via<S1, S2>(
127    input_stream: S1,
128    make_via_stream: impl FnOnce(SplitCase<IntoStream<S1>, S1::Ok>) -> S2,
129) -> impl Stream<Item = Result<S2::Ok, S1::Error>>
130where
131    S1: TryStream + Unpin,
132    S2: TryStream<Error = S1::Error>,
133    S1::Ok: Debug,
134    S1::Error: Debug,
135{
136    let (oks, errs) = trystream_split_result(input_stream); // the select -> SplitCase
137    let via = make_via_stream(oks); // the map_ok/err function
138    stream::select(via.into_stream(), errs.map(Err)) // recombine
139}
140
141/// A [`JoinHandle`] that cancels the [`Future`] when dropped, rather than detaching it
142pub struct CancelableJoinHandle<T> {
143    inner: JoinHandle<T>,
144}
145
146impl<T> CancelableJoinHandle<T>
147where
148    T: Send + 'static,
149{
150    /// Wrap a future in a cancelable handle, and spawn in a runtime
151    pub fn spawn(future: impl Future<Output = T> + Send + 'static, runtime: &Handle) -> Self {
152        CancelableJoinHandle {
153            inner: runtime.spawn(future),
154        }
155    }
156}
157
158impl<T> Drop for CancelableJoinHandle<T> {
159    fn drop(&mut self) {
160        self.inner.abort()
161    }
162}
163
164impl<T> Future for CancelableJoinHandle<T> {
165    type Output = T;
166
167    fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
168        self.inner.poll_unpin(cx).map(
169            // JoinError => underlying future was either aborted (which should only happen when the handle is dropped), or
170            // panicked (which should be propagated)
171            Result::unwrap,
172        )
173    }
174}
175
176#[pin_project]
177pub(crate) struct OnComplete<S, F> {
178    #[pin]
179    stream: stream::Fuse<S>,
180    #[pin]
181    on_complete: F,
182}
183
184impl<S: Stream, F: Future<Output = ()>> Stream for OnComplete<S, F> {
185    type Item = S::Item;
186
187    fn poll_next(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Option<Self::Item>> {
188        let this = self.project();
189        match this.stream.poll_next(cx) {
190            Poll::Ready(None) => match this.on_complete.poll(cx) {
191                Poll::Pending => Poll::Pending,
192                Poll::Ready(()) => Poll::Ready(None),
193            },
194            x => x,
195        }
196    }
197}
198
199pub(crate) trait KubeRuntimeStreamExt: Stream + Sized {
200    /// Runs the [`Future`] `on_complete` once the [`Stream`] finishes (by returning [`None`]).
201    fn on_complete<F: Future<Output = ()>>(self, on_complete: F) -> OnComplete<Self, F> {
202        OnComplete {
203            stream: self.fuse(),
204            on_complete,
205        }
206    }
207}
208
209impl<S: Stream> KubeRuntimeStreamExt for S {}
210
211#[cfg(test)]
212mod tests {
213    use std::convert::Infallible;
214
215    use futures::stream::{self, StreamExt};
216
217    use super::trystream_try_via;
218
219    // Type-level test does not need to be executed
220    #[allow(dead_code)]
221    fn trystream_try_via_should_be_able_to_borrow() {
222        struct WeirdComplexObject {}
223        impl Drop for WeirdComplexObject {
224            fn drop(&mut self) {}
225        }
226
227        let mut x = WeirdComplexObject {};
228        let y = WeirdComplexObject {};
229        drop(trystream_try_via(
230            Box::pin(stream::once(async {
231                let _ = &mut x;
232                Result::<_, Infallible>::Ok(())
233            })),
234            |s| {
235                s.map(|()| {
236                    let _ = &y;
237                    Ok(())
238                })
239            },
240        ));
241    }
242}
kube_runtime/utils/mod.rs

kube_runtime/utils/
mod.rs
