parallel_stream/par_stream/
for_each.rs

1use async_std::channel::{self, Receiver, Sender};
2use async_std::prelude::*;
3use async_std::task::{self, Context, Poll};
4
5use std::pin::Pin;
6use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
7use std::sync::Arc;
8
9use crate::ParallelStream;
10
11pin_project_lite::pin_project! {
12    /// Call a closure on each element of the stream.
13    #[derive(Debug)]
14    pub struct ForEach {
15        // Receiver that tracks whether all tasks have finished executing.
16        #[pin]
17        receiver: Receiver<()>,
18        // Track whether the input stream has been exhausted.
19        exhausted: Arc<AtomicBool>,
20        // Count how many tasks are executing.
21        ref_count: Arc<AtomicU64>,
22    }
23}
24
25impl ForEach {
26    /// Create a new instance of `ForEach`.
27    pub fn new<S, F, Fut>(mut stream: S, mut f: F) -> Self
28    where
29        S: ParallelStream,
30        F: FnMut(S::Item) -> Fut + Send + Sync + Copy + 'static,
31        Fut: Future<Output = ()> + Send,
32    {
33        let exhausted = Arc::new(AtomicBool::new(false));
34        let ref_count = Arc::new(AtomicU64::new(0));
35        let (sender, receiver): (Sender<()>, Receiver<()>) = channel::bounded(1);
36        let _limit = stream.get_limit();
37
38        // Initialize the return type here to prevent borrowing issues.
39        let this = Self {
40            receiver,
41            exhausted: exhausted.clone(),
42            ref_count: ref_count.clone(),
43        };
44
45        task::spawn(async move {
46            while let Some(item) = stream.next().await {
47                let sender = sender.clone();
48                let exhausted = exhausted.clone();
49                let ref_count = ref_count.clone();
50
51                ref_count.fetch_add(1, Ordering::SeqCst);
52
53                task::spawn(async move {
54                    // Execute the closure.
55                    f(item).await;
56
57                    // Wake up the receiver if we know we're done.
58                    ref_count.fetch_sub(1, Ordering::SeqCst);
59                    if exhausted.load(Ordering::SeqCst) && ref_count.load(Ordering::SeqCst) == 0 {
60                        sender.send(()).await.expect("message failed to send");
61                    }
62                });
63            }
64
65            // The input stream will no longer yield items.
66            exhausted.store(true, Ordering::SeqCst);
67        });
68
69        this
70    }
71}
72
73impl Future for ForEach {
74    type Output = ();
75
76    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
77        let this = self.project();
78        task::ready!(this.receiver.poll_next(cx));
79        Poll::Ready(())
80    }
81}
82
83#[async_std::test]
84async fn smoke() {
85    let s = async_std::stream::repeat(5usize);
86    crate::from_stream(s)
87        .take(3)
88        .for_each(|n| async move {
89            // TODO: assert that this is called 3 times.
90            dbg!(n);
91        })
92        .await;
93}
parallel_stream/par_stream/for_each.rs

parallel_stream/par_stream/
for_each.rs
