Implement simple work stealing
Stjepan Glavina
5 years ago
parent
dd92d8dc61
commit
0924911ac3
@ -0,0 +1,32 @@
|
||||
use super::pool;
|
||||
use super::JoinHandle;
|
||||
use crate::future::Future;
|
||||
use crate::io;
|
||||
|
||||
/// Task builder that configures the settings of a new task.
|
||||
#[derive(Debug)]
|
||||
pub struct Builder {
|
||||
pub(crate) name: Option<String>,
|
||||
}
|
||||
|
||||
impl Builder {
|
||||
/// Creates a new builder.
|
||||
pub fn new() -> Builder {
|
||||
Builder { name: None }
|
||||
}
|
||||
|
||||
/// Configures the name of the task.
|
||||
pub fn name(mut self, name: String) -> Builder {
|
||||
self.name = Some(name);
|
||||
self
|
||||
}
|
||||
|
||||
/// Spawns a task with the configured settings.
|
||||
pub fn spawn<F, T>(self, future: F) -> io::Result<JoinHandle<T>>
|
||||
where
|
||||
F: Future<Output = T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
Ok(pool::get().spawn(future, self))
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,52 @@
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Condvar, Mutex};
|
||||
|
||||
/// The place where worker threads go to sleep.
|
||||
///
|
||||
/// Similar to how thread parking works, if a notification comes up while no threads are sleeping,
|
||||
/// the next thread that attempts to go to sleep will pick up the notification immediately.
|
||||
pub struct Sleepers {
|
||||
/// How many threads are currently a sleep.
|
||||
sleep: Mutex<usize>,
|
||||
|
||||
/// A condvar for notifying sleeping threads.
|
||||
wake: Condvar,
|
||||
|
||||
/// Set to `true` if a notification came up while nobody was sleeping.
|
||||
notified: AtomicBool,
|
||||
}
|
||||
|
||||
impl Sleepers {
|
||||
/// Creates a new `Sleepers`.
|
||||
pub fn new() -> Sleepers {
|
||||
Sleepers {
|
||||
sleep: Mutex::new(0),
|
||||
wake: Condvar::new(),
|
||||
notified: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
/// Puts the current thread to sleep.
|
||||
pub fn wait(&self) {
|
||||
let mut sleep = self.sleep.lock().unwrap();
|
||||
|
||||
if self.notified.swap(false, Ordering::SeqCst) == false {
|
||||
*sleep += 1;
|
||||
let _ = self.wake.wait(sleep).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
/// Notifies one thread.
|
||||
pub fn notify_one(&self) {
|
||||
if self.notified.load(Ordering::SeqCst) == false {
|
||||
let mut sleep = self.sleep.lock().unwrap();
|
||||
|
||||
if *sleep > 0 {
|
||||
*sleep -= 1;
|
||||
self.wake.notify_one();
|
||||
} else {
|
||||
self.notified.store(true, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,110 @@
|
||||
use std::cell::Cell;
|
||||
use std::ptr;
|
||||
|
||||
use crossbeam_deque::Worker;
|
||||
|
||||
use super::pool;
|
||||
use super::task;
|
||||
use super::Task;
|
||||
use crate::utils::abort_on_panic;
|
||||
|
||||
/// Returns a handle to the current task.
|
||||
///
|
||||
/// # Panics
|
||||
///
|
||||
/// This function will panic if not called within the context of a task created by [`block_on`],
|
||||
/// [`spawn`], or [`Builder::spawn`].
|
||||
///
|
||||
/// [`block_on`]: fn.block_on.html
|
||||
/// [`spawn`]: fn.spawn.html
|
||||
/// [`Builder::spawn`]: struct.Builder.html#method.spawn
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # fn main() { async_std::task::block_on(async {
|
||||
/// #
|
||||
/// use async_std::task;
|
||||
///
|
||||
/// println!("The name of this task is {:?}", task::current().name());
|
||||
/// #
|
||||
/// # }) }
|
||||
/// ```
|
||||
pub fn current() -> Task {
|
||||
get_task(|task| task.clone()).expect("`task::current()` called outside the context of a task")
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
static TAG: Cell<*const task::Tag> = Cell::new(ptr::null_mut());
|
||||
}
|
||||
|
||||
pub(crate) fn set_tag<F, R>(tag: *const task::Tag, f: F) -> R
|
||||
where
|
||||
F: FnOnce() -> R,
|
||||
{
|
||||
struct ResetTag<'a>(&'a Cell<*const task::Tag>);
|
||||
|
||||
impl Drop for ResetTag<'_> {
|
||||
fn drop(&mut self) {
|
||||
self.0.set(ptr::null());
|
||||
}
|
||||
}
|
||||
|
||||
TAG.with(|t| {
|
||||
t.set(tag);
|
||||
let _guard = ResetTag(t);
|
||||
|
||||
f()
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn get_task<F, R>(f: F) -> Option<R>
|
||||
where
|
||||
F: FnOnce(&Task) -> R,
|
||||
{
|
||||
let res = TAG.try_with(|tag| unsafe { tag.get().as_ref().map(task::Tag::task).map(f) });
|
||||
|
||||
match res {
|
||||
Ok(Some(val)) => Some(val),
|
||||
Ok(None) | Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
thread_local! {
|
||||
static IS_WORKER: Cell<bool> = Cell::new(false);
|
||||
static QUEUE: Cell<Option<Worker<task::Runnable>>> = Cell::new(None);
|
||||
}
|
||||
|
||||
pub fn is_worker() -> bool {
|
||||
IS_WORKER.with(|is_worker| is_worker.get())
|
||||
}
|
||||
|
||||
fn get_queue<F: FnOnce(&Worker<task::Runnable>) -> T, T>(f: F) -> T {
|
||||
QUEUE.with(|queue| {
|
||||
let q = queue.take().unwrap();
|
||||
let ret = f(&q);
|
||||
queue.set(Some(q));
|
||||
ret
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn schedule(task: task::Runnable) {
|
||||
if is_worker() {
|
||||
get_queue(|q| q.push(task));
|
||||
} else {
|
||||
pool::get().injector.push(task);
|
||||
}
|
||||
pool::get().sleepers.notify_one();
|
||||
}
|
||||
|
||||
pub(crate) fn main_loop(worker: Worker<task::Runnable>) {
|
||||
IS_WORKER.with(|is_worker| is_worker.set(true));
|
||||
QUEUE.with(|queue| queue.set(Some(worker)));
|
||||
|
||||
loop {
|
||||
match get_queue(|q| pool::get().find_task(q)) {
|
||||
Some(task) => set_tag(task.tag(), || abort_on_panic(|| task.run())),
|
||||
None => pool::get().sleepers.wait(),
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue