aws_smithy_async/future/
pagination_stream.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6//! Provides types to support stream-like operations for paginators.
7
8use crate::future::pagination_stream::collect::sealed::Collectable;
9use std::future::Future;
10use std::pin::Pin;
11use std::task::{Context, Poll};
12
13pub mod collect;
14pub mod fn_stream;
15use fn_stream::FnStream;
16
17/// Stream specifically made to support paginators.
18///
19/// `PaginationStream` provides two primary mechanisms for accessing stream of data.
20/// 1. With [`.next()`](PaginationStream::next) (or [`try_next()`](PaginationStream::try_next)):
21///
22/// ```no_run
23/// # async fn docs() {
24/// # use aws_smithy_async::future::pagination_stream::PaginationStream;
25/// # fn operation_to_yield_paginator<T>() -> PaginationStream<T> {
26/// #     todo!()
27/// # }
28/// # struct Page;
29/// let mut stream: PaginationStream<Page> = operation_to_yield_paginator();
30/// while let Some(page) = stream.next().await {
31///     // process `page`
32/// }
33/// # }
34/// ```
35/// 2. With [`.collect()`](PaginationStream::collect) (or [`try_collect()`](PaginationStream::try_collect)):
36///
37/// ```no_run
38/// # async fn docs() {
39/// # use aws_smithy_async::future::pagination_stream::PaginationStream;
40/// # fn operation_to_yield_paginator<T>() -> PaginationStream<T> {
41/// #     todo!()
42/// # }
43/// # struct Page;
44/// let mut stream: PaginationStream<Page> = operation_to_yield_paginator();
45/// let result = stream.collect::<Vec<Page>>().await;
46/// # }
47/// ```
48///
49/// ## Why isn't `PaginationStream` a `futures::stream::Stream`?
50///
51/// `PaginationStream` does not implement the `futures::stream::Stream` trait because
52/// the `futures` crate is 0.x, and we cannot expose 0.x dependencies in a 1.0 library.
53/// If you need `Stream` trait compatibility, use the [`PaginationStreamExt`] trait from
54/// the `aws-smithy-types-convert` crate, which provides conversion to `Stream`.
55///
56/// [`PaginationStreamExt`]: https://docs.rs/aws-smithy-types-convert/latest/aws_smithy_types_convert/stream/trait.PaginationStreamExt.html
57///
58/// [`PaginationStream`] is implemented in terms of [`FnStream`], but the latter is meant to be
59/// used internally and not by external users.
60#[derive(Debug)]
61pub struct PaginationStream<Item>(FnStream<Item>);
62
63impl<Item> PaginationStream<Item> {
64    /// Creates a `PaginationStream` from the given [`FnStream`].
65    pub fn new(stream: FnStream<Item>) -> Self {
66        Self(stream)
67    }
68
69    /// Consumes and returns the next `Item` from this stream.
70    pub async fn next(&mut self) -> Option<Item> {
71        self.0.next().await
72    }
73
74    /// Poll an item from the stream
75    pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Item>> {
76        Pin::new(&mut self.0).poll_next(cx)
77    }
78
79    /// Consumes this stream and gathers elements into a collection.
80    pub async fn collect<T: Collectable<Item>>(self) -> T {
81        self.0.collect().await
82    }
83}
84
85impl<T, E> PaginationStream<Result<T, E>> {
86    /// Yields the next item in the stream or returns an error if an error is encountered.
87    pub async fn try_next(&mut self) -> Result<Option<T>, E> {
88        self.next().await.transpose()
89    }
90
91    /// Convenience method for `.collect::<Result<Vec<_>, _>()`.
92    pub async fn try_collect(self) -> Result<Vec<T>, E> {
93        self.collect::<Result<Vec<T>, E>>().await
94    }
95}
96
97/// Utility wrapper to flatten paginated results
98///
99/// When flattening paginated results, it's most convenient to produce an iterator where the `Result`
100/// is present in each item. This provides `items()` which can wrap an stream of `Result<Page, Err>`
101/// and produce a stream of `Result<Item, Err>`.
102#[derive(Debug)]
103pub struct TryFlatMap<Page, Err>(PaginationStream<Result<Page, Err>>);
104
105impl<Page, Err> TryFlatMap<Page, Err> {
106    /// Creates a `TryFlatMap` that wraps the input.
107    pub fn new(stream: PaginationStream<Result<Page, Err>>) -> Self {
108        Self(stream)
109    }
110
111    /// Produces a new [`PaginationStream`] by mapping this stream with `map` then flattening the result.
112    pub fn flat_map<M, Item, Iter>(mut self, map: M) -> PaginationStream<Result<Item, Err>>
113    where
114        Page: Send + 'static,
115        Err: Send + 'static,
116        M: Fn(Page) -> Iter + Send + 'static,
117        Item: Send + 'static,
118        Iter: IntoIterator<Item = Item> + Send,
119        <Iter as IntoIterator>::IntoIter: Send,
120    {
121        PaginationStream::new(FnStream::new(|tx| {
122            Box::pin(async move {
123                while let Some(page) = self.0.next().await {
124                    match page {
125                        Ok(page) => {
126                            let mapped = map(page);
127                            for item in mapped.into_iter() {
128                                let _ = tx.send(Ok(item)).await;
129                            }
130                        }
131                        Err(e) => {
132                            let _ = tx.send(Err(e)).await;
133                            break;
134                        }
135                    }
136                }
137            }) as Pin<Box<dyn Future<Output = ()> + Send>>
138        }))
139    }
140}
141
142#[cfg(test)]
143mod test {
144    use crate::future::pagination_stream::{FnStream, PaginationStream, TryFlatMap};
145    use std::sync::{Arc, Mutex};
146    use std::time::Duration;
147
148    /// basic test of FnStream functionality
149    #[tokio::test]
150    async fn fn_stream_returns_results() {
151        tokio::time::pause();
152        let mut stream = FnStream::new(|tx| {
153            Box::pin(async move {
154                tx.send("1").await.expect("failed to send");
155                tokio::time::sleep(Duration::from_secs(1)).await;
156                tokio::time::sleep(Duration::from_secs(1)).await;
157                tx.send("2").await.expect("failed to send");
158                tokio::time::sleep(Duration::from_secs(1)).await;
159                tx.send("3").await.expect("failed to send");
160            })
161        });
162        let mut out = vec![];
163        while let Some(value) = stream.next().await {
164            out.push(value);
165        }
166        assert_eq!(vec!["1", "2", "3"], out);
167    }
168
169    #[tokio::test]
170    async fn fn_stream_try_next() {
171        tokio::time::pause();
172        let mut stream = FnStream::new(|tx| {
173            Box::pin(async move {
174                tx.send(Ok(1)).await.unwrap();
175                tx.send(Ok(2)).await.unwrap();
176                tx.send(Err("err")).await.unwrap();
177            })
178        });
179        let mut out = vec![];
180        while let Ok(value) = stream.try_next().await {
181            out.push(value);
182        }
183        assert_eq!(vec![Some(1), Some(2)], out);
184    }
185
186    // smithy-rs#1902: there was a bug where we could continue to poll the generator after it
187    // had returned Poll::Ready. This test case leaks the tx half so that the channel stays open
188    // but the send side generator completes. By calling `poll` multiple times on the resulting future,
189    // we can trigger the bug and validate the fix.
190    #[tokio::test]
191    async fn fn_stream_doesnt_poll_after_done() {
192        let mut stream = FnStream::new(|tx| {
193            Box::pin(async move {
194                assert!(tx.send("blah").await.is_ok());
195                Box::leak(Box::new(tx));
196            })
197        });
198        assert_eq!(Some("blah"), stream.next().await);
199        let mut test_stream = tokio_test::task::spawn(stream);
200        // `tokio_test::task::Spawn::poll_next` can only be invoked when the wrapped
201        // type implements the `Stream` trait. Here, `FnStream` does not implement it,
202        // so we work around it by using the `enter` method.
203        test_stream.enter(|ctx, pin| {
204            let polled = pin.poll_next(ctx);
205            assert!(polled.is_pending());
206        });
207        test_stream.enter(|ctx, pin| {
208            let polled = pin.poll_next(ctx);
209            assert!(polled.is_pending());
210        });
211    }
212
213    /// Tests that the generator will not advance until demand exists
214    #[tokio::test]
215    async fn waits_for_reader() {
216        let progress = Arc::new(Mutex::new(0));
217        let mut stream = FnStream::new(|tx| {
218            let progress = progress.clone();
219            Box::pin(async move {
220                *progress.lock().unwrap() = 1;
221                tx.send("1").await.expect("failed to send");
222                *progress.lock().unwrap() = 2;
223                tx.send("2").await.expect("failed to send");
224                *progress.lock().unwrap() = 3;
225                tx.send("3").await.expect("failed to send");
226                *progress.lock().unwrap() = 4;
227            })
228        });
229        assert_eq!(*progress.lock().unwrap(), 0);
230        stream.next().await.expect("ready");
231        assert_eq!(*progress.lock().unwrap(), 1);
232
233        assert_eq!("2", stream.next().await.expect("ready"));
234        assert_eq!(2, *progress.lock().unwrap());
235
236        let _ = stream.next().await.expect("ready");
237        assert_eq!(3, *progress.lock().unwrap());
238        assert_eq!(None, stream.next().await);
239        assert_eq!(4, *progress.lock().unwrap());
240    }
241
242    #[tokio::test]
243    async fn generator_with_errors() {
244        let mut stream = FnStream::new(|tx| {
245            Box::pin(async move {
246                for i in 0..5 {
247                    if i != 2 {
248                        if tx.send(Ok(i)).await.is_err() {
249                            return;
250                        }
251                    } else {
252                        tx.send(Err(i)).await.unwrap();
253                        return;
254                    }
255                }
256            })
257        });
258        let mut out = vec![];
259        while let Some(Ok(value)) = stream.next().await {
260            out.push(value);
261        }
262        assert_eq!(vec![0, 1], out);
263    }
264
265    #[tokio::test]
266    async fn flatten_items_ok() {
267        #[derive(Debug)]
268        struct Output {
269            items: Vec<u8>,
270        }
271        let stream: FnStream<Result<_, &str>> = FnStream::new(|tx| {
272            Box::pin(async move {
273                tx.send(Ok(Output {
274                    items: vec![1, 2, 3],
275                }))
276                .await
277                .unwrap();
278                tx.send(Ok(Output {
279                    items: vec![4, 5, 6],
280                }))
281                .await
282                .unwrap();
283            })
284        });
285        assert_eq!(
286            Ok(vec![1, 2, 3, 4, 5, 6]),
287            TryFlatMap::new(PaginationStream::new(stream))
288                .flat_map(|output| output.items.into_iter())
289                .try_collect()
290                .await,
291        );
292    }
293
294    #[tokio::test]
295    async fn flatten_items_error() {
296        #[derive(Debug)]
297        struct Output {
298            items: Vec<u8>,
299        }
300        let stream = FnStream::new(|tx| {
301            Box::pin(async move {
302                tx.send(Ok(Output {
303                    items: vec![1, 2, 3],
304                }))
305                .await
306                .unwrap();
307                tx.send(Err("bummer")).await.unwrap();
308            })
309        });
310        assert_eq!(
311            Err("bummer"),
312            TryFlatMap::new(PaginationStream::new(stream))
313                .flat_map(|output| output.items.into_iter())
314                .try_collect()
315                .await
316        )
317    }
318}
aws_smithy_async/future/pagination_stream.rs

aws_smithy_async/future/
pagination_stream.rs
