Skip to content

Commit

Permalink
Add the ability to spawn futures
Browse files Browse the repository at this point in the history
For Rust 1.36+ with `std::future::Future`, add a way to spawn tasks with
a returned `Future`. The task is immediately queued for the thread pool
to execute.
  • Loading branch information
cuviper committed Aug 2, 2019
1 parent 678ffbb commit e3df22b
Show file tree
Hide file tree
Showing 4 changed files with 154 additions and 1 deletion.
3 changes: 3 additions & 0 deletions rayon-core/Cargo.toml
Expand Up @@ -55,3 +55,6 @@ path = "tests/simple_panic.rs"
[[test]]
name = "scoped_threadpool"
path = "tests/scoped_threadpool.rs"

[build-dependencies]
autocfg = "0.1.5"
7 changes: 6 additions & 1 deletion rayon-core/build.rs
@@ -1,7 +1,12 @@
extern crate autocfg;

// We need a build script to use `link = "rayon-core"`. But we're not
// *actually* linking to anything, just making sure that we're the only
// rayon-core in use.
fn main() {
let ac = autocfg::new();
ac.emit_path_cfg("std::future::Future", "has_future");

// we don't need to rebuild for anything else
println!("cargo:rerun-if-changed=build.rs");
autocfg::rerun_path("build.rs");
}
139 changes: 139 additions & 0 deletions rayon-core/src/future.rs
@@ -0,0 +1,139 @@
#![allow(missing_docs)]

use crate::ThreadPool;
use crate::{spawn, spawn_fifo};
use crate::{Scope, ScopeFifo};

use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};

use job::JobResult;
use unwind;

struct RayonFuture<T> {
state: Arc<Mutex<State<T>>>,
}

struct RayonFutureJob<T> {
state: Arc<Mutex<State<T>>>,
}

struct State<T> {
result: JobResult<T>,
waker: Option<Waker>,
}

fn new<T>() -> (RayonFuture<T>, RayonFutureJob<T>) {
let state = Arc::new(Mutex::new(State {
result: JobResult::None,
waker: None,
}));
(
RayonFuture {
state: state.clone(),
},
RayonFutureJob { state },
)
}

impl<T> Future for RayonFuture<T> {
type Output = T;

fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let mut guard = self.state.lock().expect("rayon future lock");
match mem::replace(&mut guard.result, JobResult::None) {
JobResult::None => {
guard.waker = Some(cx.waker().clone());
Poll::Pending
}
JobResult::Ok(x) => Poll::Ready(x),
JobResult::Panic(p) => {
drop(guard); // don't poison the lock
unwind::resume_unwinding(p);
}
}
}
}

impl<T> RayonFutureJob<T> {
fn execute(self, func: impl FnOnce() -> T) {
let result = unwind::halt_unwinding(func);
let mut guard = self.state.lock().expect("rayon future lock");
guard.result = match result {
Ok(x) => JobResult::Ok(x),
Err(p) => JobResult::Panic(p),
};
if let Some(waker) = guard.waker.take() {
waker.wake();
}
}
}

pub fn spawn_future<F, T>(func: F) -> impl Future<Output = T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let (future, job) = new();
spawn(move || job.execute(func));
future
}

pub fn spawn_fifo_future<F, T>(func: F) -> impl Future<Output = T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let (future, job) = new();
spawn_fifo(move || job.execute(func));
future
}

impl ThreadPool {
pub fn spawn_future<F, T>(&self, func: F) -> impl Future<Output = T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let (future, job) = new();
self.spawn(move || job.execute(func));
future
}

pub fn spawn_fifo_future<F, T>(&self, func: F) -> impl Future<Output = T>
where
F: FnOnce() -> T + Send + 'static,
T: Send + 'static,
{
let (future, job) = new();
self.spawn_fifo(move || job.execute(func));
future
}
}

impl<'scope> Scope<'scope> {
pub fn spawn_future<F, T>(&self, func: F) -> impl Future<Output = T>
where
F: FnOnce(&Self) -> T + Send + 'scope,
T: Send + 'scope,
{
let (future, job) = new();
self.spawn(|scope| job.execute(move || func(scope)));
future
}
}

impl<'scope> ScopeFifo<'scope> {
pub fn spawn_fifo_future<F, T>(&self, func: F) -> impl Future<Output = T>
where
F: FnOnce(&Self) -> T + Send + 'scope,
T: Send + 'scope,
{
let (future, job) = new();
self.spawn_fifo(|scope| job.execute(move || func(scope)));
future
}
}
6 changes: 6 additions & 0 deletions rayon-core/src/lib.rs
Expand Up @@ -61,6 +61,9 @@ mod thread_pool;
mod unwind;
mod util;

#[cfg(has_future)]
mod future;

mod compile_fail;
mod test;

Expand All @@ -75,6 +78,9 @@ pub use thread_pool::current_thread_has_pending_tasks;
pub use thread_pool::current_thread_index;
pub use thread_pool::ThreadPool;

#[cfg(has_future)]
pub use future::{spawn_future, spawn_fifo_future};

use registry::{CustomSpawn, DefaultSpawn, ThreadSpawn};

/// Returns the number of threads in the current registry. If this
Expand Down

0 comments on commit e3df22b

Please sign in to comment.