forked from mirror/async-std
Refactor the task module (#421)
* Refactor the task module * Fix clippy warning * Simplify task-local entries * Reduce the amount of future wrapping * Cleanup * Simplify stealingpoc-serde-support
Stjepan Glavina
5 years ago
committed by
GitHub
parent
c1e8517959
commit
3dd59d7056
@ -0,0 +1,28 @@
|
||||
use crate::task::Task;
|
||||
|
||||
/// 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
|
||||
///
|
||||
/// ```
|
||||
/// # async_std::task::block_on(async {
|
||||
/// #
|
||||
/// use async_std::task;
|
||||
///
|
||||
/// println!("The name of this task is {:?}", task::current().name());
|
||||
/// #
|
||||
/// # })
|
||||
/// ```
|
||||
pub fn current() -> Task {
|
||||
Task::get_current(|t| t.clone())
|
||||
.expect("`task::current()` called outside the context of a task")
|
||||
}
|
||||
@ -0,0 +1,13 @@
|
||||
//! Task executor.
|
||||
//!
|
||||
//! API bindings between `crate::task` and this module are very simple:
|
||||
//!
|
||||
//! * The only export is the `schedule` function.
|
||||
//! * The only import is the `crate::task::Runnable` type.
|
||||
|
||||
pub(crate) use pool::schedule;
|
||||
|
||||
use sleepers::Sleepers;
|
||||
|
||||
mod pool;
|
||||
mod sleepers;
|
||||
@ -0,0 +1,140 @@
|
||||
use std::cell::UnsafeCell;
|
||||
use std::iter;
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use crossbeam_deque::{Injector, Stealer, Worker};
|
||||
use once_cell::sync::Lazy;
|
||||
|
||||
use crate::task::executor::Sleepers;
|
||||
use crate::task::Runnable;
|
||||
use crate::utils::{abort_on_panic, random};
|
||||
|
||||
/// The state of an executor.
|
||||
struct Pool {
|
||||
/// The global queue of tasks.
|
||||
injector: Injector<Runnable>,
|
||||
|
||||
/// Handles to local queues for stealing work from worker threads.
|
||||
stealers: Vec<Stealer<Runnable>>,
|
||||
|
||||
/// Used for putting idle workers to sleep and notifying them when new tasks come in.
|
||||
sleepers: Sleepers,
|
||||
}
|
||||
|
||||
/// Global executor that runs spawned tasks.
|
||||
static POOL: Lazy<Pool> = Lazy::new(|| {
|
||||
let num_threads = num_cpus::get().max(1);
|
||||
let mut stealers = Vec::new();
|
||||
|
||||
// Spawn worker threads.
|
||||
for _ in 0..num_threads {
|
||||
let worker = Worker::new_fifo();
|
||||
stealers.push(worker.stealer());
|
||||
|
||||
thread::Builder::new()
|
||||
.name("async-std/executor".to_string())
|
||||
.spawn(|| abort_on_panic(|| main_loop(worker)))
|
||||
.expect("cannot start a thread driving tasks");
|
||||
}
|
||||
|
||||
Pool {
|
||||
injector: Injector::new(),
|
||||
stealers,
|
||||
sleepers: Sleepers::new(),
|
||||
}
|
||||
});
|
||||
|
||||
thread_local! {
|
||||
/// Local task queue associated with the current worker thread.
|
||||
static QUEUE: UnsafeCell<Option<Worker<Runnable>>> = UnsafeCell::new(None);
|
||||
}
|
||||
|
||||
/// Schedules a new runnable task for execution.
|
||||
pub(crate) fn schedule(task: Runnable) {
|
||||
QUEUE.with(|queue| {
|
||||
let local = unsafe { (*queue.get()).as_ref() };
|
||||
|
||||
// If the current thread is a worker thread, push the task into its local task queue.
|
||||
// Otherwise, push it into the global task queue.
|
||||
match local {
|
||||
None => POOL.injector.push(task),
|
||||
Some(q) => q.push(task),
|
||||
}
|
||||
});
|
||||
|
||||
// Notify a sleeping worker that new work just came in.
|
||||
POOL.sleepers.notify_one();
|
||||
}
|
||||
|
||||
/// Main loop running a worker thread.
|
||||
fn main_loop(local: Worker<Runnable>) {
|
||||
// Initialize the local task queue.
|
||||
QUEUE.with(|queue| unsafe { *queue.get() = Some(local) });
|
||||
|
||||
// The number of times the thread didn't find work in a row.
|
||||
let mut step = 0;
|
||||
|
||||
loop {
|
||||
// Try to find a runnable task.
|
||||
match find_runnable() {
|
||||
Some(task) => {
|
||||
// Found. Now run the task.
|
||||
task.run();
|
||||
step = 0;
|
||||
}
|
||||
None => {
|
||||
// Yield the current thread or put it to sleep.
|
||||
match step {
|
||||
0..=2 => {
|
||||
thread::yield_now();
|
||||
step += 1;
|
||||
}
|
||||
3 => {
|
||||
thread::sleep(Duration::from_micros(10));
|
||||
step += 1;
|
||||
}
|
||||
_ => {
|
||||
POOL.sleepers.wait();
|
||||
step = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the next runnable task.
|
||||
fn find_runnable() -> Option<Runnable> {
|
||||
let pool = &*POOL;
|
||||
|
||||
QUEUE.with(|queue| {
|
||||
let local = unsafe { (*queue.get()).as_ref().unwrap() };
|
||||
|
||||
// Pop a task from the local queue, if not empty.
|
||||
local.pop().or_else(|| {
|
||||
// Otherwise, we need to look for a task elsewhere.
|
||||
iter::repeat_with(|| {
|
||||
// Try stealing a batch of tasks from the global queue.
|
||||
pool.injector
|
||||
.steal_batch_and_pop(&local)
|
||||
// Or try stealing a batch of tasks from one of the other threads.
|
||||
.or_else(|| {
|
||||
// First, pick a random starting point in the list of local queues.
|
||||
let len = pool.stealers.len();
|
||||
let start = random(len as u32) as usize;
|
||||
|
||||
// Try stealing a batch of tasks from each local queue starting from the
|
||||
// chosen point.
|
||||
let (l, r) = pool.stealers.split_at(start);
|
||||
let rotated = r.iter().chain(l.iter());
|
||||
rotated.map(|s| s.steal_batch_and_pop(&local)).collect()
|
||||
})
|
||||
})
|
||||
// Loop while no task was stolen and any steal operation needs to be retried.
|
||||
.find(|s| !s.is_retry())
|
||||
// Extract the stolen task, if there is one.
|
||||
.and_then(|s| s.success())
|
||||
})
|
||||
})
|
||||
}
|
||||
@ -0,0 +1,56 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use crate::task::{Context, Poll, Task};
|
||||
|
||||
/// A handle that awaits the result of a task.
|
||||
///
|
||||
/// Dropping a [`JoinHandle`] will detach the task, meaning that there is no longer
|
||||
/// a handle to the task and no way to `join` on it.
|
||||
///
|
||||
/// Created when a task is [spawned].
|
||||
///
|
||||
/// [spawned]: fn.spawn.html
|
||||
#[derive(Debug)]
|
||||
pub struct JoinHandle<T>(async_task::JoinHandle<T, Task>);
|
||||
|
||||
unsafe impl<T> Send for JoinHandle<T> {}
|
||||
unsafe impl<T> Sync for JoinHandle<T> {}
|
||||
|
||||
impl<T> JoinHandle<T> {
|
||||
/// Creates a new `JoinHandle`.
|
||||
pub(crate) fn new(inner: async_task::JoinHandle<T, Task>) -> JoinHandle<T> {
|
||||
JoinHandle(inner)
|
||||
}
|
||||
|
||||
/// Returns a handle to the underlying task.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # async_std::task::block_on(async {
|
||||
/// #
|
||||
/// use async_std::task;
|
||||
///
|
||||
/// let handle = task::spawn(async {
|
||||
/// 1 + 2
|
||||
/// });
|
||||
/// println!("id = {}", handle.task().id());
|
||||
/// #
|
||||
/// # })
|
||||
pub fn task(&self) -> &Task {
|
||||
self.0.tag()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Future for JoinHandle<T> {
|
||||
type Output = T;
|
||||
|
||||
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
match Pin::new(&mut self.0).poll(cx) {
|
||||
Poll::Pending => Poll::Pending,
|
||||
Poll::Ready(None) => panic!("cannot await the result of a panicked task"),
|
||||
Poll::Ready(Some(val)) => Poll::Ready(val),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,136 +0,0 @@
|
||||
use std::iter;
|
||||
use std::thread;
|
||||
|
||||
use crossbeam_deque::{Injector, Stealer, Worker};
|
||||
use kv_log_macro::trace;
|
||||
use once_cell::sync::Lazy;
|
||||
|
||||
use super::sleepers::Sleepers;
|
||||
use super::task;
|
||||
use super::task_local;
|
||||
use super::worker;
|
||||
use super::{Builder, JoinHandle};
|
||||
use crate::future::Future;
|
||||
use crate::utils::abort_on_panic;
|
||||
|
||||
/// Spawns a task.
|
||||
///
|
||||
/// This function is similar to [`std::thread::spawn`], except it spawns an asynchronous task.
|
||||
///
|
||||
/// [`std::thread`]: https://doc.rust-lang.org/std/thread/fn.spawn.html
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # async_std::task::block_on(async {
|
||||
/// #
|
||||
/// use async_std::task;
|
||||
///
|
||||
/// let handle = task::spawn(async {
|
||||
/// 1 + 2
|
||||
/// });
|
||||
///
|
||||
/// assert_eq!(handle.await, 3);
|
||||
/// #
|
||||
/// # })
|
||||
/// ```
|
||||
pub fn spawn<F, T>(future: F) -> JoinHandle<T>
|
||||
where
|
||||
F: Future<Output = T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
Builder::new().spawn(future).expect("cannot spawn future")
|
||||
}
|
||||
|
||||
pub(crate) struct Pool {
|
||||
pub injector: Injector<task::Runnable>,
|
||||
pub stealers: Vec<Stealer<task::Runnable>>,
|
||||
pub sleepers: Sleepers,
|
||||
}
|
||||
|
||||
impl Pool {
|
||||
/// Spawn a future onto the pool.
|
||||
pub fn spawn<F, T>(&self, future: F, builder: Builder) -> JoinHandle<T>
|
||||
where
|
||||
F: Future<Output = T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
let tag = task::Tag::new(builder.name);
|
||||
|
||||
// Log this `spawn` operation.
|
||||
let child_id = tag.task_id().as_u64();
|
||||
let parent_id = worker::get_task(|t| t.id().as_u64()).unwrap_or(0);
|
||||
|
||||
trace!("spawn", {
|
||||
parent_id: parent_id,
|
||||
child_id: child_id,
|
||||
});
|
||||
|
||||
// Wrap the future into one that drops task-local variables on exit.
|
||||
let future = unsafe { task_local::add_finalizer(future) };
|
||||
|
||||
// Wrap the future into one that logs completion on exit.
|
||||
let future = async move {
|
||||
let res = future.await;
|
||||
trace!("spawn completed", {
|
||||
parent_id: parent_id,
|
||||
child_id: child_id,
|
||||
});
|
||||
res
|
||||
};
|
||||
|
||||
let (task, handle) = async_task::spawn(future, worker::schedule, tag);
|
||||
task.schedule();
|
||||
JoinHandle::new(handle)
|
||||
}
|
||||
|
||||
/// Find the next runnable task to run.
|
||||
pub fn find_task(&self, local: &Worker<task::Runnable>) -> Option<task::Runnable> {
|
||||
// Pop a task from the local queue, if not empty.
|
||||
local.pop().or_else(|| {
|
||||
// Otherwise, we need to look for a task elsewhere.
|
||||
iter::repeat_with(|| {
|
||||
// Try stealing a batch of tasks from the injector queue.
|
||||
self.injector
|
||||
.steal_batch_and_pop(local)
|
||||
// Or try stealing a bach of tasks from one of the other threads.
|
||||
.or_else(|| {
|
||||
self.stealers
|
||||
.iter()
|
||||
.map(|s| s.steal_batch_and_pop(local))
|
||||
.collect()
|
||||
})
|
||||
})
|
||||
// Loop while no task was stolen and any steal operation needs to be retried.
|
||||
.find(|s| !s.is_retry())
|
||||
// Extract the stolen task, if there is one.
|
||||
.and_then(|s| s.success())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn get() -> &'static Pool {
|
||||
static POOL: Lazy<Pool> = Lazy::new(|| {
|
||||
let num_threads = num_cpus::get().max(1);
|
||||
let mut stealers = Vec::new();
|
||||
|
||||
// Spawn worker threads.
|
||||
for _ in 0..num_threads {
|
||||
let worker = Worker::new_fifo();
|
||||
stealers.push(worker.stealer());
|
||||
|
||||
thread::Builder::new()
|
||||
.name("async-task-driver".to_string())
|
||||
.spawn(|| abort_on_panic(|| worker::main_loop(worker)))
|
||||
.expect("cannot start a thread driving tasks");
|
||||
}
|
||||
|
||||
Pool {
|
||||
injector: Injector::new(),
|
||||
stealers,
|
||||
sleepers: Sleepers::new(),
|
||||
}
|
||||
});
|
||||
&*POOL
|
||||
}
|
||||
@ -0,0 +1,31 @@
|
||||
use crate::future::Future;
|
||||
use crate::task::{Builder, JoinHandle};
|
||||
|
||||
/// Spawns a task.
|
||||
///
|
||||
/// This function is similar to [`std::thread::spawn`], except it spawns an asynchronous task.
|
||||
///
|
||||
/// [`std::thread`]: https://doc.rust-lang.org/std/thread/fn.spawn.html
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// # async_std::task::block_on(async {
|
||||
/// #
|
||||
/// use async_std::task;
|
||||
///
|
||||
/// let handle = task::spawn(async {
|
||||
/// 1 + 2
|
||||
/// });
|
||||
///
|
||||
/// assert_eq!(handle.await, 3);
|
||||
/// #
|
||||
/// # })
|
||||
/// ```
|
||||
pub fn spawn<F, T>(future: F) -> JoinHandle<T>
|
||||
where
|
||||
F: Future<Output = T> + Send + 'static,
|
||||
T: Send + 'static,
|
||||
{
|
||||
Builder::new().spawn(future).expect("cannot spawn task")
|
||||
}
|
||||
@ -0,0 +1,35 @@
|
||||
use std::fmt;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
/// A unique identifier for a task.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use async_std::task;
|
||||
///
|
||||
/// task::block_on(async {
|
||||
/// println!("id = {:?}", task::current().id());
|
||||
/// })
|
||||
/// ```
|
||||
#[derive(Eq, PartialEq, Clone, Copy, Hash, Debug)]
|
||||
pub struct TaskId(pub(crate) u64);
|
||||
|
||||
impl TaskId {
|
||||
/// Generates a new `TaskId`.
|
||||
pub(crate) fn generate() -> TaskId {
|
||||
static COUNTER: AtomicU64 = AtomicU64::new(1);
|
||||
|
||||
let id = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||
if id > u64::max_value() / 2 {
|
||||
std::process::abort();
|
||||
}
|
||||
TaskId(id)
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for TaskId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.0)
|
||||
}
|
||||
}
|
||||
@ -1,110 +0,0 @@
|
||||
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
|
||||
///
|
||||
/// ```
|
||||
/// # 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(crate) 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