rayon_core/
job.rs

1use crate::latch::Latch;
2use crate::unwind;
3use crossbeam_deque::{Injector, Steal};
4use std::any::Any;
5use std::cell::UnsafeCell;
6use std::mem;
7use std::sync::Arc;
8
9pub(super) enum JobResult<T> {
10    None,
11    Ok(T),
12    Panic(Box<dyn Any + Send>),
13}
14
15/// A `Job` is used to advertise work for other threads that they may
16/// want to steal. In accordance with time honored tradition, jobs are
17/// arranged in a deque, so that thieves can take from the top of the
18/// deque while the main worker manages the bottom of the deque. This
19/// deque is managed by the `thread_pool` module.
20pub(super) trait Job {
21    /// Unsafe: this may be called from a different thread than the one
22    /// which scheduled the job, so the implementer must ensure the
23    /// appropriate traits are met, whether `Send`, `Sync`, or both.
24    unsafe fn execute(this: *const ());
25}
26
27/// Effectively a Job trait object. Each JobRef **must** be executed
28/// exactly once, or else data may leak.
29///
30/// Internally, we store the job's data in a `*const ()` pointer.  The
31/// true type is something like `*const StackJob<...>`, but we hide
32/// it. We also carry the "execute fn" from the `Job` trait.
33pub(super) struct JobRef {
34    pointer: *const (),
35    execute_fn: unsafe fn(*const ()),
36}
37
38unsafe impl Send for JobRef {}
39unsafe impl Sync for JobRef {}
40
41impl JobRef {
42    /// Unsafe: caller asserts that `data` will remain valid until the
43    /// job is executed.
44    pub(super) unsafe fn new<T>(data: *const T) -> JobRef
45    where
46        T: Job,
47    {
48        // erase types:
49        JobRef {
50            pointer: data as *const (),
51            execute_fn: <T as Job>::execute,
52        }
53    }
54
55    /// Returns an opaque handle that can be saved and compared,
56    /// without making `JobRef` itself `Copy + Eq`.
57    #[inline]
58    pub(super) fn id(&self) -> impl Eq {
59        (self.pointer, self.execute_fn)
60    }
61
62    #[inline]
63    pub(super) unsafe fn execute(self) {
64        (self.execute_fn)(self.pointer)
65    }
66}
67
68/// A job that will be owned by a stack slot. This means that when it
69/// executes it need not free any heap data, the cleanup occurs when
70/// the stack frame is later popped.  The function parameter indicates
71/// `true` if the job was stolen -- executed on a different thread.
72pub(super) struct StackJob<L, F, R>
73where
74    L: Latch + Sync,
75    F: FnOnce(bool) -> R + Send,
76    R: Send,
77{
78    pub(super) latch: L,
79    func: UnsafeCell<Option<F>>,
80    result: UnsafeCell<JobResult<R>>,
81}
82
83impl<L, F, R> StackJob<L, F, R>
84where
85    L: Latch + Sync,
86    F: FnOnce(bool) -> R + Send,
87    R: Send,
88{
89    pub(super) fn new(func: F, latch: L) -> StackJob<L, F, R> {
90        StackJob {
91            latch,
92            func: UnsafeCell::new(Some(func)),
93            result: UnsafeCell::new(JobResult::None),
94        }
95    }
96
97    pub(super) unsafe fn as_job_ref(&self) -> JobRef {
98        JobRef::new(self)
99    }
100
101    pub(super) unsafe fn run_inline(self, stolen: bool) -> R {
102        self.func.into_inner().unwrap()(stolen)
103    }
104
105    pub(super) unsafe fn into_result(self) -> R {
106        self.result.into_inner().into_return_value()
107    }
108}
109
110impl<L, F, R> Job for StackJob<L, F, R>
111where
112    L: Latch + Sync,
113    F: FnOnce(bool) -> R + Send,
114    R: Send,
115{
116    unsafe fn execute(this: *const ()) {
117        let this = &*(this as *const Self);
118        let abort = unwind::AbortIfPanic;
119        let func = (*this.func.get()).take().unwrap();
120        (*this.result.get()) = JobResult::call(func);
121        Latch::set(&this.latch);
122        mem::forget(abort);
123    }
124}
125
126/// Represents a job stored in the heap. Used to implement
127/// `scope`. Unlike `StackJob`, when executed, `HeapJob` simply
128/// invokes a closure, which then triggers the appropriate logic to
129/// signal that the job executed.
130///
131/// (Probably `StackJob` should be refactored in a similar fashion.)
132pub(super) struct HeapJob<BODY>
133where
134    BODY: FnOnce() + Send,
135{
136    job: BODY,
137}
138
139impl<BODY> HeapJob<BODY>
140where
141    BODY: FnOnce() + Send,
142{
143    pub(super) fn new(job: BODY) -> Box<Self> {
144        Box::new(HeapJob { job })
145    }
146
147    /// Creates a `JobRef` from this job -- note that this hides all
148    /// lifetimes, so it is up to you to ensure that this JobRef
149    /// doesn't outlive any data that it closes over.
150    pub(super) unsafe fn into_job_ref(self: Box<Self>) -> JobRef {
151        JobRef::new(Box::into_raw(self))
152    }
153
154    /// Creates a static `JobRef` from this job.
155    pub(super) fn into_static_job_ref(self: Box<Self>) -> JobRef
156    where
157        BODY: 'static,
158    {
159        unsafe { self.into_job_ref() }
160    }
161}
162
163impl<BODY> Job for HeapJob<BODY>
164where
165    BODY: FnOnce() + Send,
166{
167    unsafe fn execute(this: *const ()) {
168        let this = Box::from_raw(this as *mut Self);
169        (this.job)();
170    }
171}
172
173/// Represents a job stored in an `Arc` -- like `HeapJob`, but may
174/// be turned into multiple `JobRef`s and called multiple times.
175pub(super) struct ArcJob<BODY>
176where
177    BODY: Fn() + Send + Sync,
178{
179    job: BODY,
180}
181
182impl<BODY> ArcJob<BODY>
183where
184    BODY: Fn() + Send + Sync,
185{
186    pub(super) fn new(job: BODY) -> Arc<Self> {
187        Arc::new(ArcJob { job })
188    }
189
190    /// Creates a `JobRef` from this job -- note that this hides all
191    /// lifetimes, so it is up to you to ensure that this JobRef
192    /// doesn't outlive any data that it closes over.
193    pub(super) unsafe fn as_job_ref(this: &Arc<Self>) -> JobRef {
194        JobRef::new(Arc::into_raw(Arc::clone(this)))
195    }
196
197    /// Creates a static `JobRef` from this job.
198    pub(super) fn as_static_job_ref(this: &Arc<Self>) -> JobRef
199    where
200        BODY: 'static,
201    {
202        unsafe { Self::as_job_ref(this) }
203    }
204}
205
206impl<BODY> Job for ArcJob<BODY>
207where
208    BODY: Fn() + Send + Sync,
209{
210    unsafe fn execute(this: *const ()) {
211        let this = Arc::from_raw(this as *mut Self);
212        (this.job)();
213    }
214}
215
216impl<T> JobResult<T> {
217    fn call(func: impl FnOnce(bool) -> T) -> Self {
218        match unwind::halt_unwinding(|| func(true)) {
219            Ok(x) => JobResult::Ok(x),
220            Err(x) => JobResult::Panic(x),
221        }
222    }
223
224    /// Convert the `JobResult` for a job that has finished (and hence
225    /// its JobResult is populated) into its return value.
226    ///
227    /// NB. This will panic if the job panicked.
228    pub(super) fn into_return_value(self) -> T {
229        match self {
230            JobResult::None => unreachable!(),
231            JobResult::Ok(x) => x,
232            JobResult::Panic(x) => unwind::resume_unwinding(x),
233        }
234    }
235}
236
237/// Indirect queue to provide FIFO job priority.
238pub(super) struct JobFifo {
239    inner: Injector<JobRef>,
240}
241
242impl JobFifo {
243    pub(super) fn new() -> Self {
244        JobFifo {
245            inner: Injector::new(),
246        }
247    }
248
249    pub(super) unsafe fn push(&self, job_ref: JobRef) -> JobRef {
250        // A little indirection ensures that spawns are always prioritized in FIFO order.  The
251        // jobs in a thread's deque may be popped from the back (LIFO) or stolen from the front
252        // (FIFO), but either way they will end up popping from the front of this queue.
253        self.inner.push(job_ref);
254        JobRef::new(self)
255    }
256}
257
258impl Job for JobFifo {
259    unsafe fn execute(this: *const ()) {
260        // We "execute" a queue by executing its first job, FIFO.
261        let this = &*(this as *const Self);
262        loop {
263            match this.inner.steal() {
264                Steal::Success(job_ref) => break job_ref.execute(),
265                Steal::Empty => panic!("FIFO is empty"),
266                Steal::Retry => {}
267            }
268        }
269    }
270}