Merge rayon-rs#1026

1026: Add `yield_now` and `yield_local` r=cuviper a=cuviper

* `yield_now` looks for work anywhere in the thread pool and executes it.
* `yield_local` only looks in the thread-local deque (including broadcasts).

In either case, they return:
* `Some(Yield::Executed)` if work was found and executed.
* `Some(Yield::Idle)` if no work was found in their pool/thread.
* `None` if the current thread is not part of a thread pool.

These do not call `std::thread::yield_now()`, but the docs suggest polling loops might want to add that as a fallback.

Co-authored-by: Josh Stone <cuviper@gmail.com>

Author: bors[bot]

rayon-core/src/lib.rs

@@ -38,7 +38,8 @@
 //! there is no other thread to share the work. However, since the pool is not running independent
 //! of the main thread, non-blocking calls like `spawn` may not execute at all, unless a lower-
 //! priority call like `broadcast` gives them an opening. The fallback mode does not try to emulate
-//! anything like thread preemption or `async` task switching.
+//! anything like thread preemption or `async` task switching, but `yield_now` or `yield_local`
+//! can also volunteer execution time.
 //!
 //! Explicit `ThreadPoolBuilder` methods always report their error without any fallback.
 //!
@@ -101,6 +102,7 @@ pub use self::spawn::{spawn, spawn_fifo};
 pub use self::thread_pool::current_thread_has_pending_tasks;
 pub use self::thread_pool::current_thread_index;
 pub use self::thread_pool::ThreadPool;
+pub use self::thread_pool::{yield_local, yield_now, Yield};
 
 use self::registry::{CustomSpawn, DefaultSpawn, ThreadSpawn};
 

rayon-core/src/registry.rs

@@ -6,6 +6,7 @@ use crate::sleep::Sleep;
 use crate::unwind;
 use crate::{
     ErrorKind, ExitHandler, PanicHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder,
+    Yield,
 };
 use crossbeam_deque::{Injector, Steal, Stealer, Worker};
 use std::cell::Cell;
@@ -772,7 +773,7 @@ impl WorkerThread {
     /// for breadth-first execution, it would mean dequeuing from the
     /// bottom.
     #[inline]
-    pub(super) unsafe fn take_local_job(&self) -> Option<JobRef> {
+    pub(super) fn take_local_job(&self) -> Option<JobRef> {
         let popped_job = self.worker.pop();
 
         if popped_job.is_some() {
@@ -814,16 +815,7 @@ impl WorkerThread {
 
         let mut idle_state = self.registry.sleep.start_looking(self.index, latch);
         while !latch.probe() {
-            // Try to find some work to do. We give preference first
-            // to things in our local deque, then in other workers
-            // deques, and finally to injected jobs from the
-            // outside. The idea is to finish what we started before
-            // we take on something new.
-            if let Some(job) = self
-                .take_local_job()
-                .or_else(|| self.steal())
-                .or_else(|| self.registry.pop_injected_job(self.index))
-            {
+            if let Some(job) = self.find_work() {
                 self.registry.sleep.work_found(idle_state);
                 self.execute(job);
                 idle_state = self.registry.sleep.start_looking(self.index, latch);
@@ -846,6 +838,37 @@ impl WorkerThread {
         mem::forget(abort_guard); // successful execution, do not abort
     }
 
+    fn find_work(&self) -> Option<JobRef> {
+        // Try to find some work to do. We give preference first
+        // to things in our local deque, then in other workers
+        // deques, and finally to injected jobs from the
+        // outside. The idea is to finish what we started before
+        // we take on something new.
+        self.take_local_job()
+            .or_else(|| self.steal())
+            .or_else(|| self.registry.pop_injected_job(self.index))
+    }
+
+    pub(super) fn yield_now(&self) -> Yield {
+        match self.find_work() {
+            Some(job) => unsafe {
+                self.execute(job);
+                Yield::Executed
+            },
+            None => Yield::Idle,
+        }
+    }
+
+    pub(super) fn yield_local(&self) -> Yield {
+        match self.take_local_job() {
+            Some(job) => unsafe {
+                self.execute(job);
+                Yield::Executed
+            },
+            None => Yield::Idle,
+        }
+    }
+
     #[inline]
     pub(super) unsafe fn execute(&self, job: JobRef) {
         job.execute();
@@ -855,7 +878,7 @@ impl WorkerThread {
     ///
     /// This should only be done as a last resort, when there is no
     /// local work to do.
-    unsafe fn steal(&self) -> Option<JobRef> {
+    fn steal(&self) -> Option<JobRef> {
         // we only steal when we don't have any work to do locally
         debug_assert!(self.local_deque_is_empty());
 

rayon-core/src/thread_pool/mod.rs

@@ -339,6 +339,30 @@ impl ThreadPool {
         // We assert that `self.registry` has not terminated.
         unsafe { broadcast::spawn_broadcast_in(op, &self.registry) }
     }
+
+    /// Cooperatively yields execution to Rayon.
+    ///
+    /// This is similar to the general [`yield_now()`], but only if the current
+    /// thread is part of *this* thread pool.
+    ///
+    /// Returns `Some(Yield::Executed)` if anything was executed, `Some(Yield::Idle)` if
+    /// nothing was available, or `None` if the current thread is not part this pool.
+    pub fn yield_now(&self) -> Option<Yield> {
+        let curr = self.registry.current_thread()?;
+        Some(curr.yield_now())
+    }
+
+    /// Cooperatively yields execution to local Rayon work.
+    ///
+    /// This is similar to the general [`yield_local()`], but only if the current
+    /// thread is part of *this* thread pool.
+    ///
+    /// Returns `Some(Yield::Executed)` if anything was executed, `Some(Yield::Idle)` if
+    /// nothing was available, or `None` if the current thread is not part this pool.
+    pub fn yield_local(&self) -> Option<Yield> {
+        let curr = self.registry.current_thread()?;
+        Some(curr.yield_local())
+    }
 }
 
 impl Drop for ThreadPool {
@@ -400,3 +424,48 @@ pub fn current_thread_has_pending_tasks() -> Option<bool> {
         Some(!curr.local_deque_is_empty())
     }
 }
+
+/// Cooperatively yields execution to Rayon.
+///
+/// If the current thread is part of a rayon thread pool, this looks for a
+/// single unit of pending work in the pool, then executes it. Completion of
+/// that work might include nested work or further work stealing.
+///
+/// This is similar to [`std::thread::yield_now()`], but does not literally make
+/// that call. If you are implementing a polling loop, you may want to also
+/// yield to the OS scheduler yourself if no Rayon work was found.
+///
+/// Returns `Some(Yield::Executed)` if anything was executed, `Some(Yield::Idle)` if
+/// nothing was available, or `None` if this thread is not part of any pool at all.
+pub fn yield_now() -> Option<Yield> {
+    unsafe {
+        let thread = WorkerThread::current().as_ref()?;
+        Some(thread.yield_now())
+    }
+}
+
+/// Cooperatively yields execution to local Rayon work.
+///
+/// If the current thread is part of a rayon thread pool, this looks for a
+/// single unit of pending work in this thread's queue, then executes it.
+/// Completion of that work might include nested work or further work stealing.
+///
+/// This is similar to [`yield_now()`], but does not steal from other threads.
+///
+/// Returns `Some(Yield::Executed)` if anything was executed, `Some(Yield::Idle)` if
+/// nothing was available, or `None` if this thread is not part of any pool at all.
+pub fn yield_local() -> Option<Yield> {
+    unsafe {
+        let thread = WorkerThread::current().as_ref()?;
+        Some(thread.yield_local())
+    }
+}
+
+/// Result of [`yield_now()`] or [`yield_local()`].
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum Yield {
+    /// Work was found and executed.
+    Executed,
+    /// No available work was found.
+    Idle,
+}

rayon-core/src/thread_pool/test.rs

@@ -380,3 +380,39 @@ fn in_place_scope_fifo_no_deadlock() {
         rx_ref.recv().unwrap();
     });
 }
+
+#[test]
+fn yield_now_to_spawn() {
+    let (tx, rx) = crossbeam_channel::bounded(1);
+
+    // Queue a regular spawn.
+    crate::spawn(move || tx.send(22).unwrap());
+
+    // The single-threaded fallback mode (for wasm etc.) won't
+    // get a chance to run the spawn if we never yield to it.
+    crate::registry::in_worker(move |_, _| {
+        crate::yield_now();
+    });
+
+    // The spawn **must** have started by now, but we still might have to wait
+    // for it to finish if a different thread stole it first.
+    assert_eq!(22, rx.recv().unwrap());
+}
+
+#[test]
+fn yield_local_to_spawn() {
+    let (tx, rx) = crossbeam_channel::bounded(1);
+
+    // Queue a regular spawn.
+    crate::spawn(move || tx.send(22).unwrap());
+
+    // The single-threaded fallback mode (for wasm etc.) won't
+    // get a chance to run the spawn if we never yield to it.
+    crate::registry::in_worker(move |_, _| {
+        crate::yield_local();
+    });
+
+    // The spawn **must** have started by now, but we still might have to wait
+    // for it to finish if a different thread stole it first.
+    assert_eq!(22, rx.recv().unwrap());
+}

src/lib.rs

@@ -124,6 +124,7 @@ pub use rayon_core::{in_place_scope, scope, Scope};
 pub use rayon_core::{in_place_scope_fifo, scope_fifo, ScopeFifo};
 pub use rayon_core::{join, join_context};
 pub use rayon_core::{spawn, spawn_fifo};
+pub use rayon_core::{yield_local, yield_now, Yield};
 
 /// We need to transmit raw pointers across threads. It is possible to do this
 /// without any unsafe code by converting pointers to usize or to AtomicPtr<T>