rayon/iter/
par_bridge.rs

1#[cfg(not(feature = "web_spin_lock"))]
2use std::sync::Mutex;
3
4#[cfg(feature = "web_spin_lock")]
5use wasm_sync::Mutex;
6
7use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
8
9use crate::iter::plumbing::{bridge_unindexed, Folder, UnindexedConsumer, UnindexedProducer};
10use crate::iter::ParallelIterator;
11use crate::{current_num_threads, current_thread_index};
12
13/// Conversion trait to convert an `Iterator` to a `ParallelIterator`.
14///
15/// This creates a "bridge" from a sequential iterator to a parallel one, by distributing its items
16/// across the Rayon thread pool. This has the advantage of being able to parallelize just about
17/// anything, but the resulting `ParallelIterator` can be less efficient than if you started with
18/// `par_iter` instead. However, it can still be useful for iterators that are difficult to
19/// parallelize by other means, like channels or file or network I/O.
20///
21/// Iterator items are pulled by `next()` one at a time, synchronized from each thread that is
22/// ready for work, so this may become a bottleneck if the serial iterator can't keep up with the
23/// parallel demand. The items are not buffered by `IterBridge`, so it's fine to use this with
24/// large or even unbounded iterators.
25///
26/// The resulting iterator is not guaranteed to keep the order of the original iterator.
27///
28/// # Examples
29///
30/// To use this trait, take an existing `Iterator` and call `par_bridge` on it. After that, you can
31/// use any of the `ParallelIterator` methods:
32///
33/// ```
34/// use rayon::iter::ParallelBridge;
35/// use rayon::prelude::ParallelIterator;
36/// use std::sync::mpsc::channel;
37///
38/// let rx = {
39///     let (tx, rx) = channel();
40///
41///     tx.send("one!");
42///     tx.send("two!");
43///     tx.send("three!");
44///
45///     rx
46/// };
47///
48/// let mut output: Vec<&'static str> = rx.into_iter().par_bridge().collect();
49/// output.sort_unstable();
50///
51/// assert_eq!(&*output, &["one!", "three!", "two!"]);
52/// ```
53pub trait ParallelBridge: Sized {
54    /// Creates a bridge from this type to a `ParallelIterator`.
55    fn par_bridge(self) -> IterBridge<Self>;
56}
57
58impl<T> ParallelBridge for T
59where
60    T: Iterator<Item: Send> + Send,
61{
62    fn par_bridge(self) -> IterBridge<Self> {
63        IterBridge { iter: self }
64    }
65}
66
67/// `IterBridge` is a parallel iterator that wraps a sequential iterator.
68///
69/// This type is created when using the `par_bridge` method on `ParallelBridge`. See the
70/// [`ParallelBridge`] documentation for details.
71#[derive(Debug, Clone)]
72pub struct IterBridge<Iter> {
73    iter: Iter,
74}
75
76impl<Iter> ParallelIterator for IterBridge<Iter>
77where
78    Iter: Iterator<Item: Send> + Send,
79{
80    type Item = Iter::Item;
81
82    fn drive_unindexed<C>(self, consumer: C) -> C::Result
83    where
84        C: UnindexedConsumer<Self::Item>,
85    {
86        let num_threads = current_num_threads();
87        let threads_started: Vec<_> = (0..num_threads).map(|_| AtomicBool::new(false)).collect();
88
89        bridge_unindexed(
90            &IterParallelProducer {
91                split_count: AtomicUsize::new(num_threads),
92                iter: Mutex::new(self.iter.fuse()),
93                threads_started: &threads_started,
94            },
95            consumer,
96        )
97    }
98}
99
100struct IterParallelProducer<'a, Iter> {
101    split_count: AtomicUsize,
102    iter: Mutex<std::iter::Fuse<Iter>>,
103    threads_started: &'a [AtomicBool],
104}
105
106impl<Iter: Iterator + Send> UnindexedProducer for &IterParallelProducer<'_, Iter> {
107    type Item = Iter::Item;
108
109    fn split(self) -> (Self, Option<Self>) {
110        // Check if the iterator is exhausted
111        let update = self
112            .split_count
113            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |c| c.checked_sub(1));
114        (self, update.is_ok().then_some(self))
115    }
116
117    fn fold_with<F>(self, mut folder: F) -> F
118    where
119        F: Folder<Self::Item>,
120    {
121        // Guard against work-stealing-induced recursion, in case `Iter::next()`
122        // calls rayon internally, so we don't deadlock our mutex. We might also
123        // be recursing via `folder` methods, which doesn't present a mutex hazard,
124        // but it's lower overhead for us to just check this once, rather than
125        // updating additional shared state on every mutex lock/unlock.
126        // (If this isn't a rayon thread, then there's no work-stealing anyway...)
127        if let Some(i) = current_thread_index() {
128            // Note: If the number of threads in the pool ever grows dynamically, then
129            // we'll end up sharing flags and may falsely detect recursion -- that's
130            // still fine for overall correctness, just not optimal for parallelism.
131            let thread_started = &self.threads_started[i % self.threads_started.len()];
132            if thread_started.swap(true, Ordering::Relaxed) {
133                // We can't make progress with a nested mutex, so just return and let
134                // the outermost loop continue with the rest of the iterator items.
135                return folder;
136            }
137        }
138
139        loop {
140            if let Ok(mut iter) = self.iter.lock() {
141                if let Some(it) = iter.next() {
142                    drop(iter);
143                    folder = folder.consume(it);
144                    if folder.full() {
145                        return folder;
146                    }
147                } else {
148                    return folder;
149                }
150            } else {
151                // any panics from other threads will have been caught by the pool,
152                // and will be re-thrown when joined - just exit
153                return folder;
154            }
155        }
156    }
157}
rayon/iter/par_bridge.rs

rayon/iter/
par_bridge.rs
