async_utils/stream/
short_circuit.rs

1// Copyright 2024 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use futures::stream::FusedStream;
6use futures::task::Poll;
7use futures::{ready, Stream, TryStream};
8use pin_project::pin_project;
9use std::pin::Pin;
10use std::task::Context;
11
12/// Short circuits a [`TryStream`] stream when the first error occurs.
13///
14/// The [`futures::TryStreamExt`] extension crate provides combinators that
15/// pass through errors. However, sometimes we want the stream to short circurt
16/// when an error occurs.
17#[pin_project]
18pub struct ShortCircuit<St> {
19    #[pin]
20    stream: Option<St>,
21}
22
23impl<St: TryStream> Stream for ShortCircuit<St> {
24    type Item = Result<St::Ok, St::Error>;
25
26    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
27        let mut this = self.project();
28        let Some(stream) = this.stream.as_mut().as_pin_mut() else { return Poll::Ready(None) };
29        Poll::Ready(match ready!(stream.try_poll_next(cx)) {
30            Some(Ok(val)) => Some(Ok(val)),
31            Some(Err(err)) => {
32                this.stream.set(None);
33                Some(Err(err))
34            }
35            None => None,
36        })
37    }
38}
39
40impl<St: FusedStream + TryStream> FusedStream for ShortCircuit<St> {
41    fn is_terminated(&self) -> bool {
42        self.stream.as_ref().map(|st| st.is_terminated()).unwrap_or(true)
43    }
44}
45
46impl<St> ShortCircuit<St> {
47    /// Creates a new [`TryStream`] that will terminate after the first error
48    /// is observed.
49    pub fn new(stream: St) -> Self {
50        Self { stream: Some(stream) }
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use futures::{FutureExt, StreamExt as _};
57    use test_case::test_case;
58
59    use super::*;
60
61    #[test_case(vec![Ok(1), Err(()), Ok(2)] => vec![Ok(1), Err(())])]
62    #[test_case(vec![Ok(1), Ok(2), Ok(3)] => vec![Ok(1), Ok(2), Ok(3)])]
63    fn short_circuit(input: Vec<Result<usize, ()>>) -> Vec<Result<usize, ()>> {
64        let stream = futures::stream::iter(input).fuse();
65        let mut short_circuited = ShortCircuit::new(stream);
66        assert!(!short_circuited.is_terminated());
67        let output = (&mut short_circuited)
68            .collect::<Vec<_>>()
69            .now_or_never()
70            .expect("all items should be ready");
71        assert!(short_circuited.is_terminated());
72        output
73    }
74}
async_utils/stream/short_circuit.rs

async_utils/stream/
short_circuit.rs
