asupersync/stream/
for_each.rs

1//! ForEach combinator for streams.
2//!
3//! The `ForEach` future consumes a stream and executes a closure for each item.
4
5use super::Stream;
6use pin_project::pin_project;
7use std::future::Future;
8use std::pin::Pin;
9use std::task::{Context, Poll};
10
11/// Cooperative budget for items processed in a single poll.
12///
13/// Without this bound, always-ready streams can monopolize one executor turn.
14const FOR_EACH_COOPERATIVE_BUDGET: usize = 1024;
15
16/// A future that executes a closure for each item in a stream.
17///
18/// Created by [`StreamExt::for_each`](super::StreamExt::for_each).
19#[pin_project]
20#[derive(Debug)]
21#[must_use = "futures do nothing unless polled"]
22pub struct ForEach<S, F> {
23    #[pin]
24    stream: S,
25    f: F,
26    completed: bool,
27}
28
29impl<S, F> ForEach<S, F> {
30    /// Creates a new `ForEach` future.
31    #[inline]
32    pub(crate) fn new(stream: S, f: F) -> Self {
33        Self {
34            stream,
35            f,
36            completed: false,
37        }
38    }
39}
40
41impl<S, F> Future for ForEach<S, F>
42where
43    S: Stream,
44    F: FnMut(S::Item),
45{
46    type Output = ();
47
48    #[inline]
49    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
50        let mut this = self.project();
51        if *this.completed {
52            return Poll::Ready(());
53        }
54        let mut processed_this_poll = 0usize;
55        loop {
56            match this.stream.as_mut().poll_next(cx) {
57                Poll::Ready(Some(item)) => {
58                    (this.f)(item);
59                    processed_this_poll += 1;
60                    if processed_this_poll >= FOR_EACH_COOPERATIVE_BUDGET {
61                        cx.waker().wake_by_ref();
62                        return Poll::Pending;
63                    }
64                }
65                Poll::Ready(None) => {
66                    *this.completed = true;
67                    return Poll::Ready(());
68                }
69                Poll::Pending => return Poll::Pending,
70            }
71        }
72    }
73}
74
75/// A future that executes an async closure for each item in a stream.
76///
77/// Created by [`StreamExt::for_each_async`](super::StreamExt::for_each_async).
78#[pin_project]
79#[derive(Debug)]
80#[must_use = "futures do nothing unless polled"]
81pub struct ForEachAsync<S, F, Fut> {
82    #[pin]
83    stream: S,
84    f: F,
85    #[pin]
86    pending: Option<Fut>,
87    completed: bool,
88}
89
90impl<S, F, Fut> ForEachAsync<S, F, Fut> {
91    #[inline]
92    pub(crate) fn new(stream: S, f: F) -> Self {
93        Self {
94            stream,
95            f,
96            pending: None,
97            completed: false,
98        }
99    }
100}
101
102impl<S, F, Fut> Future for ForEachAsync<S, F, Fut>
103where
104    S: Stream,
105    F: FnMut(S::Item) -> Fut,
106    Fut: Future<Output = ()>,
107{
108    type Output = ();
109
110    #[inline]
111    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
112        let mut this = self.project();
113        if *this.completed {
114            return Poll::Ready(());
115        }
116        let mut processed_this_poll = 0usize;
117        loop {
118            // Complete pending future first
119            if let Some(fut) = this.pending.as_mut().as_pin_mut() {
120                match fut.poll(cx) {
121                    Poll::Ready(()) => {
122                        this.pending.set(None);
123                        processed_this_poll += 1;
124                        if processed_this_poll >= FOR_EACH_COOPERATIVE_BUDGET {
125                            cx.waker().wake_by_ref();
126                            return Poll::Pending;
127                        }
128                    }
129                    Poll::Pending => return Poll::Pending,
130                }
131            }
132
133            // Get next item
134            match this.stream.as_mut().poll_next(cx) {
135                Poll::Ready(Some(item)) => {
136                    this.pending.set(Some((this.f)(item)));
137                }
138                Poll::Ready(None) => {
139                    *this.completed = true;
140                    return Poll::Ready(());
141                }
142                Poll::Pending => return Poll::Pending,
143            }
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    #![allow(
151        clippy::pedantic,
152        clippy::nursery,
153        clippy::expect_fun_call,
154        clippy::map_unwrap_or,
155        clippy::cast_possible_wrap,
156        clippy::future_not_send
157    )]
158    use super::*;
159    use crate::stream::iter;
160    use std::cell::RefCell;
161    use std::sync::Arc;
162    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
163    use std::task::{Poll, Waker};
164
165    fn noop_waker() -> Waker {
166        std::task::Waker::noop().clone()
167    }
168
169    struct TrackWaker(Arc<AtomicBool>);
170
171    use std::task::Wake;
172    impl Wake for TrackWaker {
173        fn wake(self: Arc<Self>) {
174            self.0.store(true, Ordering::SeqCst);
175        }
176
177        fn wake_by_ref(self: &Arc<Self>) {
178            self.0.store(true, Ordering::SeqCst);
179        }
180    }
181
182    #[derive(Debug, Default)]
183    struct AlwaysReadyCounter {
184        next: usize,
185        end: usize,
186    }
187
188    impl AlwaysReadyCounter {
189        fn new(end: usize) -> Self {
190            Self { next: 0, end }
191        }
192    }
193
194    impl Stream for AlwaysReadyCounter {
195        type Item = usize;
196
197        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
198            if self.next >= self.end {
199                return Poll::Ready(None);
200            }
201
202            let item = self.next;
203            self.next += 1;
204            Poll::Ready(Some(item))
205        }
206    }
207
208    #[derive(Debug)]
209    struct PollCountingEmptyStream {
210        polls: Arc<AtomicUsize>,
211    }
212
213    impl PollCountingEmptyStream {
214        fn new(polls: Arc<AtomicUsize>) -> Self {
215            Self { polls }
216        }
217    }
218
219    impl Stream for PollCountingEmptyStream {
220        type Item = usize;
221
222        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
223            self.polls.fetch_add(1, Ordering::SeqCst);
224            Poll::Ready(None)
225        }
226    }
227
228    #[derive(Debug)]
229    struct PendingOnceThenItems<T> {
230        items: Vec<T>,
231        index: usize,
232        pending_first: bool,
233    }
234
235    impl<T> PendingOnceThenItems<T> {
236        fn new(items: Vec<T>) -> Self {
237            Self {
238                items,
239                index: 0,
240                pending_first: true,
241            }
242        }
243    }
244
245    impl<T: Clone + Unpin> Stream for PendingOnceThenItems<T> {
246        type Item = T;
247
248        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
249            if self.pending_first {
250                self.pending_first = false;
251                return Poll::Pending;
252            }
253            if self.index >= self.items.len() {
254                return Poll::Ready(None);
255            }
256
257            let item = self.items[self.index].clone();
258            self.index += 1;
259            Poll::Ready(Some(item))
260        }
261    }
262
263    fn init_test(name: &str) {
264        crate::test_utils::init_test_logging();
265        crate::test_phase!(name);
266    }
267
268    fn poll_unit_future_to_completion<F>(future: &mut F)
269    where
270        F: std::future::Future<Output = ()> + Unpin,
271    {
272        let waker = noop_waker();
273        let mut cx = Context::from_waker(&waker);
274        loop {
275            match Pin::new(&mut *future).poll(&mut cx) {
276                Poll::Ready(()) => return,
277                Poll::Pending => {}
278            }
279        }
280    }
281
282    fn for_each_side_effects(input: Vec<i32>) -> Vec<i32> {
283        let results = RefCell::new(Vec::new());
284        let mut future = ForEach::new(iter(input), |item| {
285            results.borrow_mut().push(item);
286        });
287        poll_unit_future_to_completion(&mut future);
288        drop(future);
289        results.into_inner()
290    }
291
292    fn for_each_async_side_effects(input: Vec<i32>) -> Vec<i32> {
293        let results = RefCell::new(Vec::new());
294        let mut future = ForEachAsync::new(iter(input), |item| {
295            let results = &results;
296            Box::pin(async move {
297                results.borrow_mut().push(item);
298            })
299        });
300        poll_unit_future_to_completion(&mut future);
301        drop(future);
302        results.into_inner()
303    }
304
305    #[test]
306    fn for_each_collects_side_effects() {
307        init_test("for_each_collects_side_effects");
308        let results = RefCell::new(Vec::new());
309        let mut future = ForEach::new(iter(vec![1i32, 2, 3]), |x| {
310            results.borrow_mut().push(x);
311        });
312        let waker = noop_waker();
313        let mut cx = Context::from_waker(&waker);
314
315        match Pin::new(&mut future).poll(&mut cx) {
316            Poll::Ready(()) => {
317                let collected = results.borrow().clone();
318                let ok = collected == vec![1, 2, 3];
319                crate::assert_with_log!(ok, "collected", vec![1, 2, 3], collected);
320            }
321            Poll::Pending => panic!("expected Ready"),
322        }
323        crate::test_complete!("for_each_collects_side_effects");
324    }
325
326    #[test]
327    fn for_each_empty() {
328        init_test("for_each_empty");
329        let mut called = false;
330        let mut future = ForEach::new(iter(Vec::<i32>::new()), |_| {
331            called = true;
332        });
333        let waker = noop_waker();
334        let mut cx = Context::from_waker(&waker);
335
336        match Pin::new(&mut future).poll(&mut cx) {
337            Poll::Ready(()) => {
338                crate::assert_with_log!(!called, "not called", false, called);
339            }
340            Poll::Pending => panic!("expected Ready"),
341        }
342        crate::test_complete!("for_each_empty");
343    }
344
345    #[test]
346    fn for_each_async() {
347        init_test("for_each_async");
348        let results = RefCell::new(Vec::new());
349        let mut future = ForEachAsync::new(iter(vec![1i32, 2, 3]), |x| {
350            let res = &results;
351            Box::pin(async move {
352                res.borrow_mut().push(x);
353            })
354        });
355        let waker = noop_waker();
356        let mut cx = Context::from_waker(&waker);
357
358        // This test requires re-polling because async block yields?
359        // No, Box::pin(async { ... }) is ready immediately if no await.
360        // But ForEachAsync needs to poll the future.
361
362        // We simulate polling loop
363        loop {
364            match Pin::new(&mut future).poll(&mut cx) {
365                Poll::Ready(()) => break,
366                Poll::Pending => {} // Should not happen for immediate futures but safe
367            }
368        }
369
370        let collected = results.borrow().clone();
371        let ok = collected == vec![1, 2, 3];
372        crate::assert_with_log!(ok, "collected", vec![1, 2, 3], collected);
373        crate::test_complete!("for_each_async");
374    }
375
376    /// Invariant: ForEachAsync with empty stream completes without calling the closure.
377    #[test]
378    fn for_each_async_empty() {
379        init_test("for_each_async_empty");
380        let mut called = false;
381        let mut future = ForEachAsync::new(iter(Vec::<i32>::new()), |_x| {
382            called = true;
383            Box::pin(async {})
384        });
385        let waker = noop_waker();
386        let mut cx = Context::from_waker(&waker);
387
388        let poll = Pin::new(&mut future).poll(&mut cx);
389        let completed = matches!(poll, Poll::Ready(()));
390        crate::assert_with_log!(completed, "async empty completes", true, completed);
391        crate::assert_with_log!(!called, "closure not called", false, called);
392
393        crate::test_complete!("for_each_async_empty");
394    }
395
396    #[test]
397    fn for_each_yields_after_budget_on_always_ready_stream() {
398        init_test("for_each_yields_after_budget_on_always_ready_stream");
399        let seen = RefCell::new(Vec::new());
400        let mut future = ForEach::new(
401            AlwaysReadyCounter::new(FOR_EACH_COOPERATIVE_BUDGET + 5),
402            |x| seen.borrow_mut().push(x),
403        );
404        let woke = Arc::new(AtomicBool::new(false));
405        let waker = Waker::from(Arc::new(TrackWaker(woke.clone())));
406        let mut cx = Context::from_waker(&waker);
407
408        let first = Pin::new(&mut future).poll(&mut cx);
409        crate::assert_with_log!(
410            matches!(first, Poll::Pending),
411            "first poll yields cooperatively",
412            "Poll::Pending",
413            first
414        );
415        crate::assert_with_log!(
416            future.stream.next == FOR_EACH_COOPERATIVE_BUDGET,
417            "upstream advanced only to budget",
418            FOR_EACH_COOPERATIVE_BUDGET,
419            future.stream.next
420        );
421        crate::assert_with_log!(
422            seen.borrow().len() == FOR_EACH_COOPERATIVE_BUDGET,
423            "side effects applied to budget items",
424            FOR_EACH_COOPERATIVE_BUDGET,
425            seen.borrow().len()
426        );
427        crate::assert_with_log!(
428            woke.load(Ordering::SeqCst),
429            "self-wake requested",
430            true,
431            woke.load(Ordering::SeqCst)
432        );
433
434        let second = Pin::new(&mut future).poll(&mut cx);
435        crate::assert_with_log!(
436            matches!(second, Poll::Ready(())),
437            "second poll completes",
438            "Poll::Ready(())",
439            second
440        );
441        crate::assert_with_log!(
442            seen.borrow().len() == FOR_EACH_COOPERATIVE_BUDGET + 5,
443            "all side effects complete",
444            FOR_EACH_COOPERATIVE_BUDGET + 5,
445            seen.borrow().len()
446        );
447        crate::test_complete!("for_each_yields_after_budget_on_always_ready_stream");
448    }
449
450    #[test]
451    fn for_each_async_yields_after_budget_on_immediate_futures() {
452        init_test("for_each_async_yields_after_budget_on_immediate_futures");
453        let seen = RefCell::new(Vec::new());
454        let mut future = ForEachAsync::new(
455            AlwaysReadyCounter::new(FOR_EACH_COOPERATIVE_BUDGET + 5),
456            |x| {
457                let seen = &seen;
458                Box::pin(async move {
459                    seen.borrow_mut().push(x);
460                })
461            },
462        );
463        let woke = Arc::new(AtomicBool::new(false));
464        let waker = Waker::from(Arc::new(TrackWaker(woke.clone())));
465        let mut cx = Context::from_waker(&waker);
466
467        let first = Pin::new(&mut future).poll(&mut cx);
468        crate::assert_with_log!(
469            matches!(first, Poll::Pending),
470            "first poll yields cooperatively",
471            "Poll::Pending",
472            first
473        );
474        crate::assert_with_log!(
475            future.stream.next == FOR_EACH_COOPERATIVE_BUDGET,
476            "upstream advanced only to budget",
477            FOR_EACH_COOPERATIVE_BUDGET,
478            future.stream.next
479        );
480        crate::assert_with_log!(
481            future.pending.is_none(),
482            "no pending future left at cooperative boundary",
483            true,
484            future.pending.is_none()
485        );
486        crate::assert_with_log!(
487            seen.borrow().len() == FOR_EACH_COOPERATIVE_BUDGET,
488            "side effects applied to budget items",
489            FOR_EACH_COOPERATIVE_BUDGET,
490            seen.borrow().len()
491        );
492        crate::assert_with_log!(
493            woke.load(Ordering::SeqCst),
494            "self-wake requested",
495            true,
496            woke.load(Ordering::SeqCst)
497        );
498
499        let second = Pin::new(&mut future).poll(&mut cx);
500        crate::assert_with_log!(
501            matches!(second, Poll::Ready(())),
502            "second poll completes",
503            "Poll::Ready(())",
504            second
505        );
506        crate::assert_with_log!(
507            seen.borrow().len() == FOR_EACH_COOPERATIVE_BUDGET + 5,
508            "all side effects complete",
509            FOR_EACH_COOPERATIVE_BUDGET + 5,
510            seen.borrow().len()
511        );
512        crate::test_complete!("for_each_async_yields_after_budget_on_immediate_futures");
513    }
514
515    #[test]
516    fn for_each_repoll_after_completion_fails_closed_without_repolling_upstream() {
517        init_test("for_each_repoll_after_completion_fails_closed_without_repolling_upstream");
518        let polls = Arc::new(AtomicUsize::new(0));
519        let mut future = ForEach::new(PollCountingEmptyStream::new(Arc::clone(&polls)), |_| {});
520        let waker = noop_waker();
521        let mut cx = Context::from_waker(&waker);
522
523        let first = Pin::new(&mut future).poll(&mut cx);
524        crate::assert_with_log!(
525            matches!(first, Poll::Ready(())),
526            "first poll completes",
527            "Poll::Ready(())",
528            first
529        );
530        crate::assert_with_log!(
531            polls.load(Ordering::SeqCst) == 1,
532            "first completion polls upstream once",
533            1,
534            polls.load(Ordering::SeqCst)
535        );
536
537        // Fail-closed: repoll returns Ready(()) instead of panicking
538        let repoll = Pin::new(&mut future).poll(&mut cx);
539        crate::assert_with_log!(
540            matches!(repoll, Poll::Ready(())),
541            "repoll returns Ready(())",
542            "Poll::Ready(())",
543            repoll
544        );
545        crate::assert_with_log!(
546            polls.load(Ordering::SeqCst) == 1,
547            "repoll does not touch upstream again",
548            1,
549            polls.load(Ordering::SeqCst)
550        );
551        crate::test_complete!(
552            "for_each_repoll_after_completion_fails_closed_without_repolling_upstream"
553        );
554    }
555
556    #[test]
557    fn for_each_async_repoll_after_completion_fails_closed_without_repolling_upstream() {
558        init_test("for_each_async_repoll_after_completion_fails_closed_without_repolling_upstream");
559        let polls = Arc::new(AtomicUsize::new(0));
560        let mut future =
561            ForEachAsync::new(PollCountingEmptyStream::new(Arc::clone(&polls)), |_| {
562                Box::pin(async {})
563            });
564        let waker = noop_waker();
565        let mut cx = Context::from_waker(&waker);
566
567        let first = Pin::new(&mut future).poll(&mut cx);
568        crate::assert_with_log!(
569            matches!(first, Poll::Ready(())),
570            "first poll completes",
571            "Poll::Ready(())",
572            first
573        );
574        crate::assert_with_log!(
575            polls.load(Ordering::SeqCst) == 1,
576            "first completion polls upstream once",
577            1,
578            polls.load(Ordering::SeqCst)
579        );
580
581        // Fail-closed: repoll returns Ready(()) instead of panicking
582        let repoll = Pin::new(&mut future).poll(&mut cx);
583        crate::assert_with_log!(
584            matches!(repoll, Poll::Ready(())),
585            "repoll returns Ready(())",
586            "Poll::Ready(())",
587            repoll
588        );
589        crate::assert_with_log!(
590            polls.load(Ordering::SeqCst) == 1,
591            "repoll does not touch upstream again",
592            1,
593            polls.load(Ordering::SeqCst)
594        );
595        crate::test_complete!(
596            "for_each_async_repoll_after_completion_fails_closed_without_repolling_upstream"
597        );
598    }
599
600    #[test]
601    fn mr_for_each_partitioned_inputs_match_unsplit_side_effects() {
602        init_test("mr_for_each_partitioned_inputs_match_unsplit_side_effects");
603        let input = vec![-3, -1, 0, 1, 2, 5, 8];
604        let expected = for_each_side_effects(input.clone());
605
606        for split in 0..=input.len() {
607            let mut partitioned = for_each_side_effects(input[..split].to_vec());
608            partitioned.extend(for_each_side_effects(input[split..].to_vec()));
609            crate::assert_with_log!(
610                partitioned == expected,
611                format!("split at {split}"),
612                expected.clone(),
613                partitioned
614            );
615        }
616        crate::test_complete!("mr_for_each_partitioned_inputs_match_unsplit_side_effects");
617    }
618
619    #[test]
620    fn mr_for_each_async_immediate_matches_sync_side_effects() {
621        init_test("mr_for_each_async_immediate_matches_sync_side_effects");
622        for input in [
623            Vec::new(),
624            vec![1],
625            vec![1, 1, 2, 3, 5, 8],
626            vec![-10, 0, 10, 20],
627        ] {
628            let sync = for_each_side_effects(input.clone());
629            let async_immediate = for_each_async_side_effects(input.clone());
630            crate::assert_with_log!(
631                async_immediate == sync,
632                format!("input {input:?}"),
633                sync,
634                async_immediate
635            );
636        }
637        crate::test_complete!("mr_for_each_async_immediate_matches_sync_side_effects");
638    }
639
640    #[test]
641    fn mr_for_each_pending_cancellation_before_first_item_has_no_side_effects() {
642        init_test("mr_for_each_pending_cancellation_before_first_item_has_no_side_effects");
643        let input = vec![4, 8, 15, 16, 23, 42];
644        let results = RefCell::new(Vec::new());
645        let mut future = ForEach::new(PendingOnceThenItems::new(input.clone()), |item| {
646            results.borrow_mut().push(item);
647        });
648        let waker = noop_waker();
649        let mut cx = Context::from_waker(&waker);
650
651        let first_poll = Pin::new(&mut future).poll(&mut cx);
652        crate::assert_with_log!(
653            matches!(first_poll, Poll::Pending),
654            "first poll is pending",
655            "Poll::Pending",
656            first_poll
657        );
658        crate::assert_with_log!(
659            results.borrow().is_empty(),
660            "no side effects before first item is ready",
661            true,
662            results.borrow().is_empty()
663        );
664
665        drop(future);
666        let cancelled_effects = results.into_inner();
667        crate::assert_with_log!(
668            cancelled_effects.is_empty(),
669            "cancellation before first ready item has no effects",
670            Vec::<i32>::new(),
671            cancelled_effects
672        );
673
674        let fresh = for_each_side_effects(input.clone());
675        crate::assert_with_log!(
676            fresh == input,
677            "fresh equivalent run still observes all items",
678            input,
679            fresh
680        );
681        crate::test_complete!(
682            "mr_for_each_pending_cancellation_before_first_item_has_no_side_effects"
683        );
684    }
685}
asupersync/stream/for_each.rs

asupersync/stream/
for_each.rs
