hataori/
local.rs

1use crate::domain::{Domain, DomainAdmission, DomainBusy, LocalMode};
2use crate::lanes::{map_in_lanes, LaneControl, Lanes, LanesError, MapInError};
3use crate::map::{map, MapError};
4use std::any::Any;
5use std::convert::Infallible;
6use std::fmt;
7use std::panic::AssertUnwindSafe;
8use std::sync::{Mutex, PoisonError};
9
10pub(crate) fn in_rayon_worker_context() -> bool {
11    rayon::current_thread_index().is_some()
12}
13
14/// Runs `items` one at a time inside `pool` without stopping early.
15///
16/// [`LocalMode::Sequential`] and [`LocalMode::Inner`] share this path: both run
17/// one callback at a time on a pool worker, and `Inner` keeps nested same-pool
18/// work available because the callback already runs there.
19fn run_sequential_in_pool<T, U, E, F>(
20    pool: &rayon::ThreadPool,
21    items: Vec<T>,
22    f: F,
23) -> Result<Vec<U>, MapError>
24where
25    T: Send,
26    U: Send,
27    E: fmt::Display + Send,
28    F: Fn(T) -> Result<U, E> + Send + Sync,
29{
30    pool.install(|| map(items, f))
31}
32
33/// Why a [`RayonLanes`] executor could not be created for a domain.
34#[derive(Debug)]
35pub enum RayonLanesError {
36    /// The domain has no explicit Rayon pool.
37    MissingPool,
38    /// The call originated on a different Rayon pool.
39    ForeignPool,
40    /// Another operation already owns the domain admission slot.
41    DomainBusy(DomainBusy),
42}
43
44impl fmt::Display for RayonLanesError {
45    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
46        match self {
47            Self::MissingPool => formatter.write_str("domain has no Rayon pool"),
48            Self::ForeignPool => formatter.write_str("lanes entered from a foreign Rayon pool"),
49            Self::DomainBusy(error) => error.fmt(formatter),
50        }
51    }
52}
53
54impl std::error::Error for RayonLanesError {
55    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
56        match self {
57            Self::DomainBusy(error) => Some(error),
58            Self::MissingPool | Self::ForeignPool => None,
59        }
60    }
61}
62
63impl From<RayonLanesError> for MapInError {
64    fn from(error: RayonLanesError) -> Self {
65        match error {
66            RayonLanesError::MissingPool => Self::MissingPool,
67            RayonLanesError::ForeignPool => Self::ForeignPool,
68            RayonLanesError::DomainBusy(error) => Self::DomainBusy(error),
69        }
70    }
71}
72
73/// The [`Lanes`] implementation over a domain's Rayon pool.
74///
75/// Creating one admits the domain, so the executor holds the domain's one coarse
76/// operation slot for its whole lifetime. Each lane is one Rayon worker; the
77/// lane resource is `()` because the callback already runs on that worker.
78///
79/// Entry from a Rayon worker that does not belong to the domain pool is rejected,
80/// and so is entry while another operation owns the domain.
81///
82/// # Example
83///
84/// ```
85/// # #[cfg(feature = "rayon")]
86/// # {
87/// use hataori::{Domain, Lanes, RayonLanes};
88/// use std::sync::Arc;
89///
90/// let pool = Arc::new(rayon::ThreadPoolBuilder::new().num_threads(2).build().unwrap());
91/// let domain = Domain::external(Arc::clone(&pool), vec![0, 1], 2).unwrap();
92/// let mut lanes = RayonLanes::new(&domain).unwrap();
93/// assert_eq!(lanes.lane_count(), 2);
94/// let visited = std::sync::atomic::AtomicUsize::new(0);
95/// lanes
96///     .run(|_index, (), _control| {
97///         visited.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
98///         Ok::<_, std::convert::Infallible>(())
99///     })
100///     .unwrap();
101/// assert_eq!(visited.into_inner(), 2);
102/// # }
103/// ```
104pub struct RayonLanes<'d> {
105    pool: &'d rayon::ThreadPool,
106    worker_count: usize,
107    // Held for the executor's whole lifetime; dropping it releases admission.
108    _admission: DomainAdmission<'d>,
109}
110
111impl<'d> RayonLanes<'d> {
112    /// Admits `domain` and builds a lane executor over its pool.
113    ///
114    /// # Errors
115    ///
116    /// Returns [`RayonLanesError::MissingPool`] for a domain without a pool,
117    /// [`RayonLanesError::ForeignPool`] when called from a different Rayon pool,
118    /// and [`RayonLanesError::DomainBusy`] when the domain is already admitted.
119    pub fn new(domain: &'d Domain) -> Result<Self, RayonLanesError> {
120        let pool = domain.rayon_pool().ok_or(RayonLanesError::MissingPool)?;
121        if in_rayon_worker_context() && pool.current_thread_index().is_none() {
122            return Err(RayonLanesError::ForeignPool);
123        }
124        let admission = domain.try_admit().map_err(RayonLanesError::DomainBusy)?;
125        Ok(Self {
126            pool,
127            worker_count: domain.worker_count(),
128            _admission: admission,
129        })
130    }
131}
132
133/// What one lane reported back to the executor.
134enum LaneReport<E> {
135    Error(E),
136    Panic(Box<dyn Any + Send + 'static>),
137}
138
139fn record_lane_report<E>(report: &Mutex<Option<LaneReport<E>>>, value: LaneReport<E>) {
140    let mut report = report.lock().unwrap_or_else(PoisonError::into_inner);
141    // A panic payload always wins over an earlier ordinary lane error.
142    let replace = matches!(
143        (&*report, &value),
144        (None, _) | (Some(LaneReport::Error(_)), LaneReport::Panic(_))
145    );
146    if replace {
147        *report = Some(value);
148    }
149}
150
151impl Lanes for RayonLanes<'_> {
152    type Lane<'l> = ();
153    type Error = Infallible;
154
155    fn lane_count(&self) -> usize {
156        self.worker_count
157    }
158
159    fn run<E: Send>(
160        &mut self,
161        lane: impl Fn(usize, &mut (), &LaneControl<'_>) -> Result<(), E> + Sync,
162    ) -> Result<(), LanesError<E, Self::Error>> {
163        let control = LaneControl::new();
164        let report: Mutex<Option<LaneReport<E>>> = Mutex::new(None);
165        // `broadcast` returns only after every worker finished its callback, and
166        // a failing lane cancels the shared control so peers waiting on it stop.
167        self.pool.broadcast(|context| {
168            let outcome = std::panic::catch_unwind(AssertUnwindSafe(|| {
169                lane(context.index(), &mut (), &control)
170            }));
171            match outcome {
172                Ok(Ok(())) => {}
173                Ok(Err(error)) => {
174                    control.cancel();
175                    record_lane_report(&report, LaneReport::Error(error));
176                }
177                Err(payload) => {
178                    control.cancel();
179                    record_lane_report(&report, LaneReport::Panic(payload));
180                }
181            }
182        });
183        match report.into_inner().unwrap_or_else(PoisonError::into_inner) {
184            Some(LaneReport::Panic(payload)) => std::panic::resume_unwind(payload),
185            Some(LaneReport::Error(error)) => Err(LanesError::Lane(error)),
186            None => Ok(()),
187        }
188    }
189
190    fn run_with_coordinator<C, E: Send>(
191        &mut self,
192        coordinator: impl FnOnce(&LaneControl<'_>) -> Result<C, E>,
193        lane: impl Fn(usize, &mut (), &LaneControl<'_>) -> Result<(), E> + Sync,
194    ) -> Result<C, LanesError<E, Self::Error>> {
195        if self.pool.current_thread_index().is_some() {
196            // A coordinator on one of this pool's own workers cannot make
197            // progress while its lanes are joined.
198            return Err(LanesError::Reentrant);
199        }
200        let control = LaneControl::new();
201        let report: Mutex<Option<LaneReport<E>>> = Mutex::new(None);
202        // The calling thread stays outside the pool, so even a one-worker pool
203        // runs its lane concurrently with the coordinator.
204        let coordinator_outcome = self.pool.in_place_scope(|scope| {
205            let lane = &lane;
206            let control = &control;
207            let report = &report;
208            scope.spawn_broadcast(move |_scope, context| {
209                let outcome = std::panic::catch_unwind(AssertUnwindSafe(|| {
210                    lane(context.index(), &mut (), control)
211                }));
212                match outcome {
213                    Ok(Ok(())) => {}
214                    Ok(Err(error)) => {
215                        control.cancel();
216                        record_lane_report(report, LaneReport::Error(error));
217                    }
218                    Err(payload) => {
219                        control.cancel();
220                        record_lane_report(report, LaneReport::Panic(payload));
221                    }
222                }
223            });
224            let outcome = std::panic::catch_unwind(AssertUnwindSafe(|| coordinator(control)));
225            // Lanes stop when the coordinator stops feeding them.
226            control.cancel();
227            outcome
228        });
229
230        // `in_place_scope` joined every lane before returning.
231        let report = report.into_inner().unwrap_or_else(PoisonError::into_inner);
232        match coordinator_outcome {
233            Err(payload) => std::panic::resume_unwind(payload),
234            Ok(result) => match report {
235                Some(LaneReport::Panic(payload)) => std::panic::resume_unwind(payload),
236                Some(LaneReport::Error(error)) => Err(LanesError::Lane(error)),
237                None => result.map_err(LanesError::Coordinator),
238            },
239        }
240    }
241}
242
243/// Maps `items` through an explicit Rayon domain.
244///
245/// The target domain must have a pool, the call must not originate in another
246/// Rayon pool, and the domain must not already be admitted. These preconditions
247/// are checked in that order, before `f` is called. The callback and input,
248/// output, and error types must satisfy `Send`/`Sync` as required by Rayon, but
249/// none needs to be `'static`.
250///
251/// [`LocalMode::Sequential`] and [`LocalMode::Inner`] run callbacks one at a
252/// time in the target pool and stop at the first callback error. `Inner` leaves
253/// nested same-pool Rayon work available. [`LocalMode::Outer`] fans out through
254/// [`RayonLanes`] and [`map_in_lanes`]: every input is evaluated exactly once,
255/// result order is preserved, and the lowest callback-error index is reported
256/// only after all callbacks have finished.
257///
258/// # Errors
259///
260/// Returns [`MapInError::MissingPool`], [`MapInError::ForeignPool`], or
261/// [`MapInError::DomainBusy`] for the preconditions above,
262/// [`MapInError::Callback`] with the lowest failing input index when a callback
263/// fails, and [`MapInError::Executor`] when the lane executor cannot run.
264///
265/// # Example
266///
267/// ```
268/// # #[cfg(feature = "rayon")]
269/// # {
270/// use hataori::{map_in, Domain, LocalMode};
271/// use std::sync::Arc;
272/// use rayon::ThreadPoolBuilder;
273///
274/// let pool = Arc::new(ThreadPoolBuilder::new().num_threads(2).build().unwrap());
275/// let domain = Domain::external(Arc::clone(&pool), vec![0, 1], 2).unwrap();
276/// let result = map_in(&domain, LocalMode::Outer, vec![1, 2, 3], |item| {
277///     Ok::<_, std::convert::Infallible>(item * 2)
278/// }).unwrap();
279/// assert_eq!(result, vec![2, 4, 6]);
280/// # }
281/// ```
282pub fn map_in<T, U, E, F>(
283    domain: &Domain,
284    mode: LocalMode,
285    items: Vec<T>,
286    f: F,
287) -> Result<Vec<U>, MapInError>
288where
289    T: Send,
290    U: Send,
291    E: fmt::Display + Send,
292    F: Fn(T) -> Result<U, E> + Send + Sync,
293{
294    let mut lanes = RayonLanes::new(domain)?;
295    match mode {
296        LocalMode::Outer => map_in_lanes(&mut lanes, items, |item, ()| f(item)),
297        LocalMode::Sequential | LocalMode::Inner => {
298            let pool = domain.rayon_pool().ok_or(MapInError::MissingPool)?;
299            run_sequential_in_pool(pool, items, f).map_err(MapInError::Callback)
300        }
301    }
302}
303
304#[cfg(test)]
305mod tests;
hataori/local.rs

hataori/
local.rs
