tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-09-02 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

executor: move into tokio crate (#1702)

Browse Source 
A step towards collapsing Tokio sub crates into a single `tokio`
crate (#1318).

The executor implementation is now provided by the main `tokio` crate.
Functionality can be opted out of by using the various net related
feature flags.
This commit is contained in:
Carl Lerche
2019-10-28 21:40:29 -07:00
committed by GitHub
parent 7eb264a0d0
commit c62ef2d232
101 changed files with 387 additions and 690 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/Cargo.toml
+9 -12
View File
@@ -25,6 +25,7 @@ keywords = ["io", "async", "non-blocking", "futures"]
[features]
default = [
"blocking",
"fs",
"io",
"net-full",
@@ -35,25 +36,25 @@ default = [
"timer",
]
fs = ["tokio-executor/blocking"]
blocking = []
fs = ["blocking"]
io-traits = ["bytes", "iovec"]
io-util = ["io-traits", "pin-project", "memchr"]
io = ["io-traits", "io-util"]
macros = ["tokio-macros"]
net-full = ["tcp", "udp", "uds"]
net-driver = ["mio", "tokio-executor/blocking", "lazy_static"]
net-driver = ["mio", "blocking", "lazy_static"]
rt-current-thread = [
"crossbeam-channel",
"timer",
"tokio-executor/current-thread",
]
rt-full = [
"macros",
"num_cpus",
"net-full",
"rt-current-thread",
"sync",
"timer",
"tokio-executor/current-thread",
"tokio-executor/thread-pool",
]
signal = [
"lazy_static",
@@ -82,10 +83,11 @@ process = [
[dependencies]
futures-core-preview = "=0.3.0-alpha.19"
futures-sink-preview = "=0.3.0-alpha.19"
futures-util-preview = { version = "=0.3.0-alpha.19", features = ["sink"] }
futures-util-preview = { version = "=0.3.0-alpha.19", features = ["sink", "channel"] }
# Everything else is optional...
bytes = { version = "0.4", optional = true }
crossbeam-channel = { version = "0.3.8", optional = true }
crossbeam-utils = { version = "0.6.0", optional = true }
iovec = { version = "0.1", optional = true }
lazy_static = { version = "1.0.2", optional = true }
@@ -95,7 +97,6 @@ num_cpus = { version = "1.8.0", optional = true }
pin-project = { version = "0.4", optional = true }
# Backs `DelayQueue`
slab = { version = "0.4.1", optional = true }
tokio-executor = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-executor" }
tokio-macros = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-macros" }
tokio-sync = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-sync", features = ["async-traits"] }
@@ -113,10 +114,6 @@ version = "0.3.8"
default-features = false
optional = true
[target.'cfg(loom)'.dependencies]
# play nice with loom tests in other crates.
loom = "0.2.11"
[dev-dependencies]
tokio-test = { version = "=0.2.0-alpha.6", path = "../tokio-test" }
tokio-util = { version = "=0.2.0-alpha.6", path = "../tokio-util" }
@@ -127,6 +124,7 @@ flate2 = { version = "1", features = ["tokio"] }
http = "0.1"
httparse = "1.0"
libc = "0.2"
loom = { version = "0.2.11", features = ["futures", "checkpoint"] }
num_cpus = "1.0"
rand = "0.7.2"
serde = { version = "1.0", features = ["derive"] }
@@ -135,7 +133,6 @@ tempfile = "3.1.0"
time = "0.1"
# sharded slab tests
loom = "0.2.11"
proptest = "0.9.4"
[package.metadata.docs.rs]
tokio/benches/thread_pool.rs
+161
View File
@@ -0,0 +1,161 @@
#![feature(test)]
extern crate test;
use tokio::executor::thread_pool::{Builder, Spawner};
use tokio_sync::oneshot;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::{mpsc, Arc};
use std::task::{Context, Poll};
struct Backoff(usize);
impl Future for Backoff {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.0 == 0 {
Poll::Ready(())
} else {
self.0 -= 1;
cx.waker().wake_by_ref();
Poll::Pending
}
}
}
const NUM_THREADS: usize = 6;
#[bench]
fn spawn_many(b: &mut test::Bencher) {
const NUM_SPAWN: usize = 10_000;
let threadpool = Builder::new().num_threads(NUM_THREADS).build();
let (tx, rx) = mpsc::sync_channel(1000);
let rem = Arc::new(AtomicUsize::new(0));
b.iter(|| {
rem.store(NUM_SPAWN, Relaxed);
for _ in 0..NUM_SPAWN {
let tx = tx.clone();
let rem = rem.clone();
threadpool.spawn(async move {
if 1 == rem.fetch_sub(1, Relaxed) {
tx.send(()).unwrap();
}
});
}
let _ = rx.recv().unwrap();
});
}
#[bench]
fn yield_many(b: &mut test::Bencher) {
const NUM_YIELD: usize = 1_000;
const TASKS_PER_CPU: usize = 50;
let threadpool = Builder::new().num_threads(NUM_THREADS).build();
let tasks = TASKS_PER_CPU * num_cpus::get_physical();
let (tx, rx) = mpsc::sync_channel(tasks);
b.iter(move || {
for _ in 0..tasks {
let tx = tx.clone();
threadpool.spawn(async move {
let backoff = Backoff(NUM_YIELD);
backoff.await;
tx.send(()).unwrap();
});
}
for _ in 0..tasks {
let _ = rx.recv().unwrap();
}
});
}
#[bench]
fn ping_pong(b: &mut test::Bencher) {
const NUM_PINGS: usize = 1_000;
let threadpool = Builder::new().num_threads(NUM_THREADS).build();
let (done_tx, done_rx) = mpsc::sync_channel(1000);
let rem = Arc::new(AtomicUsize::new(0));
b.iter(|| {
let done_tx = done_tx.clone();
let rem = rem.clone();
rem.store(NUM_PINGS, Relaxed);
let spawner = threadpool.spawner().clone();
threadpool.spawn(async move {
for _ in 0..NUM_PINGS {
let rem = rem.clone();
let done_tx = done_tx.clone();
let spawner2 = spawner.clone();
spawner.spawn(async move {
let (tx1, rx1) = oneshot::channel();
let (tx2, rx2) = oneshot::channel();
spawner2.spawn(async move {
rx1.await.unwrap();
tx2.send(()).unwrap();
});
tx1.send(()).unwrap();
rx2.await.unwrap();
if 1 == rem.fetch_sub(1, Relaxed) {
done_tx.send(()).unwrap();
}
});
}
});
done_rx.recv().unwrap();
});
}
#[bench]
fn chained_spawn(b: &mut test::Bencher) {
const ITER: usize = 1_000;
let threadpool = Builder::new().num_threads(NUM_THREADS).build();
fn iter(spawner: Spawner, done_tx: mpsc::SyncSender<()>, n: usize) {
if n == 0 {
done_tx.send(()).unwrap();
} else {
let s2 = spawner.clone();
spawner.spawn(async move {
iter(s2, done_tx, n - 1);
});
}
}
let (done_tx, done_rx) = mpsc::sync_channel(1000);
b.iter(move || {
let done_tx = done_tx.clone();
let spawner = threadpool.spawner().clone();
threadpool.spawn(async move {
iter(spawner, done_tx, ITER);
});
done_rx.recv().unwrap();
});
}
tokio/src/executor.rs
-104
View File
@@ -1,104 +0,0 @@
//! Task execution utilities.
//!
//! In the Tokio execution model, futures are lazy. When a future is created, no
//! work is performed. In order for the work defined by the future to happen,
//! the future must be submitted to an executor. A future that is submitted to
//! an executor is called a "task".
//!
//! The executor is responsible for ensuring that [`Future::poll`] is
//! called whenever the task is [notified]. Notification happens when the
//! internal state of a task transitions from "not ready" to ready. For
//! example, a socket might have received data and a call to `read` will now be
//! able to succeed.
//!
//! The specific strategy used to manage the tasks is left up to the
//! executor. There are two main flavors of executors: single-threaded and
//! multi-threaded. Tokio provides implementation for both of these in the
//! [`runtime`] module.
//!
//! # `Executor` trait.
//!
//! This module provides the [`Executor`] trait (re-exported from
//! [`tokio-executor`]), which describes the API that all executors must
//! implement.
//!
//! A free [`spawn`] function is provided that allows spawning futures onto the
//! default executor (tracked via a thread-local variable) without referencing a
//! handle. It is expected that all executors will set a value for the default
//! executor. This value will often be set to the executor itself, but it is
//! possible that the default executor might be set to a different executor.
//!
//! For example, a single threaded executor might set the default executor to a
//! thread pool instead of itself, allowing futures to spawn new tasks onto the
//! thread pool when those tasks are `Send`.
//!
//! [`Future::poll`]: https://docs.rs/futures/0.1/futures/future/trait.Future.html#tymethod.poll
//! [notified]: https://docs.rs/futures/0.1/futures/executor/trait.Notify.html#tymethod.notify
//! [`runtime`]: ../runtime/index.html
//! [`tokio-executor`]: https://docs.rs/tokio-executor/0.1
//! [`Executor`]: trait.Executor.html
//! [`spawn`]: fn.spawn.html
use std::future::Future;
pub use tokio_executor::{DefaultExecutor, Executor, SpawnError, TypedExecutor};
/// Return value from the `spawn` function.
///
/// Currently this value doesn't actually provide any functionality. However, it
/// provides a way to add functionality later without breaking backwards
/// compatibility.
///
/// See [`spawn`] for more details.
///
/// [`spawn`]: fn.spawn.html
#[derive(Debug)]
pub struct Spawn(());
/// Spawns a future on the default executor.
///
/// In order for a future to do work, it must be spawned on an executor. The
/// `spawn` function is the easiest way to do this. It spawns a future on the
/// [default executor] for the current execution context (tracked using a
/// thread-local variable).
///
/// The default executor is **usually** a thread pool.
///
/// # Examples
///
/// In this example, a server is started and `spawn` is used to start a new task
/// that processes each received connection.
///
/// ```
/// use tokio::net::TcpListener;
///
/// # async fn process<T>(_t: T) {}
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let mut listener = TcpListener::bind("127.0.0.1:8080").await?;
///
/// loop {
/// let (socket, _) = listener.accept().await?;
///
/// tokio::spawn(async move {
/// // Process each socket concurrently.
/// process(socket).await
/// });
/// }
/// # }
/// ```
///
/// [default executor]: struct.DefaultExecutor.html
///
/// # Panics
///
/// This function will panic if the default executor is not set or if spawning
/// onto the default executor returns an error. To avoid the panic, use
/// [`DefaultExecutor`].
///
/// [`DefaultExecutor`]: struct.DefaultExecutor.html
pub fn spawn<F>(f: F) -> Spawn
where
F: Future<Output = ()> + 'static + Send,
{
::tokio_executor::spawn(f);
Spawn(())
}
tokio/src/executor/blocking/builder.rs
+58
View File
@@ -0,0 +1,58 @@
use crate::executor::blocking::Pool;
use crate::executor::loom::thread;
use std::usize;
/// Builds a blocking thread pool with custom configuration values.
pub(crate) struct Builder {
/// Thread name
name: String,
/// Thread stack size
stack_size: Option<usize>,
}
impl Default for Builder {
fn default() -> Self {
Builder {
name: "tokio-blocking-thread".to_string(),
stack_size: None,
}
}
}
impl Builder {
/// Set name of threads spawned by the pool
///
/// If this configuration is not set, then the thread will use the system
/// default naming scheme.
pub(crate) fn name<S: Into<String>>(&mut self, val: S) -> &mut Self {
self.name = val.into();
self
}
/// Set the stack size (in bytes) for worker threads.
///
/// The actual stack size may be greater than this value if the platform
/// specifies minimal stack size.
///
/// The default stack size for spawned threads is 2 MiB, though this
/// particular stack size is subject to change in the future.
pub(crate) fn stack_size(&mut self, val: usize) -> &mut Self {
self.stack_size = Some(val);
self
}
pub(crate) fn build(self) -> Pool {
let mut p = Pool::default();
let Builder { stack_size, name } = self;
p.new_thread = Box::new(move || {
let mut b = thread::Builder::new().name(name.clone());
if let Some(stack_size) = stack_size {
b = b.stack_size(stack_size);
}
b
});
p
}
}
tokio/src/executor/blocking/mod.rs
+332
View File
@@ -0,0 +1,332 @@
//! Thread pool for blocking operations
use crate::executor::loom::sync::{Arc, Condvar, Mutex};
use crate::executor::loom::thread;
#[cfg(feature = "blocking")]
use tokio_sync::oneshot;
use std::cell::Cell;
use std::collections::VecDeque;
use std::fmt;
#[cfg(feature = "blocking")]
use std::future::Future;
use std::ops::Deref;
#[cfg(feature = "blocking")]
use std::pin::Pin;
#[cfg(feature = "blocking")]
use std::task::{Context, Poll};
use std::time::Duration;
#[cfg(feature = "rt-full")]
mod builder;
#[cfg(feature = "rt-full")]
pub(crate) use builder::Builder;
#[derive(Clone, Copy)]
enum State {
Empty,
Ready(*const Arc<Pool>),
}
thread_local! {
/// Thread-local tracking the current executor
static BLOCKING: Cell<State> = Cell::new(State::Empty)
}
/// Set the blocking pool for the duration of the closure
///
/// If a blocking pool is already set, it will be restored when the closure returns or if it
/// panics.
#[allow(dead_code)] // we allow dead code since this won't be called if no executors are enabled
pub(crate) fn with_pool<F, R>(pool: &Arc<Pool>, f: F) -> R
where
F: FnOnce() -> R,
{
// While scary, this is safe. The function takes a `&Pool`, which guarantees
// that the reference lives for the duration of `with_pool`.
//
// Because we are always clearing the TLS value at the end of the
// function, we can cast the reference to 'static which thread-local
// cells require.
BLOCKING.with(|cell| {
let was = cell.replace(State::Empty);
// Ensure that the pool is removed from the thread-local context
// when leaving the scope. This handles cases that involve panicking.
struct Reset<'a>(&'a Cell<State>, State);
impl Drop for Reset<'_> {
fn drop(&mut self) {
self.0.set(self.1);
}
}
let _reset = Reset(cell, was);
cell.set(State::Ready(pool as *const _));
f()
})
}
pub(crate) struct Pool {
shared: Mutex<Shared>,
condvar: Condvar,
new_thread: Box<dyn Fn() -> thread::Builder + Send + Sync + 'static>,
}
impl fmt::Debug for Pool {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Pool").finish()
}
}
struct Shared {
queue: VecDeque<Box<dyn FnOnce() + Send>>,
num_th: u32,
num_idle: u32,
num_notify: u32,
shutdown: bool,
}
const MAX_THREADS: u32 = 1_000;
const KEEP_ALIVE: Duration = Duration::from_secs(10);
/// Result of a blocking operation running on the blocking thread pool.
#[cfg(feature = "blocking")]
#[derive(Debug)]
pub struct Blocking<T> {
rx: oneshot::Receiver<T>,
}
impl Pool {
/// Run the provided function on an executor dedicated to blocking operations.
pub(crate) fn spawn(this: &Arc<Self>, f: Box<dyn FnOnce() + Send + 'static>) {
let should_spawn = {
let mut shared = this.shared.lock().unwrap();
if shared.shutdown {
// no need to even push this task; it would never get picked up
return;
}
shared.queue.push_back(f);
if shared.num_idle == 0 {
// No threads are able to process the task.
if shared.num_th == MAX_THREADS {
// At max number of threads
false
} else {
shared.num_th += 1;
true
}
} else {
// Notify an idle worker thread. The notification counter
// is used to count the needed amount of notifications
// exactly. Thread libraries may generate spurious
// wakeups, this counter is used to keep us in a
// consistent state.
shared.num_idle -= 1;
shared.num_notify += 1;
this.condvar.notify_one();
false
}
};
if should_spawn {
Pool::spawn_thread(Arc::clone(this), (this.new_thread)());
}
}
// NOTE: we cannot use self here w/o arbitrary_self_types since Arc is loom::Arc
fn spawn_thread(this: Arc<Self>, builder: thread::Builder) {
builder
.spawn(move || {
let mut shared = this.shared.lock().unwrap();
'main: loop {
// BUSY
while let Some(task) = shared.queue.pop_front() {
drop(shared);
run_task(task);
shared = this.shared.lock().unwrap();
if shared.shutdown {
break; // Need to increment idle before we exit
}
}
// IDLE
shared.num_idle += 1;
while !shared.shutdown {
let lock_result = this.condvar.wait_timeout(shared, KEEP_ALIVE).unwrap();
shared = lock_result.0;
let timeout_result = lock_result.1;
if shared.num_notify != 0 {
// We have received a legitimate wakeup,
// acknowledge it by decrementing the counter
// and transition to the BUSY state.
shared.num_notify -= 1;
break;
}
if timeout_result.timed_out() {
break 'main;
}
// Spurious wakeup detected, go back to sleep.
}
if shared.shutdown {
// Work was produced, and we "took" it (by decrementing num_notify).
// This means that num_idle was decremented once for our wakeup.
// But, since we are exiting, we need to "undo" that, as we'll stay idle.
shared.num_idle += 1;
// NOTE: Technically we should also do num_notify++ and notify again,
// but since we're shutting down anyway, that won't be necessary.
break;
}
}
// Thread exit
shared.num_th -= 1;
// num_idle should now be tracked exactly, panic
// with a descriptive message if it is not the
// case.
shared.num_idle = shared
.num_idle
.checked_sub(1)
.expect("num_idle underflowed on thread exit");
if shared.shutdown && shared.num_th == 0 {
this.condvar.notify_one();
}
})
.unwrap();
}
/// Shut down all workers in the pool the next time they are idle.
///
/// Blocks until all threads have exited.
pub(crate) fn shutdown(&self) {
let mut shared = self.shared.lock().unwrap();
shared.shutdown = true;
self.condvar.notify_all();
while shared.num_th > 0 {
shared = self.condvar.wait(shared).unwrap();
}
}
}
pub(crate) struct PoolWaiter(Arc<Pool>);
impl From<Pool> for PoolWaiter {
fn from(p: Pool) -> Self {
Self::from(Arc::new(p))
}
}
impl From<Arc<Pool>> for PoolWaiter {
fn from(p: Arc<Pool>) -> Self {
Self(p)
}
}
impl Deref for PoolWaiter {
type Target = Arc<Pool>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl Drop for PoolWaiter {
fn drop(&mut self) {
self.0.shutdown();
}
}
/// Run the provided closure on a thread where blocking is acceptable.
///
/// In general, issuing a blocking call or performing a lot of compute in a future without
/// yielding is not okay, as it may prevent the executor from driving other futures forward.
/// A closure that is run through this method will instead be run on a dedicated thread pool for
/// such blocking tasks without holding up the main futures executor.
///
/// # Examples
///
/// ```
/// # async fn docs() {
/// tokio::executor::blocking::run(move || {
/// // do some compute-heavy work or call synchronous code
/// }).await;
/// # }
/// ```
#[cfg(feature = "blocking")]
pub fn run<F, R>(f: F) -> Blocking<R>
where
F: FnOnce() -> R + Send + 'static,
R: Send + 'static,
{
let (tx, rx) = oneshot::channel();
BLOCKING.with(|current_pool| match current_pool.get() {
State::Ready(pool) => {
let pool = unsafe { &*pool };
Pool::spawn(
pool,
Box::new(move || {
// receiver may have gone away
let _ = tx.send(f());
}),
);
}
State::Empty => panic!("must be called from the context of Tokio runtime"),
});
Blocking { rx }
}
#[cfg(feature = "blocking")]
impl<T> Future for Blocking<T> {
type Output = T;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
use std::task::Poll::*;
match Pin::new(&mut self.rx).poll(cx) {
Ready(Ok(v)) => Ready(v),
Ready(Err(_)) => panic!(
"the blocking operation has been dropped before completing. \
This should not happen and is a bug."
),
Pending => Pending,
}
}
}
fn run_task(f: Box<dyn FnOnce() + Send>) {
use std::panic::{catch_unwind, AssertUnwindSafe};
let _ = catch_unwind(AssertUnwindSafe(|| f()));
}
impl Default for Pool {
fn default() -> Self {
Pool {
shared: Mutex::new(Shared {
queue: VecDeque::new(),
num_th: 0,
num_idle: 0,
num_notify: 0,
shutdown: false,
}),
condvar: Condvar::new(),
new_thread: Box::new(|| {
thread::Builder::new().name("tokio-blocking-driver".to_string())
}),
}
}
}
tokio/src/executor/current_thread/mod.rs
+873
View File
@@ -0,0 +1,873 @@
//! A single-threaded executor which executes tasks on the same thread from which
//! they are spawned.
//!
//! [`CurrentThread`] is the main type of this crate. It executes tasks on the
//! current thread. The easiest way to start a new [`CurrentThread`] executor
//! is to call [`block_on_all`] with an initial task to seed the executor. All
//! tasks that are being managed by a [`CurrentThread`] executor are able to
//! spawn additional tasks by calling [`spawn`].
//!
//! Application authors will not use this crate directly. Instead, they will use
//! the `tokio` crate. Library authors should only depend on
//! `tokio-current-thread` if they are building a custom task executor.
//!
//! [`CurrentThread`]: struct.CurrentThread.html
//! [`spawn`]: fn.spawn.html
//! [`block_on_all`]: fn.block_on_all.html
mod scheduler;
use self::scheduler::{Scheduler, TickArgs};
use crate::executor::{EnterError, Executor, SpawnError, TypedExecutor};
#[cfg(feature = "blocking")]
use crate::executor::blocking::{Pool, PoolWaiter};
use crate::executor::park::{Park, ParkThread, Unpark};
use std::cell::Cell;
use std::error::Error;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::rc::Rc;
use std::sync::{atomic, Arc};
use std::task::{Context, Poll, Waker};
use std::thread;
use std::time::{Duration, Instant};
/// Executes tasks on the current thread
pub struct CurrentThread<P: Park = ParkThread> {
/// Execute futures and receive unpark notifications.
scheduler: Scheduler<P::Unpark>,
/// Current number of futures being executed.
///
/// The LSB is used to indicate that the runtime is preparing to shut down.
/// Thus, to get the actual number of pending futures, `>>1`.
num_futures: Arc<atomic::AtomicUsize>,
/// Thread park handle
park: P,
/// Handle for spawning new futures from other threads
spawn_handle: Handle,
/// Receiver for futures spawned from other threads
spawn_receiver: crossbeam_channel::Receiver<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
/// Handle to pool for handling blocking tasks
#[cfg(feature = "blocking")]
blocking: PoolWaiter,
/// The thread-local ID assigned to this executor.
id: u64,
}
/// Executes futures on the current thread.
///
/// All futures executed using this executor will be executed on the current
/// thread. As such, `run` will wait for these futures to complete before
/// returning.
///
/// For more details, see the [module level](index.html) documentation.
#[derive(Debug, Clone)]
pub struct TaskExecutor {
// Prevent the handle from moving across threads.
_p: ::std::marker::PhantomData<Rc<()>>,
}
/// Returned by the `turn` function.
#[derive(Debug)]
pub struct Turn {
polled: bool,
}
impl Turn {
/// `true` if any futures were polled at all and `false` otherwise.
pub fn has_polled(&self) -> bool {
self.polled
}
}
/// A `CurrentThread` instance bound to a supplied execution context.
pub struct Entered<'a, P: Park> {
executor: &'a mut CurrentThread<P>,
}
/// Error returned by the `run` function.
#[derive(Debug)]
pub struct RunError {
_p: (),
}
impl fmt::Display for RunError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "Run error")
}
}
impl Error for RunError {}
/// Error returned by the `run_timeout` function.
#[derive(Debug)]
pub struct RunTimeoutError {
timeout: bool,
}
impl fmt::Display for RunTimeoutError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
let descr = if self.timeout {
"Run timeout error (timeout)"
} else {
"Run timeout error (not timeout)"
};
write!(fmt, "{}", descr)
}
}
impl Error for RunTimeoutError {}
/// Error returned by the `turn` function.
#[derive(Debug)]
pub struct TurnError {
_p: (),
}
impl fmt::Display for TurnError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "Turn error")
}
}
impl Error for TurnError {}
/// Error returned by the `block_on` function.
#[derive(Debug)]
pub struct BlockError<T> {
inner: Option<T>,
}
impl<T> fmt::Display for BlockError<T> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "Block error")
}
}
impl<T: fmt::Debug> Error for BlockError<T> {}
/// This is mostly split out to make the borrow checker happy.
struct Borrow<'a, U> {
spawner: BorrowSpawner<'a, U>,
#[cfg(feature = "blocking")]
blocking: &'a PoolWaiter,
}
/// As is this.
struct BorrowSpawner<'a, U> {
id: u64,
num_futures: &'a atomic::AtomicUsize,
scheduler: &'a mut Scheduler<U>,
}
trait SpawnLocal {
fn spawn_local(&mut self, future: Pin<Box<dyn Future<Output = ()>>>, already_counted: bool);
}
struct CurrentRunner {
spawn: Cell<Option<*mut dyn SpawnLocal>>,
id: Cell<Option<u64>>,
}
thread_local! {
/// Current thread's task runner. This is set in `TaskRunner::with`
static CURRENT: CurrentRunner = CurrentRunner {
spawn: Cell::new(None),
id: Cell::new(None),
}
}
thread_local! {
/// Unique ID to assign to each new executor launched on this thread.
///
/// The unique ID is used to determine if the currently running executor matches the one
/// referred to by a `Handle` so that direct task dispatch can be used.
static EXECUTOR_ID: Cell<u64> = Cell::new(0)
}
/// Run the executor bootstrapping the execution with the provided future.
///
/// This creates a new [`CurrentThread`] executor, spawns the provided future,
/// and blocks the current thread until the provided future and **all**
/// subsequently spawned futures complete. In other words:
///
/// * If the provided bootstrap future does **not** spawn any additional tasks,
/// `block_on_all` returns once `future` completes.
/// * If the provided bootstrap future **does** spawn additional tasks, then
/// `block_on_all` returns once **all** spawned futures complete.
///
/// See [module level][mod] documentation for more details.
///
/// [`CurrentThread`]: struct.CurrentThread.html
/// [mod]: index.html
pub fn block_on_all<F>(future: F) -> F::Output
where
F: Future,
{
let mut current_thread = CurrentThread::new();
let ret = current_thread.block_on(future);
current_thread.run().unwrap();
ret
}
/// Executes a future on the current thread.
///
/// The provided future must complete or be canceled before `run` will return.
///
/// Unlike [`tokio::spawn`], this function will always spawn on a
/// `CurrentThread` executor and is able to spawn futures that are not `Send`.
///
/// # Panics
///
/// This function can only be invoked from the context of a `run` call; any
/// other use will result in a panic.
///
/// [`tokio::spawn`]: ../fn.spawn.html
pub fn spawn<F>(future: F)
where
F: Future<Output = ()> + 'static,
{
TaskExecutor::current()
.spawn_local(Box::pin(future))
.unwrap();
}
// ===== impl CurrentThread =====
impl CurrentThread<ParkThread> {
/// Create a new instance of `CurrentThread`.
pub fn new() -> Self {
CurrentThread::new_with_park(ParkThread::new())
}
}
impl<P: Park> CurrentThread<P> {
/// Create a new instance of `CurrentThread` backed by the given park
/// handle.
pub fn new_with_park(park: P) -> Self {
let unpark = park.unpark();
let (spawn_sender, spawn_receiver) = crossbeam_channel::unbounded();
let thread = thread::current().id();
let id = EXECUTOR_ID.with(|idc| {
let id = idc.get();
idc.set(id + 1);
id
});
let scheduler = Scheduler::new(unpark);
let waker = scheduler.waker();
let num_futures = Arc::new(atomic::AtomicUsize::new(0));
CurrentThread {
scheduler,
num_futures: num_futures.clone(),
park,
id,
spawn_handle: Handle {
sender: spawn_sender,
num_futures,
waker,
thread,
id,
},
spawn_receiver,
#[cfg(feature = "blocking")]
blocking: PoolWaiter::from(Pool::default()),
}
}
/// Returns `true` if the executor is currently idle.
///
/// An idle executor is defined by not currently having any spawned tasks.
///
/// Note that this method is inherently racy -- if a future is spawned from a remote `Handle`,
/// this method may return `true` even though there are more futures to be executed.
pub fn is_idle(&self) -> bool {
self.num_futures.load(atomic::Ordering::SeqCst) <= 1
}
/// Spawn the future on the executor.
///
/// This internally queues the future to be executed once `run` is called.
pub fn spawn<F>(&mut self, future: F) -> &mut Self
where
F: Future<Output = ()> + 'static,
{
self.borrow().spawner.spawn_local(Box::pin(future), false);
self
}
/// Synchronously waits for the provided `future` to complete.
///
/// This function can be used to synchronously block the current thread
/// until the provided `future` has resolved either successfully or with an
/// error. The result of the future is then returned from this function
/// call.
///
/// Note that this function will **also** execute any spawned futures on the
/// current thread, but will **not** block until these other spawned futures
/// have completed.
///
/// The caller is responsible for ensuring that other spawned futures
/// complete execution.
pub fn block_on<F>(&mut self, future: F) -> F::Output
where
F: Future,
{
let _enter = crate::executor::enter().expect("failed to start `current_thread::Runtime`");
self.enter().block_on(future)
}
/// Run the executor to completion, blocking the thread until **all**
/// spawned futures have completed.
pub fn run(&mut self) -> Result<(), RunError> {
let _enter = crate::executor::enter().expect("failed to start `current_thread::Runtime`");
self.enter().run()
}
/// Run the executor to completion, blocking the thread until all
/// spawned futures have completed **or** `duration` time has elapsed.
pub fn run_timeout(&mut self, duration: Duration) -> Result<(), RunTimeoutError> {
let _enter = crate::executor::enter().expect("failed to start `current_thread::Runtime`");
self.enter().run_timeout(duration)
}
/// Perform a single iteration of the event loop.
///
/// This function blocks the current thread even if the executor is idle.
pub fn turn(&mut self, duration: Option<Duration>) -> Result<Turn, TurnError> {
let _enter = crate::executor::enter().expect("failed to start `current_thread::Runtime`");
self.enter().turn(duration)
}
/// Bind `CurrentThread` instance with an execution context.
fn enter(&mut self) -> Entered<'_, P> {
Entered { executor: self }
}
/// Returns a reference to the underlying `Park` instance.
pub fn get_park(&self) -> &P {
&self.park
}
/// Returns a mutable reference to the underlying `Park` instance.
pub fn get_park_mut(&mut self) -> &mut P {
&mut self.park
}
fn borrow(&mut self) -> Borrow<'_, P::Unpark> {
Borrow {
spawner: BorrowSpawner {
id: self.id,
scheduler: &mut self.scheduler,
num_futures: &*self.num_futures,
},
#[cfg(feature = "blocking")]
blocking: &self.blocking,
}
}
/// Get a new handle to spawn futures on the executor
///
/// Different to the executor itself, the handle can be sent to different
/// threads and can be used to spawn futures on the executor.
pub fn handle(&self) -> Handle {
self.spawn_handle.clone()
}
}
impl<P: Park> Drop for CurrentThread<P> {
fn drop(&mut self) {
// Signal to Handles that no more futures can be spawned by setting LSB.
//
// NOTE: this isn't technically necessary since the send on the mpsc will fail once the
// receiver is dropped, but it's useful to illustrate how clean shutdown will be
// implemented (e.g., by setting the LSB).
let pending = self.num_futures.fetch_add(1, atomic::Ordering::SeqCst);
// TODO: We currently ignore any pending futures at the time we shut down.
//
// The "proper" fix for this is to have an explicit shutdown phase (`shutdown_on_idle`)
// which sets LSB (as above) do make Handle::spawn stop working, and then runs until
// num_futures.load() == 1.
let _ = pending;
// We will wait for any blocking ops by virtue of dropping `blocking`.
}
}
impl Executor for CurrentThread {
fn spawn(
&mut self,
future: Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), SpawnError> {
self.borrow().spawner.spawn_local(future, false);
Ok(())
}
}
impl<T> TypedExecutor<T> for CurrentThread
where
T: Future<Output = ()> + 'static,
{
fn spawn(&mut self, future: T) -> Result<(), SpawnError> {
self.borrow().spawner.spawn_local(Box::pin(future), false);
Ok(())
}
}
impl<P: Park> fmt::Debug for CurrentThread<P> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("CurrentThread")
.field("scheduler", &self.scheduler)
.field(
"num_futures",
&self.num_futures.load(atomic::Ordering::SeqCst),
)
.finish()
}
}
impl<P: Park + Default> Default for CurrentThread<P> {
fn default() -> Self {
CurrentThread::new_with_park(P::default())
}
}
// ===== impl Entered =====
impl<P: Park> Entered<'_, P> {
/// Spawn the future on the executor.
///
/// This internally queues the future to be executed once `run` is called.
pub fn spawn<F>(&mut self, future: F) -> &mut Self
where
F: Future<Output = ()> + 'static,
{
self.executor
.borrow()
.spawner
.spawn_local(Box::pin(future), false);
self
}
/// Synchronously waits for the provided `future` to complete.
///
/// This function can be used to synchronously block the current thread
/// until the provided `future` has resolved either successfully or with an
/// error. The result of the future is then returned from this function
/// call.
///
/// Note that this function will **also** execute any spawned futures on the
/// current thread, but will **not** block until these other spawned futures
/// have completed.
///
/// The caller is responsible for ensuring that other spawned futures
/// complete execution.
///
/// # Panics
///
/// This function will panic if the `Park` call returns an error.
pub fn block_on<F>(&mut self, mut future: F) -> F::Output
where
F: Future,
{
// Safety: we shadow the original `future`, so it will never move
// again.
let mut future = unsafe { Pin::new_unchecked(&mut future) };
let waker = self.executor.scheduler.waker();
let mut cx = Context::from_waker(&waker);
loop {
let res = self
.executor
.borrow()
.enter(|| future.as_mut().poll(&mut cx));
match res {
Poll::Ready(e) => return e,
Poll::Pending => {}
}
self.tick();
if self.executor.park.park().is_err() {
panic!("block_on park failed");
}
}
}
/// Run the executor to completion, blocking the thread until **all**
/// spawned futures have completed.
pub fn run(&mut self) -> Result<(), RunError> {
self.run_timeout2(None).map_err(|_| RunError { _p: () })
}
/// Run the executor to completion, blocking the thread until all
/// spawned futures have completed **or** `duration` time has elapsed.
pub fn run_timeout(&mut self, duration: Duration) -> Result<(), RunTimeoutError> {
self.run_timeout2(Some(duration))
}
/// Perform a single iteration of the event loop.
///
/// This function blocks the current thread even if the executor is idle.
pub fn turn(&mut self, duration: Option<Duration>) -> Result<Turn, TurnError> {
let res = if self.executor.scheduler.has_pending_futures() {
self.executor.park.park_timeout(Duration::from_millis(0))
} else {
match duration {
Some(duration) => self.executor.park.park_timeout(duration),
None => self.executor.park.park(),
}
};
if res.is_err() {
return Err(TurnError { _p: () });
}
let polled = self.tick();
Ok(Turn { polled })
}
/// Returns a reference to the underlying `Park` instance.
pub fn get_park(&self) -> &P {
&self.executor.park
}
/// Returns a mutable reference to the underlying `Park` instance.
pub fn get_park_mut(&mut self) -> &mut P {
&mut self.executor.park
}
fn run_timeout2(&mut self, dur: Option<Duration>) -> Result<(), RunTimeoutError> {
if self.executor.is_idle() {
// Nothing to do
return Ok(());
}
let mut time = dur.map(|dur| (Instant::now() + dur, dur));
loop {
self.tick();
if self.executor.is_idle() {
return Ok(());
}
match time {
Some((until, rem)) => {
if self.executor.park.park_timeout(rem).is_err() {
return Err(RunTimeoutError::new(false));
}
let now = Instant::now();
if now >= until {
return Err(RunTimeoutError::new(true));
}
time = Some((until, until - now));
}
None => {
if self.executor.park.park().is_err() {
return Err(RunTimeoutError::new(false));
}
}
}
}
}
/// Returns `true` if any futures were processed
fn tick(&mut self) -> bool {
// Spawn any futures that were spawned from other threads by manually
// looping over the receiver stream
// FIXME: Slightly ugly but needed to make the borrow checker happy
let (mut borrow, spawn_receiver) = (
Borrow {
spawner: BorrowSpawner {
id: self.executor.id,
scheduler: &mut self.executor.scheduler,
num_futures: &*self.executor.num_futures,
},
#[cfg(feature = "blocking")]
blocking: &self.executor.blocking,
},
&mut self.executor.spawn_receiver,
);
while let Ok(future) = spawn_receiver.try_recv() {
borrow.spawner.spawn_local(future, true);
}
// After any pending futures were scheduled, do the actual tick
borrow.spawner.scheduler.tick(TickArgs {
id: borrow.spawner.id,
num_futures: borrow.spawner.num_futures,
#[cfg(feature = "blocking")]
blocking: borrow.blocking,
})
}
}
impl<P: Park> fmt::Debug for Entered<'_, P> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Entered")
.field("executor", &self.executor)
.finish()
}
}
// ===== impl Handle =====
/// Handle to spawn a future on the corresponding `CurrentThread` instance
#[derive(Clone)]
pub struct Handle {
sender: crossbeam_channel::Sender<Pin<Box<dyn Future<Output = ()> + Send + 'static>>>,
num_futures: Arc<atomic::AtomicUsize>,
/// Waker to the Scheduler
waker: Waker,
thread: thread::ThreadId,
/// The thread-local ID assigned to this Handle's executor.
id: u64,
}
// Manual implementation because the Sender does not implement Debug
impl fmt::Debug for Handle {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Handle")
.field("shut_down", &self.is_shut_down())
.finish()
}
}
impl Handle {
/// Spawn a future onto the `CurrentThread` instance corresponding to this handle
///
/// # Panics
///
/// This function panics if the spawn fails. Failure occurs if the `CurrentThread`
/// instance of the `Handle` does not exist anymore.
pub fn spawn<F>(&self, future: F) -> Result<(), SpawnError>
where
F: Future<Output = ()> + Send + 'static,
{
if thread::current().id() == self.thread {
let mut e = TaskExecutor::current();
if e.id() == Some(self.id) {
return e.spawn_local(Box::pin(future));
}
}
// NOTE: += 2 since LSB is the shutdown bit
let pending = self.num_futures.fetch_add(2, atomic::Ordering::SeqCst);
if pending % 2 == 1 {
// Bring the count back so we still know when the Runtime is idle.
self.num_futures.fetch_sub(2, atomic::Ordering::SeqCst);
return Err(SpawnError::shutdown());
}
self.sender
.send(Box::pin(future))
.expect("CurrentThread does not exist anymore");
self.waker.wake_by_ref();
Ok(())
}
/// Provides a best effort **hint** to whether or not `spawn` will succeed.
///
/// This function may return both false positives **and** false negatives.
/// If `status` returns `Ok`, then a call to `spawn` will *probably*
/// succeed, but may fail. If `status` returns `Err`, a call to `spawn` will
/// *probably* fail, but may succeed.
///
/// This allows a caller to avoid creating the task if the call to `spawn`
/// has a high likelihood of failing.
pub fn status(&self) -> Result<(), SpawnError> {
if self.is_shut_down() {
return Err(SpawnError::shutdown());
}
Ok(())
}
fn is_shut_down(&self) -> bool {
// LSB of "num_futures" is the shutdown bit
let num_futures = self.num_futures.load(atomic::Ordering::SeqCst);
num_futures % 2 == 1
}
}
// ===== impl TaskExecutor =====
impl TaskExecutor {
/// Returns an executor that executes futures on the current thread.
///
/// The user of `TaskExecutor` must ensure that when a future is submitted,
/// that it is done within the context of a call to `run`.
///
/// For more details, see the [module level](index.html) documentation.
pub fn current() -> TaskExecutor {
TaskExecutor {
_p: ::std::marker::PhantomData,
}
}
/// Get the current executor's thread-local ID.
fn id(&self) -> Option<u64> {
CURRENT.with(|current| current.id.get())
}
/// Spawn a future onto the current `CurrentThread` instance.
pub fn spawn_local(
&mut self,
future: Pin<Box<dyn Future<Output = ()>>>,
) -> Result<(), SpawnError> {
CURRENT.with(|current| match current.spawn.get() {
Some(spawn) => {
unsafe { (*spawn).spawn_local(future, false) };
Ok(())
}
None => Err(SpawnError::shutdown()),
})
}
}
impl Executor for TaskExecutor {
fn spawn(
&mut self,
future: Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), SpawnError> {
self.spawn_local(future)
}
}
impl<F> TypedExecutor<F> for TaskExecutor
where
F: Future<Output = ()> + 'static,
{
fn spawn(&mut self, future: F) -> Result<(), SpawnError> {
self.spawn_local(Box::pin(future))
}
}
// ===== impl Borrow =====
impl<U: Unpark> Borrow<'_, U> {
fn enter<F, R>(&mut self, f: F) -> R
where
F: FnOnce() -> R,
{
CURRENT.with(|current| {
current.id.set(Some(self.spawner.id));
let Borrow {
ref mut spawner,
#[cfg(all(feature = "blocking", not(loom)))]
ref blocking,
..
} = self;
current.set_spawn(spawner, || {
#[cfg(all(feature = "blocking", not(loom)))]
let res = crate::executor::blocking::with_pool(blocking, || f());
#[cfg(any(not(feature = "blocking"), loom))]
let res = f();
res
})
})
}
}
impl<U: Unpark> SpawnLocal for BorrowSpawner<'_, U> {
fn spawn_local(&mut self, future: Pin<Box<dyn Future<Output = ()>>>, already_counted: bool) {
if !already_counted {
// NOTE: we have a borrow of the Runtime, so we know that it isn't shut down.
// NOTE: += 2 since LSB is the shutdown bit
self.num_futures.fetch_add(2, atomic::Ordering::SeqCst);
}
self.scheduler.schedule(future);
}
}
// ===== impl CurrentRunner =====
impl CurrentRunner {
fn set_spawn<F, R>(&self, spawn: &mut dyn SpawnLocal, f: F) -> R
where
F: FnOnce() -> R,
{
struct Reset<'a>(&'a CurrentRunner);
impl Drop for Reset<'_> {
fn drop(&mut self) {
self.0.spawn.set(None);
self.0.id.set(None);
}
}
let _reset = Reset(self);
let spawn = unsafe { hide_lt(spawn as *mut dyn SpawnLocal) };
self.spawn.set(Some(spawn));
f()
}
}
unsafe fn hide_lt<'a>(p: *mut (dyn SpawnLocal + 'a)) -> *mut (dyn SpawnLocal + 'static) {
use std::mem;
// false positive: https://github.com/rust-lang/rust-clippy/issues/2906
#[allow(clippy::transmute_ptr_to_ptr)]
mem::transmute(p)
}
// ===== impl RunTimeoutError =====
impl RunTimeoutError {
fn new(timeout: bool) -> Self {
RunTimeoutError { timeout }
}
/// Returns `true` if the error was caused by the operation timing out.
pub fn is_timeout(&self) -> bool {
self.timeout
}
}
impl From<EnterError> for RunTimeoutError {
fn from(_: EnterError) -> Self {
RunTimeoutError::new(false)
}
}
// ===== impl BlockError =====
impl<T> BlockError<T> {
/// Returns the error yielded by the future being blocked on
pub fn into_inner(self) -> Option<T> {
self.inner
}
}
impl<T> From<EnterError> for BlockError<T> {
fn from(_: EnterError) -> Self {
BlockError { inner: None }
}
}
tokio/src/executor/current_thread/scheduler.rs
+808
View File
@@ -0,0 +1,808 @@
use crate::executor::current_thread::{Borrow, BorrowSpawner};
use crate::executor::park::Unpark;
use std::cell::UnsafeCell;
use std::fmt::{self, Debug};
use std::future::Future;
use std::mem;
use std::pin::Pin;
use std::ptr;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release, SeqCst};
use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize};
use std::sync::{Arc, Weak};
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
use std::thread;
use std::usize;
/// A generic task-aware scheduler.
///
/// This is used both by `FuturesUnordered` and the current-thread executor.
pub(crate) struct Scheduler<U> {
inner: Arc<Inner<U>>,
nodes: List<U>,
}
// A linked-list of nodes
struct List<U> {
len: usize,
head: *const Node<U>,
tail: *const Node<U>,
}
// Scheduler is implemented using two linked lists. The first linked list tracks
// all items managed by a `Scheduler`. This list is stored on the `Scheduler`
// struct and is **not** thread safe. The second linked list is an
// implementation of the intrusive MPSC queue algorithm described by
// 1024cores.net and is stored on `Inner`. This linked list can push items to
// the back concurrently but only one consumer may pop from the front. To
// enforce this requirement, all popping will be performed via fns on
// `Scheduler` that take `&mut self`.
//
// When a item is submitted to the set a node is allocated and inserted in
// both linked lists. This means that all insertion operations **must** be
// originated from `Scheduler` with `&mut self` The next call to `tick` will
// (eventually) see this node and call `poll` on the item.
//
// Nodes are wrapped in `Arc` cells which manage the lifetime of the node.
// However, `Arc` handles are sometimes cast to `*const Node` pointers.
// Specifically, when a node is stored in at least one of the two lists
// described above, this represents a logical `Arc` handle. This is how
// `Scheduler` maintains its reference to all nodes it manages. Each
// `NotifyHandle` instance is an `Arc<Node>` as well.
//
// When `Scheduler` drops, it clears the linked list of all nodes that it
// manages. When doing so, it must attempt to decrement the reference count (by
// dropping an Arc handle). However, it can **only** decrement the reference
// count if the node is not currently stored in the mpsc channel. If the node
// **is** "queued" in the mpsc channel, then the arc reference count cannot be
// decremented. Once the node is popped from the mpsc channel, then the final
// arc reference count can be decremented, thus freeing the node.
struct Inner<U> {
// Thread unpark handle
unpark: U,
// Tick number
tick_num: AtomicUsize,
// Head/tail of the readiness queue
head_readiness: AtomicPtr<Node<U>>,
tail_readiness: UnsafeCell<*const Node<U>>,
// Used as part of the mpsc queue algorithm
stub: Arc<Node<U>>,
}
unsafe impl<U: Sync + Send> Send for Inner<U> {}
unsafe impl<U: Sync + Send> Sync for Inner<U> {}
struct Node<U> {
// The item
item: UnsafeCell<Option<Task>>,
// The tick at which this node was notified
notified_at: AtomicUsize,
// Next pointer for linked list tracking all active nodes
next_all: UnsafeCell<*const Node<U>>,
// Previous node in linked list tracking all active nodes
prev_all: UnsafeCell<*const Node<U>>,
// Next pointer in readiness queue
next_readiness: AtomicPtr<Node<U>>,
// Whether or not this node is currently in the mpsc queue.
queued: AtomicBool,
// Queue that we'll be enqueued to when notified
queue: Weak<Inner<U>>,
}
/// Returned by `Inner::dequeue`, representing either a dequeue success (with
/// the dequeued node), an empty list, or an inconsistent state.
///
/// The inconsistent state is described in more detail at [1024cores], but
/// roughly indicates that a node will be ready to dequeue sometime shortly in
/// the future and the caller should try again soon.
///
/// [1024cores]: http://www.1024cores.net/home/lock-free-algorithms/queues/intrusive-mpsc-node-based-queue
enum Dequeue<U> {
Data(*const Node<U>),
Empty,
Yield,
Inconsistent,
}
/// Wraps a spawned boxed future
struct Task(Pin<Box<dyn Future<Output = ()>>>);
/// A task that is scheduled. `turn` must be called
pub(crate) struct Scheduled<'a, U> {
task: &'a mut Task,
node: &'a Arc<Node<U>>,
done: &'a mut bool,
}
pub(super) struct TickArgs<'a> {
pub(super) id: u64,
pub(super) num_futures: &'a AtomicUsize,
#[cfg(feature = "blocking")]
pub(super) blocking: &'a crate::executor::blocking::PoolWaiter,
}
impl<U> Scheduler<U>
where
U: Unpark,
{
/// Constructs a new, empty `Scheduler`
///
/// The returned `Scheduler` does not contain any items and, in this
/// state, `Scheduler::poll` will return `Ok(Async::Ready(None))`.
pub(crate) fn new(unpark: U) -> Self {
let stub = Arc::new(Node {
item: UnsafeCell::new(None),
notified_at: AtomicUsize::new(0),
next_all: UnsafeCell::new(ptr::null()),
prev_all: UnsafeCell::new(ptr::null()),
next_readiness: AtomicPtr::new(ptr::null_mut()),
queued: AtomicBool::new(true),
queue: Weak::new(),
});
let stub_ptr = &*stub as *const Node<U>;
let inner = Arc::new(Inner {
unpark,
tick_num: AtomicUsize::new(0),
head_readiness: AtomicPtr::new(stub_ptr as *mut _),
tail_readiness: UnsafeCell::new(stub_ptr),
stub,
});
Scheduler {
inner,
nodes: List::new(),
}
}
pub(crate) fn waker(&self) -> Waker {
waker_inner(self.inner.clone())
}
pub(crate) fn schedule(&mut self, item: Pin<Box<dyn Future<Output = ()>>>) {
// Get the current scheduler tick
let tick_num = self.inner.tick_num.load(SeqCst);
let node = Arc::new(Node {
item: UnsafeCell::new(Some(Task::new(item))),
notified_at: AtomicUsize::new(tick_num),
next_all: UnsafeCell::new(ptr::null_mut()),
prev_all: UnsafeCell::new(ptr::null_mut()),
next_readiness: AtomicPtr::new(ptr::null_mut()),
queued: AtomicBool::new(true),
queue: Arc::downgrade(&self.inner),
});
// Right now our node has a strong reference count of 1. We transfer
// ownership of this reference count to our internal linked list
// and we'll reclaim ownership through the `unlink` function below.
let ptr = self.nodes.push_back(node);
// We'll need to get the item "into the system" to start tracking it,
// e.g. getting its unpark notifications going to us tracking which
// items are ready. To do that we unconditionally enqueue it for
// polling here.
self.inner.enqueue(ptr);
}
/// Returns `true` if there are currently any pending futures
pub(crate) fn has_pending_futures(&mut self) -> bool {
// See function definition for why the unsafe is needed and
// correctly used here
unsafe { self.inner.has_pending_futures() }
}
/// Advance the scheduler state, returning `true` if any futures were
/// processed.
///
/// This function should be called whenever the caller is notified via a
/// wakeup.
pub(super) fn tick(&mut self, args: TickArgs<'_>) -> bool {
let mut ret = false;
let tick = self.inner.tick_num.fetch_add(1, SeqCst).wrapping_add(1);
loop {
let node = match unsafe { self.inner.dequeue(Some(tick)) } {
Dequeue::Empty => {
return ret;
}
Dequeue::Yield => {
self.inner.unpark.unpark();
return ret;
}
Dequeue::Inconsistent => {
thread::yield_now();
continue;
}
Dequeue::Data(node) => node,
};
ret = true;
debug_assert!(node != self.inner.stub());
unsafe {
if (*(*node).item.get()).is_none() {
// The node has already been released. However, while it was
// being released, another thread notified it, which
// resulted in it getting pushed into the mpsc channel.
//
// In this case, we just decrement the ref count.
let node = ptr2arc(node);
assert!((*node.next_all.get()).is_null());
assert!((*node.prev_all.get()).is_null());
continue;
};
// We're going to need to be very careful if the `poll`
// function below panics. We need to (a) not leak memory and
// (b) ensure that we still don't have any use-after-frees. To
// manage this we do a few things:
//
// * This "bomb" here will call `release_node` if dropped
// abnormally. That way we'll be sure the memory management
// of the `node` is managed correctly.
//
// * We unlink the node from our internal queue to preemptively
// assume is is complete (will return Ready or panic), in
// which case we'll want to discard it regardless.
//
struct Bomb<'a, U: Unpark> {
borrow: &'a mut Borrow<'a, U>,
node: Option<Arc<Node<U>>>,
}
impl<U: Unpark> Drop for Bomb<'_, U> {
fn drop(&mut self) {
if let Some(node) = self.node.take() {
self.borrow.enter(|| release_node(node))
}
}
}
let node = self.nodes.remove(node);
let mut borrow = Borrow {
spawner: BorrowSpawner {
id: args.id,
scheduler: self,
num_futures: args.num_futures,
},
#[cfg(feature = "blocking")]
blocking: args.blocking,
};
let mut bomb = Bomb {
node: Some(node),
borrow: &mut borrow,
};
let mut done = false;
// Now that the bomb holds the node, create a new scope. This
// scope ensures that the borrow will go out of scope before we
// mutate the node pointer in `bomb` again
{
let node = bomb.node.as_ref().unwrap();
// Get a reference to the inner future. We already ensured
// that the item `is_some`.
let item = (*node.item.get()).as_mut().unwrap();
// Unset queued flag... this must be done before
// polling. This ensures that the item gets
// rescheduled if it is notified **during** a call
// to `poll`.
let prev = (*node).queued.swap(false, SeqCst);
assert!(prev);
// Poll the underlying item with the appropriate `notify`
// implementation. This is where a large bit of the unsafety
// starts to stem from internally. The `notify` instance itself
// is basically just our `Arc<Node>` and tracks the mpsc
// queue of ready items.
//
// Critically though `Node` won't actually access `Task`, the
// item, while it's floating around inside of `Task`
// instances. These structs will basically just use `T` to size
// the internal allocation, appropriately accessing fields and
// deallocating the node if need be.
let borrow = &mut *bomb.borrow;
let mut scheduled = Scheduled {
task: item,
node: bomb.node.as_ref().unwrap(),
done: &mut done,
};
if borrow.enter(|| scheduled.tick()) {
// we have a borrow of the Runtime, so we know it's not shut down
borrow.spawner.num_futures.fetch_sub(2, SeqCst);
}
}
if !done {
// The future is not done, push it back into the "all
// node" list.
let node = bomb.node.take().unwrap();
bomb.borrow.spawner.scheduler.nodes.push_back(node);
}
}
}
}
}
impl<U: Unpark> Scheduled<'_, U> {
/// Polls the task, returns `true` if the task has completed.
pub(crate) fn tick(&mut self) -> bool {
let waker = unsafe {
// Safety: we don't hold this waker ref longer than
// this `tick` function
waker_ref(self.node)
};
let mut cx = Context::from_waker(&waker);
let ret = match self.task.0.as_mut().poll(&mut cx) {
Poll::Ready(()) => true,
Poll::Pending => false,
};
*self.done = ret;
ret
}
}
impl Task {
pub(crate) fn new(future: Pin<Box<dyn Future<Output = ()> + 'static>>) -> Self {
Task(future)
}
}
impl fmt::Debug for Task {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Task").finish()
}
}
fn release_node<U>(node: Arc<Node<U>>) {
// The item is done, try to reset the queued flag. This will prevent
// `notify` from doing any work in the item
let prev = node.queued.swap(true, SeqCst);
// Drop the item, even if it hasn't finished yet. This is safe
// because we're dropping the item on the thread that owns
// `Scheduler`, which correctly tracks T's lifetimes and such.
unsafe {
drop((*node.item.get()).take());
}
// If the queued flag was previously set then it means that this node
// is still in our internal mpsc queue. We then transfer ownership
// of our reference count to the mpsc queue, and it'll come along and
// free it later, noticing that the item is `None`.
//
// If, however, the queued flag was *not* set then we're safe to
// release our reference count on the internal node. The queued flag
// was set above so all item `enqueue` operations will not actually
// enqueue the node, so our node will never see the mpsc queue again.
// The node itself will be deallocated once all reference counts have
// been dropped by the various owning tasks elsewhere.
if prev {
mem::forget(node);
}
}
impl<U> Debug for Scheduler<U> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "Scheduler {{ ... }}")
}
}
impl<U> Drop for Scheduler<U> {
fn drop(&mut self) {
// When a `Scheduler` is dropped we want to drop all items associated
// with it. At the same time though there may be tons of `Task` handles
// flying around which contain `Node` references inside them. We'll
// let those naturally get deallocated when the `Task` itself goes out
// of scope or gets notified.
while let Some(node) = self.nodes.pop_front() {
release_node(node);
}
// Note that at this point we could still have a bunch of nodes in the
// mpsc queue. None of those nodes, however, have items associated
// with them so they're safe to destroy on any thread. At this point
// the `Scheduler` struct, the owner of the one strong reference
// to `Inner` will drop the strong reference. At that point
// whichever thread releases the strong refcount last (be it this
// thread or some other thread as part of an `upgrade`) will clear out
// the mpsc queue and free all remaining nodes.
//
// While that freeing operation isn't guaranteed to happen here, it's
// guaranteed to happen "promptly" as no more "blocking work" will
// happen while there's a strong refcount held.
}
}
impl<U> Inner<U> {
/// The enqueue function from the 1024cores intrusive MPSC queue algorithm.
fn enqueue(&self, node: *const Node<U>) {
unsafe {
debug_assert!((*node).queued.load(Relaxed));
// This action does not require any coordination
(*node).next_readiness.store(ptr::null_mut(), Relaxed);
// Note that these atomic orderings come from 1024cores
let node = node as *mut _;
let prev = self.head_readiness.swap(node, AcqRel);
(*prev).next_readiness.store(node, Release);
}
}
/// Returns `true` if there are currently any pending futures
///
/// See `dequeue` for an explanation why this function is unsafe.
unsafe fn has_pending_futures(&self) -> bool {
let tail = *self.tail_readiness.get();
let next = (*tail).next_readiness.load(Acquire);
if tail == self.stub() && next.is_null() {
return false;
}
true
}
/// The dequeue function from the 1024cores intrusive MPSC queue algorithm
///
/// Note that this unsafe as it required mutual exclusion (only one thread
/// can call this) to be guaranteed elsewhere.
unsafe fn dequeue(&self, tick: Option<usize>) -> Dequeue<U> {
let mut tail = *self.tail_readiness.get();
let mut next = (*tail).next_readiness.load(Acquire);
if tail == self.stub() {
if next.is_null() {
return Dequeue::Empty;
}
*self.tail_readiness.get() = next;
tail = next;
next = (*next).next_readiness.load(Acquire);
}
if let Some(tick) = tick {
let actual = (*tail).notified_at.load(SeqCst);
// Only dequeue if the node was not scheduled during the current
// tick.
if actual == tick {
// Only doing the check above **should** be enough in
// practice. However, technically there is a potential for
// deadlocking if there are `usize::MAX` ticks while the thread
// scheduling the task is frozen.
//
// If, for some reason, this is not enough, calling `unpark`
// here will resolve the issue.
return Dequeue::Yield;
}
}
if !next.is_null() {
*self.tail_readiness.get() = next;
debug_assert!(tail != self.stub());
return Dequeue::Data(tail);
}
if self.head_readiness.load(Acquire) as *const _ != tail {
return Dequeue::Inconsistent;
}
self.enqueue(self.stub());
next = (*tail).next_readiness.load(Acquire);
if !next.is_null() {
*self.tail_readiness.get() = next;
return Dequeue::Data(tail);
}
Dequeue::Inconsistent
}
fn stub(&self) -> *const Node<U> {
&*self.stub
}
}
impl<U> Drop for Inner<U> {
fn drop(&mut self) {
// Once we're in the destructor for `Inner` we need to clear out the
// mpsc queue of nodes if there's anything left in there.
//
// Note that each node has a strong reference count associated with it
// which is owned by the mpsc queue. All nodes should have had their
// items dropped already by the `Scheduler` destructor above,
// so we're just pulling out nodes and dropping their refcounts.
unsafe {
loop {
match self.dequeue(None) {
Dequeue::Empty => break,
Dequeue::Yield => unreachable!(),
Dequeue::Inconsistent => abort("inconsistent in drop"),
Dequeue::Data(ptr) => drop(ptr2arc(ptr)),
}
}
}
}
}
impl<U> List<U> {
fn new() -> Self {
List {
len: 0,
head: ptr::null_mut(),
tail: ptr::null_mut(),
}
}
/// Appends an element to the back of the list
fn push_back(&mut self, node: Arc<Node<U>>) -> *const Node<U> {
let ptr = arc2ptr(node);
unsafe {
// Point to the current last node in the list
*(*ptr).prev_all.get() = self.tail;
*(*ptr).next_all.get() = ptr::null_mut();
if !self.tail.is_null() {
*(*self.tail).next_all.get() = ptr;
self.tail = ptr;
} else {
// This is the first node
self.tail = ptr;
self.head = ptr;
}
}
self.len += 1;
ptr
}
/// Pop an element from the front of the list
fn pop_front(&mut self) -> Option<Arc<Node<U>>> {
if self.head.is_null() {
// The list is empty
return None;
}
self.len -= 1;
unsafe {
// Convert the ptr to Arc<_>
let node = ptr2arc(self.head);
// Update the head pointer
self.head = *node.next_all.get();
// If the pointer is null, then the list is empty
if self.head.is_null() {
self.tail = ptr::null_mut();
} else {
*(*self.head).prev_all.get() = ptr::null_mut();
}
Some(node)
}
}
/// Remove a specific node
unsafe fn remove(&mut self, node: *const Node<U>) -> Arc<Node<U>> {
let node = ptr2arc(node);
let next = *node.next_all.get();
let prev = *node.prev_all.get();
*node.next_all.get() = ptr::null_mut();
*node.prev_all.get() = ptr::null_mut();
if !next.is_null() {
*(*next).prev_all.get() = prev;
} else {
self.tail = prev;
}
if !prev.is_null() {
*(*prev).next_all.get() = next;
} else {
self.head = next;
}
self.len -= 1;
node
}
}
unsafe fn noop(_: *const ()) {}
// ===== Raw Waker Inner<U> ======
fn waker_inner<U: Unpark>(inner: Arc<Inner<U>>) -> Waker {
let ptr = Arc::into_raw(inner) as *const ();
let vtable = &RawWakerVTable::new(
clone_inner::<U>,
wake_inner::<U>,
wake_by_ref_inner::<U>,
drop_inner::<U>,
);
unsafe { Waker::from_raw(RawWaker::new(ptr, vtable)) }
}
unsafe fn clone_inner<U: Unpark>(data: *const ()) -> RawWaker {
let arc: Arc<Inner<U>> = Arc::from_raw(data as *const Inner<U>);
let clone = arc.clone();
// forget both Arcs so the refcounts don't get decremented
mem::forget(arc);
mem::forget(clone);
let vtable = &RawWakerVTable::new(
clone_inner::<U>,
wake_inner::<U>,
wake_by_ref_inner::<U>,
drop_inner::<U>,
);
RawWaker::new(data, vtable)
}
unsafe fn wake_inner<U: Unpark>(data: *const ()) {
let arc: Arc<Inner<U>> = Arc::from_raw(data as *const Inner<U>);
arc.unpark.unpark();
}
unsafe fn wake_by_ref_inner<U: Unpark>(data: *const ()) {
let arc: Arc<Inner<U>> = Arc::from_raw(data as *const Inner<U>);
arc.unpark.unpark();
// by_ref means we don't own the Node, so forget the Arc
mem::forget(arc);
}
unsafe fn drop_inner<U>(data: *const ()) {
drop(Arc::<Inner<U>>::from_raw(data as *const Inner<U>));
}
// ===== Raw Waker Node<U> ======
unsafe fn waker_ref<U: Unpark>(node: &Arc<Node<U>>) -> Waker {
let ptr = &*node as &Node<U> as *const Node<U> as *const ();
let vtable = &RawWakerVTable::new(
clone_node::<U>,
wake_unreachable,
wake_by_ref_node::<U>,
noop,
);
Waker::from_raw(RawWaker::new(ptr, vtable))
}
unsafe fn wake_unreachable(_data: *const ()) {
unreachable!("waker_ref::wake()");
}
unsafe fn clone_node<U: Unpark>(data: *const ()) -> RawWaker {
let arc: Arc<Node<U>> = Arc::from_raw(data as *const Node<U>);
let clone = arc.clone();
// forget both Arcs so the refcounts don't get decremented
mem::forget(arc);
mem::forget(clone);
let vtable = &RawWakerVTable::new(
clone_node::<U>,
wake_node::<U>,
wake_by_ref_node::<U>,
drop_node::<U>,
);
RawWaker::new(data, vtable)
}
unsafe fn wake_node<U: Unpark>(data: *const ()) {
let arc: Arc<Node<U>> = Arc::from_raw(data as *const Node<U>);
Node::<U>::notify(&arc);
}
unsafe fn wake_by_ref_node<U: Unpark>(data: *const ()) {
let arc: Arc<Node<U>> = Arc::from_raw(data as *const Node<U>);
Node::<U>::notify(&arc);
// by_ref means we don't own the Node, so forget the Arc
mem::forget(arc);
}
unsafe fn drop_node<U>(data: *const ()) {
drop(Arc::<Node<U>>::from_raw(data as *const Node<U>));
}
impl<U: Unpark> Node<U> {
fn notify(me: &Arc<Node<U>>) {
let inner = match me.queue.upgrade() {
Some(inner) => inner,
None => return,
};
// It's our job to notify the node that it's ready to get polled,
// meaning that we need to enqueue it into the readiness queue. To
// do this we flag that we're ready to be queued, and if successful
// we then do the literal queueing operation, ensuring that we're
// only queued once.
//
// Once the node is inserted we be sure to notify the parent task,
// as it'll want to come along and pick up our node now.
//
// Note that we don't change the reference count of the node here,
// we're just enqueueing the raw pointer. The `Scheduler`
// implementation guarantees that if we set the `queued` flag true that
// there's a reference count held by the main `Scheduler` queue
// still.
let prev = me.queued.swap(true, SeqCst);
if !prev {
// Get the current scheduler tick
let tick_num = inner.tick_num.load(SeqCst);
me.notified_at.store(tick_num, SeqCst);
inner.enqueue(&**me);
inner.unpark.unpark();
}
}
}
impl<U> Drop for Node<U> {
fn drop(&mut self) {
// Currently a `Node` is sent across all threads for any lifetime,
// regardless of `T`. This means that for memory safety we can't
// actually touch `T` at any time except when we have a reference to the
// `Scheduler` itself.
//
// Consequently it *should* be the case that we always drop items from
// the `Scheduler` instance, but this is a bomb in place to catch
// any bugs in that logic.
unsafe {
if (*self.item.get()).is_some() {
abort("item still here when dropping");
}
}
}
}
fn arc2ptr<T>(ptr: Arc<T>) -> *const T {
let addr = &*ptr as *const T;
mem::forget(ptr);
addr
}
unsafe fn ptr2arc<T>(ptr: *const T) -> Arc<T> {
let anchor = mem::transmute::<usize, Arc<T>>(0x10);
let addr = &*anchor as *const T;
mem::forget(anchor);
let offset = addr as isize - 0x10;
mem::transmute::<isize, Arc<T>>(ptr as isize - offset)
}
fn abort(s: &str) -> ! {
struct DoublePanic;
impl Drop for DoublePanic {
fn drop(&mut self) {
panic!("panicking twice to abort the program");
}
}
let _bomb = DoublePanic;
panic!("{}", s);
}
tokio/src/executor/enter.rs
+139
View File
@@ -0,0 +1,139 @@
use std::cell::{Cell, RefCell};
use std::error::Error;
use std::fmt;
use std::future::Future;
use std::marker::PhantomData;
thread_local!(static ENTERED: Cell<bool> = Cell::new(false));
/// Represents an executor context.
///
/// For more details, see [`enter` documentation](fn.enter.html)
pub struct Enter {
_p: PhantomData<RefCell<()>>,
}
/// An error returned by `enter` if an execution scope has already been
/// entered.
pub struct EnterError {
_a: (),
}
impl fmt::Debug for EnterError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("EnterError")
.field("reason", &format!("{}", self))
.finish()
}
}
impl fmt::Display for EnterError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
fmt,
"attempted to run an executor while another executor is already running"
)
}
}
impl Error for EnterError {}
/// Marks the current thread as being within the dynamic extent of an
/// executor.
///
/// Executor implementations should call this function before blocking the
/// thread. If `None` is returned, the executor should fail by panicking or
/// taking some other action without blocking the current thread. This prevents
/// deadlocks due to multiple executors competing for the same thread.
///
/// # Error
///
/// Returns an error if the current thread is already marked
pub fn enter() -> Result<Enter, EnterError> {
ENTERED.with(|c| {
if c.get() {
Err(EnterError { _a: () })
} else {
c.set(true);
Ok(Enter { _p: PhantomData })
}
})
}
// Forces the current "entered" state to be cleared while the closure
// is executed.
//
// # Warning
//
// This is hidden for a reason. Do not use without fully understanding
// executors. Misuing can easily cause your program to deadlock.
#[doc(hidden)]
pub fn exit<F: FnOnce() -> R, R>(f: F) -> R {
// Reset in case the closure panics
struct Reset;
impl Drop for Reset {
fn drop(&mut self) {
ENTERED.with(|c| {
c.set(true);
});
}
}
ENTERED.with(|c| {
debug_assert!(c.get());
c.set(false);
});
let reset = Reset;
let ret = f();
::std::mem::forget(reset);
ENTERED.with(|c| {
assert!(!c.get(), "closure claimed permanent executor");
c.set(true);
});
ret
}
impl Enter {
/// Blocks the thread on the specified future, returning the value with
/// which that future completes.
pub fn block_on<F: Future>(&mut self, mut f: F) -> F::Output {
use crate::executor::park::{Park, ParkThread};
use std::pin::Pin;
use std::task::Context;
use std::task::Poll::Ready;
let mut park = ParkThread::new();
let waker = park.unpark().into_waker();
let mut cx = Context::from_waker(&waker);
// `block_on` takes ownership of `f`. Once it is pinned here, the original `f` binding can
// no longer be accessed, making the pinning safe.
let mut f = unsafe { Pin::new_unchecked(&mut f) };
loop {
if let Ready(v) = f.as_mut().poll(&mut cx) {
return v;
}
park.park().unwrap();
}
}
}
impl fmt::Debug for Enter {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Enter").finish()
}
}
impl Drop for Enter {
fn drop(&mut self) {
ENTERED.with(|c| {
assert!(c.get());
c.set(false);
});
}
}
tokio/src/executor/error.rs
+49
View File
@@ -0,0 +1,49 @@
use std::error::Error;
use std::fmt;
/// Errors returned by `Executor::spawn`.
///
/// Spawn errors should represent relatively rare scenarios. Currently, the two
/// scenarios represented by `SpawnError` are:
///
/// * An executor being at capacity or full. As such, the executor is not able
/// to accept a new future. This error state is expected to be transient.
/// * An executor has been shutdown and can no longer accept new futures. This
/// error state is expected to be permanent.
#[derive(Debug)]
pub struct SpawnError {
is_shutdown: bool,
}
impl SpawnError {
/// Return a new `SpawnError` reflecting a shutdown executor failure.
pub fn shutdown() -> Self {
SpawnError { is_shutdown: true }
}
/// Return a new `SpawnError` reflecting an executor at capacity failure.
pub fn at_capacity() -> Self {
SpawnError { is_shutdown: false }
}
/// Returns `true` if the error reflects a shutdown executor failure.
pub fn is_shutdown(&self) -> bool {
self.is_shutdown
}
/// Returns `true` if the error reflects an executor at capacity failure.
pub fn is_at_capacity(&self) -> bool {
!self.is_shutdown
}
}
impl fmt::Display for SpawnError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
fmt,
"attempted to spawn task while the executor is at capacity or shut down"
)
}
}
impl Error for SpawnError {}
tokio/src/executor/executor.rs
+181
View File
@@ -0,0 +1,181 @@
use crate::executor::SpawnError;
use futures_util::future::{FutureExt, RemoteHandle};
use std::future::Future;
use std::pin::Pin;
/// A value that executes futures.
///
/// The [`spawn`] function is used to submit a future to an executor. Once
/// submitted, the executor takes ownership of the future and becomes
/// responsible for driving the future to completion.
///
/// The strategy employed by the executor to handle the future is less defined
/// and is left up to the `Executor` implementation. The `Executor` instance is
/// expected to call [`poll`] on the future once it has been notified, however
/// the "when" and "how" can vary greatly.
///
/// For example, the executor might be a thread pool, in which case a set of
/// threads have already been spawned up and the future is inserted into a
/// queue. A thread will acquire the future and poll it.
///
/// The `Executor` trait is only for futures that **are** `Send`. These are most
/// common. There currently is no trait that describes executors that operate
/// entirely on the current thread (i.e., are able to spawn futures that are not
/// `Send`). Note that single threaded executors can still implement `Executor`,
/// but only futures that are `Send` can be spawned via the trait.
///
/// This trait is primarily intended to implemented by executors and used to
/// back `tokio::spawn`. Libraries and applications **may** use this trait to
/// bound generics, but doing so will limit usage to futures that implement
/// `Send`. Instead, libraries and applications are recommended to use
/// [`TypedExecutor`] as a bound.
///
/// # Errors
///
/// The [`spawn`] function returns `Result` with an error type of `SpawnError`.
/// This error type represents the reason that the executor was unable to spawn
/// the future. The two current represented scenarios are:
///
/// * An executor being at capacity or full. As such, the executor is not able
/// to accept a new future. This error state is expected to be transient.
/// * An executor has been shutdown and can no longer accept new futures. This
/// error state is expected to be permanent.
///
/// If a caller encounters an at capacity error, the caller should try to shed
/// load. This can be as simple as dropping the future that was spawned.
///
/// If the caller encounters a shutdown error, the caller should attempt to
/// gracefully shutdown.
///
/// # Examples
///
/// ```
/// use tokio::executor::Executor;
///
/// # fn docs(my_executor: &mut dyn Executor) {
/// my_executor.spawn(Box::pin(async {
/// println!("running on the executor");
/// })).unwrap();
/// # }
/// ```
///
/// [`spawn`]: #tymethod.spawn
/// [`poll`]: https://doc.rust-lang.org/std/future/trait.Future.html#tymethod.poll
/// [`TypedExecutor`]: ../trait.TypedExecutor.html
pub trait Executor {
/// Spawns a future object to run on this executor.
///
/// `future` is passed to the executor, which will begin running it. The
/// future may run on the current thread or another thread at the discretion
/// of the `Executor` implementation.
///
/// # Panics
///
/// Implementations are encouraged to avoid panics. However, panics are
/// permitted and the caller should check the implementation specific
/// documentation for more details on possible panics.
///
/// # Examples
///
/// ```
/// use tokio::executor::Executor;
///
/// # fn docs(my_executor: &mut dyn Executor) {
/// my_executor.spawn(Box::pin(async {
/// println!("running on the executor");
/// })).unwrap();
/// # }
/// ```
fn spawn(&mut self, future: Pin<Box<dyn Future<Output = ()> + Send>>)
-> Result<(), SpawnError>;
/// Provides a best effort **hint** to whether or not `spawn` will succeed.
///
/// This function may return both false positives **and** false negatives.
/// If `status` returns `Ok`, then a call to `spawn` will *probably*
/// succeed, but may fail. If `status` returns `Err`, a call to `spawn` will
/// *probably* fail, but may succeed.
///
/// This allows a caller to avoid creating the task if the call to `spawn`
/// has a high likelihood of failing.
///
/// # Panics
///
/// This function must not panic. Implementers must ensure that panics do
/// not happen.
///
/// # Examples
///
/// ```
/// use tokio::executor::Executor;
///
/// # fn docs(my_executor: &mut dyn Executor) {
/// if my_executor.status().is_ok() {
/// my_executor.spawn(Box::pin(async {
/// println!("running on the executor");
/// })).unwrap();
/// } else {
/// println!("the executor is not in a good state");
/// }
/// # }
/// ```
fn status(&self) -> Result<(), SpawnError> {
Ok(())
}
}
impl dyn Executor {
/// Spawns a future object to run on this executor, returning a result of
/// its `RemoteHandle`.
///
/// `future` is passed to the executor, which will begin running it. The
/// future may run on the current thread or another thread at the discretion
/// of the `Executor` implementation.
///
/// # Panics
///
/// Implementations are encouraged to avoid panics. However, panics are
/// permitted and the caller should check the implementation specific
/// documentation for more details on possible panics.
///
/// # Examples
///
/// ```
/// use tokio::executor::Executor;
/// use futures_util::future::FutureExt;
///
/// # fn docs(my_executor: &'static mut (dyn Executor + 'static)) {
/// let handle = my_executor.spawn_with_handle(Box::pin(async {
/// println!("running on the executor");
/// })).unwrap();
///
/// let handle = handle.map(|_| println!("the future has completed"));
/// # }
/// ```
pub fn spawn_with_handle<Fut>(
&mut self,
future: Fut,
) -> Result<RemoteHandle<Fut::Output>, SpawnError>
where
Fut: Future + Send + 'static,
Fut::Output: Send,
{
let (future, handle) = future.remote_handle();
self.spawn(Box::pin(future))?;
Ok(handle)
}
}
impl<E: Executor + ?Sized> Executor for Box<E> {
fn spawn(
&mut self,
future: Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), SpawnError> {
(**self).spawn(future)
}
fn status(&self) -> Result<(), SpawnError> {
(**self).status()
}
}
tokio/src/executor/global.rs
+233
View File
@@ -0,0 +1,233 @@
#[cfg(feature = "rt-full")]
use crate::executor::thread_pool::ThreadPool;
use crate::executor::{Executor, SpawnError};
use std::cell::Cell;
use std::future::Future;
use std::pin::Pin;
/// Executes futures on the default executor for the current execution context.
///
/// `DefaultExecutor` implements `Executor` and can be used to spawn futures
/// without referencing a specific executor.
///
/// When an executor starts, it sets the `DefaultExecutor` handle to point to an
/// executor (usually itself) that is used to spawn new tasks.
///
/// The current `DefaultExecutor` reference is tracked using a thread-local
/// variable and is set using `tokio::executor::with_default`
#[derive(Debug, Clone)]
pub struct DefaultExecutor {
_dummy: (),
}
impl DefaultExecutor {
/// Returns a handle to the default executor for the current context.
///
/// Futures may be spawned onto the default executor using this handle.
///
/// The returned handle will reference whichever executor is configured as
/// the default **at the time `spawn` is called**. This enables
/// `DefaultExecutor::current()` to be called before an execution context is
/// setup, then passed **into** an execution context before it is used.
///
/// This is also true for sending the handle across threads, so calling
/// `DefaultExecutor::current()` on thread A and then sending the result to
/// thread B will _not_ reference the default executor that was set on thread A.
pub fn current() -> DefaultExecutor {
DefaultExecutor { _dummy: () }
}
#[inline]
fn with_current<F: FnOnce(&mut dyn Executor) -> R, R>(f: F) -> Option<R> {
EXECUTOR.with(|current_executor| match current_executor.get() {
State::Ready(executor_ptr) => {
let executor = unsafe { &mut *executor_ptr };
Some(f(executor))
}
#[cfg(feature = "rt-full")]
State::ThreadPool(threadpool_ptr) => {
let mut thread_pool = unsafe { &*threadpool_ptr };
Some(f(&mut thread_pool))
}
State::Empty => None,
})
}
}
#[derive(Clone, Copy)]
enum State {
// default executor not defined
Empty,
// default executor is a thread pool instance.
#[cfg(feature = "rt-full")]
ThreadPool(*const ThreadPool),
// default executor is set to a custom executor.
Ready(*mut dyn Executor),
}
thread_local! {
/// Thread-local tracking the current executor
static EXECUTOR: Cell<State> = Cell::new(State::Empty)
}
// ===== impl DefaultExecutor =====
impl super::Executor for DefaultExecutor {
fn spawn(
&mut self,
future: Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), SpawnError> {
DefaultExecutor::with_current(|executor| executor.spawn(future))
.unwrap_or_else(|| Err(SpawnError::shutdown()))
}
fn status(&self) -> Result<(), SpawnError> {
DefaultExecutor::with_current(|executor| executor.status())
.unwrap_or_else(|| Err(SpawnError::shutdown()))
}
}
impl<T> super::TypedExecutor<T> for DefaultExecutor
where
T: Future<Output = ()> + Send + 'static,
{
fn spawn(&mut self, future: T) -> Result<(), SpawnError> {
super::Executor::spawn(self, Box::pin(future))
}
fn status(&self) -> Result<(), SpawnError> {
super::Executor::status(self)
}
}
// ===== global spawn fns =====
/// Spawns a future on the default executor.
///
/// In order for a future to do work, it must be spawned on an executor. The
/// `spawn` function is the easiest way to do this. It spawns a future on the
/// [default executor] for the current execution context (tracked using a
/// thread-local variable).
///
/// The default executor is **usually** a thread pool.
///
/// # Examples
///
/// In this example, a server is started and `spawn` is used to start a new task
/// that processes each received connection.
///
/// ```
/// use tokio::net::TcpListener;
///
/// # async fn process<T>(_t: T) {}
/// # async fn dox() -> Result<(), Box<dyn std::error::Error>> {
/// let mut listener = TcpListener::bind("127.0.0.1:8080").await?;
///
/// loop {
/// let (socket, _) = listener.accept().await?;
///
/// tokio::spawn(async move {
/// // Process each socket concurrently.
/// process(socket).await
/// });
/// }
/// # }
/// ```
///
/// [default executor]: struct.DefaultExecutor.html
///
/// # Panics
///
/// This function will panic if the default executor is not set or if spawning
/// onto the default executor returns an error. To avoid the panic, use
/// [`DefaultExecutor`].
///
/// [`DefaultExecutor`]: struct.DefaultExecutor.html
pub fn spawn<T>(future: T)
where
T: Future<Output = ()> + Send + 'static,
{
EXECUTOR.with(|current_executor| match current_executor.get() {
State::Ready(executor_ptr) => {
let executor = unsafe { &mut *executor_ptr };
executor.spawn(Box::pin(future)).unwrap();
}
#[cfg(feature = "rt-full")]
State::ThreadPool(threadpool_ptr) => {
let thread_pool = unsafe { &*threadpool_ptr };
thread_pool.spawn_background(future);
}
State::Empty => panic!("must be called from the context of Tokio runtime"),
})
}
#[cfg(feature = "rt-full")]
pub(crate) fn with_threadpool<F, R>(thread_pool: &ThreadPool, f: F) -> R
where
F: FnOnce() -> R,
{
with_state(State::ThreadPool(thread_pool as *const ThreadPool), f)
}
/// Set the default executor for the duration of the closure
///
/// If a default executor is already set, it will be restored when the closure returns or if it
/// panics.
pub fn with_default<T, F, R>(executor: &mut T, f: F) -> R
where
T: Executor,
F: FnOnce() -> R,
{
// While scary, this is safe. The function takes a
// `&mut Executor`, which guarantees that the reference lives for the
// duration of `with_default`.
//
// Because we are always clearing the TLS value at the end of the
// function, we can cast the reference to 'static which thread-local
// cells require.
let executor = unsafe { hide_lt(executor as &mut _ as *mut _) };
with_state(State::Ready(executor), f)
}
fn with_state<F, R>(state: State, f: F) -> R
where
F: FnOnce() -> R,
{
EXECUTOR.with(|cell| {
let was = cell.replace(State::Empty);
// Ensure that the executor is removed from the thread-local context
// when leaving the scope. This handles cases that involve panicking.
struct Reset<'a>(&'a Cell<State>, State);
impl Drop for Reset<'_> {
fn drop(&mut self) {
self.0.set(self.1);
}
}
let _reset = Reset(cell, was);
if let State::Ready(executor) = state {
let executor = unsafe { &mut *executor };
if executor.status().is_err() {
panic!("executor not active; is this because `with_default` is called with `DefaultExecutor`?");
}
}
cell.set(state);
f()
})
}
unsafe fn hide_lt<'a>(p: *mut (dyn Executor + 'a)) -> *mut (dyn Executor + 'static) {
use std::mem;
// false positive: https://github.com/rust-lang/rust-clippy/issues/2906
#[allow(clippy::transmute_ptr_to_ptr)]
mem::transmute(p)
}
tokio/src/executor/loom/mod.rs
+27
View File
@@ -0,0 +1,27 @@
//! Stub out the necessary APIs to model with loom.
#[cfg(not(all(test, loom)))]
pub(crate) mod std;
#[cfg(all(test, loom))]
pub(crate) mod std {
pub(crate) use loom::{alloc, cell, sync, thread};
pub(crate) mod rand {
pub(crate) fn seed() -> u64 {
1
}
}
pub(crate) mod sys {
pub(crate) fn num_cpus() -> usize {
2
}
}
}
pub(crate) use self::std::sync;
#[cfg(any(feature = "blocking", feature = "rt-full"))]
pub(crate) use self::std::thread;
#[cfg(feature = "rt-full")]
pub(crate) use self::std::{alloc, cell, rand, sys};
tokio/src/executor/loom/std/atomic_u32.rs
+44
View File
@@ -0,0 +1,44 @@
use std::cell::UnsafeCell;
use std::fmt;
use std::ops::Deref;
/// `AtomicU32` providing an additional `load_unsync` function.
pub(crate) struct AtomicU32 {
inner: UnsafeCell<std::sync::atomic::AtomicU32>,
}
unsafe impl Send for AtomicU32 {}
unsafe impl Sync for AtomicU32 {}
impl AtomicU32 {
pub(crate) fn new(val: u32) -> AtomicU32 {
let inner = UnsafeCell::new(std::sync::atomic::AtomicU32::new(val));
AtomicU32 { inner }
}
/// Perform an unsynchronized load.
///
/// # Safety
///
/// All mutations must have happened before the unsynchronized load.
/// Additionally, there must be no concurrent mutations.
pub(crate) unsafe fn unsync_load(&self) -> u32 {
*(*self.inner.get()).get_mut()
}
}
impl Deref for AtomicU32 {
type Target = std::sync::atomic::AtomicU32;
fn deref(&self) -> &Self::Target {
// safety: it is always safe to access `&self` fns on the inner value as
// we never perform unsafe mutations.
unsafe { &*self.inner.get() }
}
}
impl fmt::Debug for AtomicU32 {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
self.deref().fmt(fmt)
}
}
tokio/src/executor/loom/std/atomic_usize.rs
+45
View File
@@ -0,0 +1,45 @@
use std::cell::UnsafeCell;
use std::fmt;
use std::ops::Deref;
/// `AtomicUsize` providing an additional `load_unsync` function.
pub(crate) struct AtomicUsize {
inner: UnsafeCell<std::sync::atomic::AtomicUsize>,
}
unsafe impl Send for AtomicUsize {}
unsafe impl Sync for AtomicUsize {}
impl AtomicUsize {
pub(crate) fn new(val: usize) -> AtomicUsize {
let inner = UnsafeCell::new(std::sync::atomic::AtomicUsize::new(val));
AtomicUsize { inner }
}
/// Perform an unsynchronized load.
///
/// # Safety
///
/// All mutations must have happened before the unsynchronized load.
/// Additionally, there must be no concurrent mutations.
#[cfg(feature = "rt-full")]
pub(crate) unsafe fn unsync_load(&self) -> usize {
*(*self.inner.get()).get_mut()
}
}
impl Deref for AtomicUsize {
type Target = std::sync::atomic::AtomicUsize;
fn deref(&self) -> &Self::Target {
// safety: it is always safe to access `&self` fns on the inner value as
// we never perform unsafe mutations.
unsafe { &*self.inner.get() }
}
}
impl fmt::Debug for AtomicUsize {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
self.deref().fmt(fmt)
}
}
tokio/src/executor/loom/std/causal_cell.rs
+48
View File
@@ -0,0 +1,48 @@
use std::cell::UnsafeCell;
pub(crate) struct CausalCell<T>(UnsafeCell<T>);
#[derive(Default)]
pub(crate) struct CausalCheck(());
impl<T> CausalCell<T> {
pub(crate) fn new(data: T) -> CausalCell<T> {
CausalCell(UnsafeCell::new(data))
}
pub(crate) fn with<F, R>(&self, f: F) -> R
where
F: FnOnce(*const T) -> R,
{
f(self.0.get())
}
pub(crate) fn with_unchecked<F, R>(&self, f: F) -> R
where
F: FnOnce(*const T) -> R,
{
f(self.0.get())
}
pub(crate) fn check(&self) {}
pub(crate) fn with_deferred<F, R>(&self, f: F) -> (R, CausalCheck)
where
F: FnOnce(*const T) -> R,
{
(f(self.0.get()), CausalCheck::default())
}
pub(crate) fn with_mut<F, R>(&self, f: F) -> R
where
F: FnOnce(*mut T) -> R,
{
f(self.0.get())
}
}
impl CausalCheck {
pub(crate) fn check(self) {}
pub(crate) fn join(&mut self, _other: CausalCheck) {}
}
tokio/src/executor/loom/std/mod.rs
+77
View File
@@ -0,0 +1,77 @@
#[cfg(feature = "rt-full")]
mod atomic_u32;
mod atomic_usize;
#[cfg(feature = "rt-full")]
mod causal_cell;
#[cfg(feature = "rt-full")]
pub(crate) mod alloc {
#[derive(Debug)]
pub(crate) struct Track<T> {
value: T,
}
impl<T> Track<T> {
pub(crate) fn new(value: T) -> Track<T> {
Track { value }
}
pub(crate) fn get_mut(&mut self) -> &mut T {
&mut self.value
}
pub(crate) fn into_inner(self) -> T {
self.value
}
}
}
#[cfg(feature = "rt-full")]
pub(crate) mod cell {
pub(crate) use super::causal_cell::{CausalCell, CausalCheck};
}
#[cfg(feature = "rt-full")]
pub(crate) mod rand {
use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hash, Hasher};
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering::Relaxed;
static COUNTER: AtomicU32 = AtomicU32::new(1);
pub(crate) fn seed() -> u64 {
let rand_state = RandomState::new();
let mut hasher = rand_state.build_hasher();
// Hash some unique-ish data to generate some new state
COUNTER.fetch_add(1, Relaxed).hash(&mut hasher);
// Get the seed
hasher.finish()
}
}
pub(crate) mod sync {
pub(crate) use std::sync::{Arc, Condvar, Mutex};
pub(crate) mod atomic {
#[cfg(feature = "rt-full")]
pub(crate) use crate::executor::loom::std::atomic_u32::AtomicU32;
pub(crate) use crate::executor::loom::std::atomic_usize::AtomicUsize;
#[cfg(feature = "rt-full")]
pub(crate) use std::sync::atomic::{fence, spin_loop_hint, AtomicPtr};
}
}
#[cfg(feature = "rt-full")]
pub(crate) mod sys {
pub(crate) fn num_cpus() -> usize {
usize::max(1, num_cpus::get_physical())
}
}
#[cfg(any(feature = "blocking", feature = "rt-full"))]
pub(crate) use std::thread;
tokio/src/executor/mod.rs
+84
View File
@@ -0,0 +1,84 @@
//! Task execution related traits and utilities.
//!
//! In the Tokio execution model, futures are lazy. When a future is created, no
//! work is performed. In order for the work defined by the future to happen,
//! the future must be submitted to an executor. A future that is submitted to
//! an executor is called a "task".
//!
//! The executor is responsible for ensuring that [`Future::poll`] is called
//! whenever the task is notified. Notification happens when the internal
//! state of a task transitions from *not ready* to *ready*. For example, a
//! socket might have received data and a call to `read` will now be able to
//! succeed.
//!
//! The specific strategy used to manage the tasks is left up to the
//! executor. There are two main flavors of executors: single-threaded and
//! multi-threaded. Tokio provides implementation for both of these in the
//! [`runtime`] module.
//!
//! # `Executor` trait.
//!
//! This module provides the [`Executor`] trait (re-exported from
//! [`tokio-executor`]), which describes the API that all executors must
//! implement.
//!
//! A free [`spawn`] function is provided that allows spawning futures onto the
//! default executor (tracked via a thread-local variable) without referencing a
//! handle. It is expected that all executors will set a value for the default
//! executor. This value will often be set to the executor itself, but it is
//! possible that the default executor might be set to a different executor.
//!
//! For example, a single threaded executor might set the default executor to a
//! thread pool instead of itself, allowing futures to spawn new tasks onto the
//! thread pool when those tasks are `Send`.
//!
//! [`Future::poll`]: https://docs.rs/futures/0.1/futures/future/trait.Future.html#tymethod.poll
//! [notified]: https://docs.rs/futures/0.1/futures/executor/trait.Notify.html#tymethod.notify
//! [`runtime`]: ../runtime/index.html
//! [`tokio-executor`]: https://docs.rs/tokio-executor/0.1
//! [`Executor`]: trait.Executor.html
//! [`spawn`]: fn.spawn.html#[cfg(all(test, loom))]
// At the top due to macros
#[cfg(test)]
#[macro_use]
mod tests;
mod enter;
pub use self::enter::{enter, exit, Enter, EnterError};
mod error;
pub use self::error::SpawnError;
#[allow(clippy::module_inception)]
mod executor;
pub use self::executor::Executor;
mod global;
pub use self::global::{spawn, with_default, DefaultExecutor};
mod loom;
pub mod park;
#[cfg(feature = "rt-full")]
mod task;
mod typed;
pub use self::typed::TypedExecutor;
#[cfg(feature = "rt-full")]
mod util;
#[cfg(all(not(feature = "blocking"), feature = "rt-full"))]
mod blocking;
#[cfg(feature = "blocking")]
pub mod blocking;
#[cfg(feature = "rt-current-thread")]
pub mod current_thread;
#[cfg(feature = "rt-full")]
pub mod thread_pool;
pub use futures_util::future::RemoteHandle;
tokio/src/executor/park/mod.rs
+140
View File
@@ -0,0 +1,140 @@
//! Abstraction over blocking and unblocking the current thread.
//!
//! Provides an abstraction over blocking the current thread. This is similar to
//! the park / unpark constructs provided by [`std`] but made generic. This
//! allows embedding custom functionality to perform when the thread is blocked.
//!
//! A blocked [`Park`][p] instance is unblocked by calling [`unpark`] on its
//! [`Unpark`][up] handle.
//!
//! The [`ParkThread`] struct implements [`Park`][p] using
//! [`thread::park`][`std`] to put the thread to sleep. The Tokio reactor also
//! implements park, but uses [`mio::Poll`][mio] to block the thread instead.
//!
//! The [`Park`][p] trait is composable. A timer implementation might decorate a
//! [`Park`][p] implementation by checking if any timeouts have elapsed after
//! the inner [`Park`][p] implementation unblocks.
//!
//! # Model
//!
//! Conceptually, each [`Park`][p] instance has an associated token, which is
//! initially not present:
//!
//! * The [`park`] method blocks the current thread unless or until the token
//! is available, at which point it atomically consumes the token.
//! * The [`unpark`] method atomically makes the token available if it wasn't
//! already.
//!
//! Some things to note:
//!
//! * If [`unpark`] is called before [`park`], the next call to [`park`] will
//! **not** block the thread.
//! * **Spurious** wakeups are permitted, i.e., the [`park`] method may unblock
//! even if [`unpark`] was not called.
//! * [`park_timeout`] does the same as [`park`] but allows specifying a maximum
//! time to block the thread for.
//!
//! [`std`]: https://doc.rust-lang.org/std/thread/fn.park.html
//! [`thread::park`]: https://doc.rust-lang.org/std/thread/fn.park.html
//! [`ParkThread`]: struct.ParkThread.html
//! [p]: trait.Park.html
//! [`park`]: trait.Park.html#tymethod.park
//! [`park_timeout`]: trait.Park.html#tymethod.park_timeout
//! [`unpark`]: trait.Unpark.html#tymethod.unpark
//! [up]: trait.Unpark.html
//! [mio]: https://docs.rs/mio/0.6/mio/struct.Poll.html
mod thread;
pub use self::thread::{ParkError, ParkThread, UnparkThread};
use std::sync::Arc;
use std::time::Duration;
/// Block the current thread.
///
/// See [module documentation][mod] for more details.
///
/// [mod]: ../index.html
pub trait Park {
/// Unpark handle type for the `Park` implementation.
type Unpark: Unpark;
/// Error returned by `park`
type Error;
/// Get a new `Unpark` handle associated with this `Park` instance.
fn unpark(&self) -> Self::Unpark;
/// Block the current thread unless or until the token is available.
///
/// A call to `park` does not guarantee that the thread will remain blocked
/// forever, and callers should be prepared for this possibility. This
/// function may wakeup spuriously for any reason.
///
/// See [module documentation][mod] for more details.
///
/// # Panics
///
/// This function **should** not panic, but ultimately, panics are left as
/// an implementation detail. Refer to the documentation for the specific
/// `Park` implementation
///
/// [mod]: ../index.html
fn park(&mut self) -> Result<(), Self::Error>;
/// Park the current thread for at most `duration`.
///
/// This function is the same as `park` but allows specifying a maximum time
/// to block the thread for.
///
/// Same as `park`, there is no guarantee that the thread will remain
/// blocked for any amount of time. Spurious wakeups are permitted for any
/// reason.
///
/// See [module documentation][mod] for more details.
///
/// # Panics
///
/// This function **should** not panic, but ultimately, panics are left as
/// an implementation detail. Refer to the documentation for the specific
/// `Park` implementation
///
/// [mod]: ../index.html
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error>;
}
/// Unblock a thread blocked by the associated [`Park`] instance.
///
/// See [module documentation][mod] for more details.
///
/// [mod]: ../index.html
/// [`Park`]: trait.Park.html
pub trait Unpark: Sync + Send + 'static {
/// Unblock a thread that is blocked by the associated `Park` handle.
///
/// Calling `unpark` atomically makes available the unpark token, if it is
/// not already available.
///
/// See [module documentation][mod] for more details.
///
/// # Panics
///
/// This function **should** not panic, but ultimately, panics are left as
/// an implementation detail. Refer to the documentation for the specific
/// `Unpark` implementation
///
/// [mod]: ../index.html
fn unpark(&self);
}
impl Unpark for Box<dyn Unpark> {
fn unpark(&self) {
(**self).unpark()
}
}
impl Unpark for Arc<dyn Unpark> {
fn unpark(&self) {
(**self).unpark()
}
}
tokio/src/executor/park/thread.rs
+263
View File
@@ -0,0 +1,263 @@
use crate::executor::loom::sync::atomic::AtomicUsize;
use crate::executor::loom::sync::{Arc, Condvar, Mutex};
use crate::executor::park::{Park, Unpark};
use std::marker::PhantomData;
use std::mem;
use std::rc::Rc;
use std::sync::atomic::Ordering;
use std::task::{RawWaker, RawWakerVTable, Waker};
use std::time::Duration;
/// Blocks the current thread using a condition variable.
///
/// Implements the [`Park`] functionality by using a condition variable. An
/// atomic variable is also used to avoid using the condition variable if
/// possible.
///
/// The condition variable is cached in a thread-local variable and is shared
/// across all `ParkThread` instances created on the same thread. This also
/// means that an instance of `ParkThread` might be unblocked by a handle
/// associated with a different `ParkThread` instance.
#[derive(Debug)]
pub struct ParkThread {
_anchor: PhantomData<Rc<()>>,
}
/// Error returned by [`ParkThread`]
///
/// This currently is never returned, but might at some point in the future.
///
/// [`ParkThread`]: struct.ParkThread.html
#[derive(Debug)]
pub struct ParkError {
_p: (),
}
struct Parker {
unparker: Arc<Inner>,
}
/// Unblocks a thread that was blocked by `ParkThread`.
#[derive(Clone, Debug)]
pub struct UnparkThread {
inner: Arc<Inner>,
}
#[derive(Debug)]
struct Inner {
state: AtomicUsize,
mutex: Mutex<()>,
condvar: Condvar,
}
const IDLE: usize = 0;
const NOTIFY: usize = 1;
const SLEEP: usize = 2;
thread_local! {
static CURRENT_PARKER: Parker = Parker::new();
}
// ==== impl Parker ====
impl Parker {
pub(crate) fn new() -> Self {
Self {
unparker: Arc::new(Inner {
state: AtomicUsize::new(IDLE),
mutex: Mutex::new(()),
condvar: Condvar::new(),
}),
}
}
pub(crate) fn unparker(&self) -> &Arc<Inner> {
&self.unparker
}
pub(crate) fn park(&self) -> Result<(), ParkError> {
self.unparker.park(None)
}
pub(crate) fn park_timeout(&self, timeout: Duration) -> Result<(), ParkError> {
self.unparker.park(Some(timeout))
}
}
// ==== impl Inner ====
impl Inner {
#[allow(clippy::wrong_self_convention)]
pub(crate) fn into_raw(this: Arc<Inner>) -> *const () {
Arc::into_raw(this) as *const ()
}
pub(crate) unsafe fn from_raw(ptr: *const ()) -> Arc<Inner> {
Arc::from_raw(ptr as *const Inner)
}
/// Park the current thread for at most `dur`.
pub(crate) fn park(&self, timeout: Option<Duration>) -> Result<(), ParkError> {
// If currently notified, then we skip sleeping. This is checked outside
// of the lock to avoid acquiring a mutex if not necessary.
match self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) {
NOTIFY => return Ok(()),
IDLE => {}
_ => unreachable!(),
}
// The state is currently idle, so obtain the lock and then try to
// transition to a sleeping state.
let mut m = self.mutex.lock().unwrap();
// Transition to sleeping
match self.state.compare_and_swap(IDLE, SLEEP, Ordering::SeqCst) {
NOTIFY => {
// Notified before we could sleep, consume the notification and
// exit
self.state.store(IDLE, Ordering::SeqCst);
return Ok(());
}
IDLE => {}
_ => unreachable!(),
}
m = match timeout {
Some(timeout) => self.condvar.wait_timeout(m, timeout).unwrap().0,
None => self.condvar.wait(m).unwrap(),
};
// Transition back to idle. If the state has transitioned to `NOTIFY`,
// this will consume that notification
self.state.store(IDLE, Ordering::SeqCst);
// Explicitly drop the mutex guard. There is no real point in doing it
// except that I find it helpful to make it explicit where we want the
// mutex to unlock.
drop(m);
Ok(())
}
pub(crate) fn unpark(&self) {
// First, try transitioning from IDLE -> NOTIFY, this does not require a
// lock.
match self.state.compare_and_swap(IDLE, NOTIFY, Ordering::SeqCst) {
IDLE | NOTIFY => return,
SLEEP => {}
_ => unreachable!(),
}
// The other half is sleeping, this requires a lock
let _m = self.mutex.lock().unwrap();
// Transition to NOTIFY
match self.state.swap(NOTIFY, Ordering::SeqCst) {
SLEEP => {}
NOTIFY => return,
IDLE => return,
_ => unreachable!(),
}
// Wakeup the sleeper
self.condvar.notify_one();
}
}
// ===== impl ParkThread =====
impl ParkThread {
/// Create a new `ParkThread` handle for the current thread.
///
/// This type cannot be moved to other threads, so it should be created on
/// the thread that the caller intends to park.
pub fn new() -> ParkThread {
ParkThread {
_anchor: PhantomData,
}
}
/// Get a reference to the `ParkThread` handle for this thread.
fn with_current<F, R>(&self, f: F) -> R
where
F: FnOnce(&Parker) -> R,
{
CURRENT_PARKER.with(|inner| f(inner))
}
}
impl Park for ParkThread {
type Unpark = UnparkThread;
type Error = ParkError;
fn unpark(&self) -> Self::Unpark {
let inner = self.with_current(|inner| inner.unparker().clone());
UnparkThread { inner }
}
fn park(&mut self) -> Result<(), Self::Error> {
self.with_current(|inner| inner.park())?;
Ok(())
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
self.with_current(|inner| inner.park_timeout(duration))?;
Ok(())
}
}
impl Default for ParkThread {
fn default() -> Self {
Self::new()
}
}
// ===== impl UnparkThread =====
impl Unpark for UnparkThread {
fn unpark(&self) {
self.inner.unpark();
}
}
static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop_waker);
impl UnparkThread {
pub(crate) fn into_waker(self) -> Waker {
unsafe {
let raw = unparker_to_raw_waker(self.inner);
Waker::from_raw(raw)
}
}
}
unsafe fn unparker_to_raw_waker(unparker: Arc<Inner>) -> RawWaker {
RawWaker::new(Inner::into_raw(unparker), &VTABLE)
}
unsafe fn clone(raw: *const ()) -> RawWaker {
let unparker = Inner::from_raw(raw);
// Increment the ref count
mem::forget(unparker.clone());
unparker_to_raw_waker(unparker)
}
unsafe fn drop_waker(raw: *const ()) {
let _ = Inner::from_raw(raw);
}
unsafe fn wake(raw: *const ()) {
let unparker = Inner::from_raw(raw);
unparker.unpark();
}
unsafe fn wake_by_ref(raw: *const ()) {
let unparker = Inner::from_raw(raw);
unparker.unpark();
// We don't actually own a reference to the unparker
mem::forget(unparker);
}
tokio/src/executor/task/core.rs
+153
View File
@@ -0,0 +1,153 @@
use crate::executor::loom::alloc::Track;
use crate::executor::loom::cell::CausalCell;
use crate::executor::task::raw::{self, Vtable};
use crate::executor::task::state::State;
use crate::executor::task::waker::waker_ref;
use crate::executor::task::Schedule;
use std::cell::UnsafeCell;
use std::future::Future;
use std::mem::MaybeUninit;
use std::pin::Pin;
use std::ptr::{self, NonNull};
use std::task::{Context, Poll, Waker};
/// The task cell. Contains the components of the task.
///
/// It is critical for `Header` to be the first field as the task structure will
/// be referenced by both *mut Cell and *mut Header.
#[repr(C)]
pub(super) struct Cell<T: Future, S: 'static> {
/// Hot task state data
pub(super) header: Header<S>,
/// Either the future or output, depending on the execution stage.
pub(super) core: Core<T>,
/// Cold data
pub(super) trailer: Trailer,
}
/// The core of the task.
///
/// Holds the future or output, depending on the stage of execution.
pub(super) struct Core<T: Future> {
stage: Stage<T>,
}
/// Crate public as this is also needed by the pool.
#[repr(C)]
pub(crate) struct Header<S: 'static> {
/// Task state
pub(super) state: State,
/// Pointer to the executor owned by the task
pub(super) executor: CausalCell<Option<NonNull<S>>>,
/// Pointer to next task, used for misc task linked lists.
pub(crate) queue_next: UnsafeCell<*const Header<S>>,
/// Pointer to the next task in the ownership list.
pub(crate) owned_next: UnsafeCell<Option<NonNull<Header<S>>>>,
/// Pointer to the previous task in the ownership list.
pub(crate) owned_prev: UnsafeCell<Option<NonNull<Header<S>>>>,
/// Table of function pointers for executing actions on the task.
pub(super) vtable: &'static Vtable<S>,
/// Used by loom to track the causality of the future. Without loom, this is
/// unit.
pub(super) future_causality: CausalCell<()>,
}
/// Cold data is stored after the future.
pub(super) struct Trailer {
/// Consumer task waiting on completion of this task.
pub(super) waker: CausalCell<MaybeUninit<Option<Waker>>>,
}
/// Either the future or the output.
enum Stage<T: Future> {
Running(Track<T>),
Finished(Track<super::Result<T::Output>>),
Consumed,
}
impl<T: Future, S: Schedule> Cell<T, S> {
/// Allocate a new task cell, containing the header, trailer, and core
/// structures.
pub(super) fn new(future: T, state: State) -> Box<Cell<T, S>> {
Box::new(Cell {
header: Header {
state,
executor: CausalCell::new(None),
queue_next: UnsafeCell::new(ptr::null()),
owned_next: UnsafeCell::new(None),
owned_prev: UnsafeCell::new(None),
vtable: raw::vtable::<T, S>(),
future_causality: CausalCell::new(()),
},
core: Core {
stage: Stage::Running(Track::new(future)),
},
trailer: Trailer {
waker: CausalCell::new(MaybeUninit::new(None)),
},
})
}
}
impl<T: Future> Core<T> {
pub(super) fn transition_to_consumed(&mut self) {
self.stage = Stage::Consumed
}
pub(super) fn poll<S>(&mut self, header: &Header<S>) -> Poll<T::Output>
where
S: Schedule,
{
let res = {
let future = match &mut self.stage {
Stage::Running(tracked) => tracked.get_mut(),
_ => unreachable!("unexpected stage"),
};
// The future is pinned within the task. The above state transition
// has ensured the safety of this action.
let future = unsafe { Pin::new_unchecked(future) };
// The waker passed into the `poll` function does not require a ref
// count increment.
let waker_ref = waker_ref::<T, S>(header);
let mut cx = Context::from_waker(&*waker_ref);
future.poll(&mut cx)
};
if res.is_ready() {
self.stage = Stage::Consumed;
}
res
}
pub(super) fn store_output(&mut self, output: super::Result<T::Output>) {
self.stage = Stage::Finished(Track::new(output));
}
pub(super) unsafe fn read_output(&mut self, dst: *mut Track<super::Result<T::Output>>) {
use std::mem;
dst.write(match mem::replace(&mut self.stage, Stage::Consumed) {
Stage::Finished(output) => output,
_ => unreachable!("unexpected state"),
});
}
}
impl<S> Header<S> {
pub(super) fn executor(&self) -> Option<NonNull<S>> {
unsafe { self.executor.with(|ptr| *ptr) }
}
}
tokio/src/executor/task/error.rs
+48
View File
@@ -0,0 +1,48 @@
use std::any::Any;
use std::fmt;
/// Task failed to execute to completion.
pub struct Error {
repr: Repr,
}
enum Repr {
Cancelled,
Panic(Box<dyn Any + Send + 'static>),
}
impl Error {
/// Create a new `cancelled` error
pub fn cancelled() -> Error {
Error {
repr: Repr::Cancelled,
}
}
/// Create a new `panic` error
pub fn panic(err: Box<dyn Any + Send + 'static>) -> Error {
Error {
repr: Repr::Panic(err),
}
}
}
impl fmt::Display for Error {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
match &self.repr {
Repr::Cancelled => write!(fmt, "cancelled"),
Repr::Panic(_) => write!(fmt, "panic"),
}
}
}
impl fmt::Debug for Error {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
match &self.repr {
Repr::Cancelled => write!(fmt, "task::Error::Cancelled"),
Repr::Panic(_) => write!(fmt, "task::Error::Panic(...)"),
}
}
}
impl std::error::Error for Error {}
tokio/src/executor/task/harness.rs
+546
View File
@@ -0,0 +1,546 @@
use crate::executor::loom::alloc::Track;
use crate::executor::loom::cell::CausalCheck;
use crate::executor::task::core::{Cell, Core, Header, Trailer};
use crate::executor::task::state::Snapshot;
use crate::executor::task::{Error, Schedule, Task};
use std::future::Future;
use std::mem::{ManuallyDrop, MaybeUninit};
use std::ptr::NonNull;
use std::task::{Poll, Waker};
/// Typed raw task handle
pub(super) struct Harness<T: Future, S: 'static> {
cell: NonNull<Cell<T, S>>,
}
impl<T, S> Harness<T, S>
where
T: Future,
S: 'static,
{
pub(super) unsafe fn from_raw(ptr: *mut ()) -> Harness<T, S> {
debug_assert!(!ptr.is_null());
let cell = NonNull::new_unchecked(ptr as *mut Cell<T, S>);
Harness { cell }
}
fn header(&self) -> &Header<S> {
unsafe { &self.cell.as_ref().header }
}
fn trailer(&self) -> &Trailer {
unsafe { &self.cell.as_ref().trailer }
}
fn core(&mut self) -> &mut Core<T> {
unsafe { &mut self.cell.as_mut().core }
}
}
impl<T, S> Harness<T, S>
where
T: Future,
S: Schedule,
{
/// Poll the inner future.
///
/// All necessary state checks and transitions are performed.
///
/// Panics raised while polling the future are handled.
///
/// Returns `true` if the task needs to be scheduled again
pub(super) fn poll(mut self, executor: NonNull<S>) -> bool {
use std::panic;
// Transition the task to the running state.
let res = self.header().state.transition_to_running();
if res.is_canceled() {
// The task was concurrently canceled.
self.do_cancel(res);
return false;
}
let join_interest = res.is_join_interested();
debug_assert!(join_interest || !res.has_join_waker());
// Get the cell components
let cell = unsafe { &mut self.cell.as_mut() };
let header = &cell.header;
let core = &mut cell.core;
// If the task's executor pointer is not yet set, then set it here. This
// is safe because a) this is the only time the value is set. b) at this
// point, there are no outstanding wakers which might access the
// field concurrently.
if header.executor().is_none() {
unsafe {
// We don't want the destructor to run because we don't really
// own the task here.
let task = ManuallyDrop::new(Task::from_raw(header.into()));
// Call the scheduler's bind callback
executor.as_ref().bind(&task);
header.executor.with_mut(|ptr| *ptr = Some(executor));
}
}
// The transition to `Running` done above ensures that a lock on the
// future has been obtained. This also ensures the `*mut T` pointer
// contains the future (as opposed to the output) and is initialized.
let res = header.future_causality.with_mut(|_| {
panic::catch_unwind(panic::AssertUnwindSafe(|| {
struct Guard<'a, T: Future> {
core: &'a mut Core<T>,
polled: bool,
}
impl<T: Future> Drop for Guard<'_, T> {
fn drop(&mut self) {
if !self.polled {
self.core.transition_to_consumed();
}
}
}
let mut guard = Guard {
core,
polled: false,
};
let res = guard.core.poll(header);
// prevent the guard from dropping the future
guard.polled = true;
res
}))
});
match res {
Ok(Poll::Ready(out)) => {
self.complete(executor, join_interest, Ok(out));
false
}
Ok(Poll::Pending) => {
let res = self.header().state.transition_to_idle();
if res.is_canceled() {
self.do_cancel(res);
false
} else {
res.is_notified()
}
}
Err(err) => {
self.complete(executor, join_interest, Err(Error::panic(err)));
false
}
}
}
pub(super) unsafe fn drop_task(mut self) {
let might_drop_join_waker_on_release = self.might_drop_join_waker_on_release();
// Read the join waker cell just to have it
let (join_waker, check) = self.read_join_waker();
// transition the task to released
let res = self.header().state.release_task();
assert!(res.is_terminal(), "state = {:?}", res);
if might_drop_join_waker_on_release && !res.is_join_interested() {
debug_assert!(res.has_join_waker());
// Its our responsibility to drop the waker
check.check();
let _ = join_waker.assume_init();
}
if res.is_final_ref() {
self.dealloc();
}
}
unsafe fn dealloc(self) {
// Check causality
self.header().executor.with_mut(|_| {});
self.header().future_causality.with_mut(|_| {});
self.trailer().waker.with_mut(|_| {
// we can't check the contents of this cell as it is considered
// "uninitialized" data at this point.
});
drop(Box::from_raw(self.cell.as_ptr()));
}
// ===== join handle =====
pub(super) unsafe fn read_output(
mut self,
dst: *mut Track<super::Result<T::Output>>,
state: Snapshot,
) {
if state.is_canceled() {
dst.write(Track::new(Err(Error::cancelled())));
} else {
self.core().read_output(dst);
}
// Before transitioning the state, the waker must be read. It is
// possible that, after the transition, we are responsible for dropping
// the waker but before the waker can be read from the struct, the
// struct is deallocated.
let (waker, check) = self.read_join_waker();
// The operation counts as dropping the join handle
let res = self.header().state.complete_join_handle();
if res.is_released() {
// We are responsible for freeing the waker handle
check.check();
drop(waker.assume_init());
}
if res.is_final_ref() {
self.dealloc();
}
}
pub(super) fn store_join_waker(&self, waker: &Waker) -> Snapshot {
unsafe {
self.trailer().waker.with_mut(|ptr| {
(*ptr).as_mut_ptr().replace(Some(waker.clone()));
});
}
let res = self.header().state.store_join_waker();
if res.is_complete() || res.is_canceled() {
// Drop the waker here
self.trailer()
.waker
.with_mut(|ptr| unsafe { *(*ptr).as_mut_ptr() = None });
}
res
}
pub(super) fn swap_join_waker(&self, waker: &Waker, prev: Snapshot) -> Snapshot {
unsafe {
let will_wake = self
.trailer()
.waker
.with(|ptr| (*(*ptr).as_ptr()).as_ref().unwrap().will_wake(waker));
if will_wake {
return prev;
}
// Acquire the lock
let state = self.header().state.unset_waker();
if state.is_active() {
return self.store_join_waker(waker);
}
state
}
}
pub(super) fn drop_join_handle_slow(mut self) {
unsafe {
// Before transitioning the state, the waker must be read. It is
// possible that, after the transition, we are responsible for dropping
// the waker but before the waker can be read from the struct, the
// struct is deallocated.
let (waker, check) = self.read_join_waker();
// The operation counts as dropping the join handle
let res = match self.header().state.drop_join_handle_slow() {
Ok(res) => res,
Err(res) => {
// The task output must be read & dropped
debug_assert!(!(res.is_complete() && res.is_canceled()));
if res.is_complete() {
self.core().transition_to_consumed();
}
self.header().state.complete_join_handle()
}
};
if !(res.is_complete() | res.is_canceled()) || res.is_released() {
// We are responsible for freeing the waker handle
check.check();
drop(waker.assume_init());
}
if res.is_final_ref() {
self.dealloc();
}
}
}
// ===== waker behavior =====
pub(super) fn wake_by_val(self) {
self.wake_by_ref();
self.drop_waker();
}
pub(super) fn wake_by_local_ref(&self) {
self.wake_by_ref();
}
pub(super) fn wake_by_ref(&self) {
if self.header().state.transition_to_notified() {
unsafe {
let executor = match self.header().executor.with(|ptr| *ptr) {
Some(executor) => executor,
None => panic!("executor should be set"),
};
S::schedule(executor.as_ref(), self.to_task());
}
}
}
pub(super) fn drop_waker(self) {
if self.header().state.ref_dec() {
unsafe {
self.dealloc();
}
}
}
/// Cancel the task.
///
/// `from_queue` signals the caller is cancelling the task after popping it
/// from the queue. This indicates "polling" capability.
pub(super) fn cancel(self, from_queue: bool) {
let res = if from_queue {
self.header().state.transition_to_canceled_from_queue()
} else {
match self.header().state.transition_to_canceled_from_list() {
Some(res) => res,
None => return,
}
};
self.do_cancel(res);
}
fn do_cancel(mut self, res: Snapshot) {
use std::panic;
debug_assert!(!res.is_complete());
let cell = unsafe { &mut self.cell.as_mut() };
let header = &cell.header;
let core = &mut cell.core;
// Since we transitioned the task state to `canceled`, it won't ever be
// polled again. We are now responsible for all cleanup.
//
// We have to drop the future
//
header.future_causality.with_mut(|_| {
// Guard against potential panics in the drop handler
let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
// Drop the future
core.transition_to_consumed();
}));
});
// If there is a join waker, we must notify it so it can observe the
// task was canceled.
if res.is_join_interested() && res.has_join_waker() {
// Notify the join handle. The transition to cancelled obtained a
// lock on the waker cell.
unsafe {
self.wake_join();
}
// Also track that we might be responsible for releasing the waker.
self.set_might_drop_join_waker_on_release();
}
// The `RELEASED` flag is not set yet.
assert!(!res.is_final_ref());
// This **can** be null if the task is being cancelled before it was
// ever polled.
let bound_executor = unsafe { self.header().executor.with(|ptr| *ptr) };
unsafe {
let task = self.to_task();
if let Some(executor) = bound_executor {
executor.as_ref().release(task);
} else {
// Just drop the task. This will release / deallocate memory.
drop(task);
}
}
}
// ====== internal ======
fn complete(
mut self,
executor: NonNull<S>,
join_interest: bool,
output: super::Result<T::Output>,
) {
if join_interest {
// Store the output. The future has already been dropped
self.core().store_output(output);
}
let bound_executor = unsafe { self.header().executor.with(|ptr| *ptr) };
// Handle releasing the task. First, check if the current
// executor is the one that is bound to the task:
if Some(executor) == bound_executor {
unsafe {
// perform a local release
let task = ManuallyDrop::new(self.to_task());
executor.as_ref().release_local(&task);
if self.transition_to_released(join_interest).is_final_ref() {
self.dealloc();
}
}
} else {
let res = self.transition_to_complete(join_interest);
assert!(!res.is_final_ref());
if res.has_join_waker() {
// The release step happens later once the task has migrated back to
// the worker that owns it. At that point, the releaser **may** also
// be responsible for dropping. This fact must be tracked until
// the release step happens.
self.set_might_drop_join_waker_on_release();
}
unsafe {
let task = self.to_task();
let executor = match bound_executor {
Some(executor) => executor,
None => panic!("executor should be set"),
};
executor.as_ref().release(task);
}
}
}
/// Return `true` if the task structure should be deallocated
fn transition_to_complete(&mut self, join_interest: bool) -> Snapshot {
let res = self.header().state.transition_to_complete();
self.notify_join_handle(join_interest, res);
// Transition to complete last to ensure freeing does
// not happen until the above work is done.
res
}
/// Return `true` if the task structure should be deallocated
fn transition_to_released(&mut self, join_interest: bool) -> Snapshot {
if join_interest {
let res1 = self.transition_to_complete(join_interest);
// At this point, the join waker may not be changed. Once we perform
// `release_task` we may no longer read from the struct but we
// **may** be responsible for dropping the waker. We do an
// optimistic read here.
let (join_waker, check) = unsafe { self.read_join_waker() };
let res2 = self.header().state.release_task();
if res1.has_join_waker() && !res2.is_join_interested() {
debug_assert!(res2.has_join_waker());
// Its our responsibility to drop the waker
check.check();
unsafe {
drop(join_waker.assume_init());
}
}
res2
} else {
self.header().state.transition_to_released()
}
}
fn notify_join_handle(&mut self, join_interest: bool, res: Snapshot) {
if join_interest {
if !res.is_join_interested() {
debug_assert!(!res.has_join_waker());
// The join handle dropped interest before we could release
// the output. We are now responsible for releasing the
// output.
self.core().transition_to_consumed();
} else if res.has_join_waker() {
if res.is_canceled() {
// The join handle will set the output to Cancelled without
// attempting to read the output. We must drop it here.
self.core().transition_to_consumed();
}
// Notify the join handle. The previous transition obtains the
// lock on the waker cell.
unsafe {
self.wake_join();
}
}
}
}
fn might_drop_join_waker_on_release(&self) -> bool {
unsafe {
let next = *self.header().queue_next.get() as usize;
next & 1 == 1
}
}
fn set_might_drop_join_waker_on_release(&self) {
unsafe {
debug_assert!(
(*self.header().queue_next.get()).is_null(),
"the task's queue_next field must be null when releasing"
);
*self.header().queue_next.get() = 1 as *const _;
}
}
unsafe fn wake_join(&self) {
// LOOM: ensure we can make this call
self.trailer().waker.check();
self.trailer().waker.with_unchecked(|ptr| {
(*(*ptr).as_ptr())
.as_ref()
.expect("waker missing")
.wake_by_ref();
});
}
unsafe fn read_join_waker(&mut self) -> (MaybeUninit<Option<Waker>>, CausalCheck) {
self.trailer().waker.with_deferred(|ptr| ptr.read())
}
unsafe fn to_task(&self) -> Task<S> {
let ptr = self.cell.as_ptr() as *mut Header<S>;
Task::from_raw(NonNull::new_unchecked(ptr))
}
}
tokio/src/executor/task/join.rs
+74
View File
@@ -0,0 +1,74 @@
use crate::executor::loom::alloc::Track;
use crate::executor::task::raw::RawTask;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};
pub(crate) struct JoinHandle<T, S: 'static> {
raw: Option<RawTask<S>>,
_p: PhantomData<T>,
}
impl<T, S: 'static> JoinHandle<T, S> {
pub(super) fn new(raw: RawTask<S>) -> JoinHandle<T, S> {
JoinHandle {
raw: Some(raw),
_p: PhantomData,
}
}
}
impl<T, S: 'static> Unpin for JoinHandle<T, S> {}
impl<T, S: 'static> Future for JoinHandle<T, S> {
type Output = super::Result<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
use std::mem::MaybeUninit;
// Raw should always be set
let raw = self.raw.as_ref().unwrap();
// Load the current task state
let mut state = raw.header().state.load();
debug_assert!(state.is_join_interested());
if state.is_active() {
state = if state.has_join_waker() {
raw.swap_join_waker(cx.waker(), state)
} else {
raw.store_join_waker(cx.waker())
};
if state.is_active() {
return Poll::Pending;
}
}
let mut out = MaybeUninit::<Track<Self::Output>>::uninit();
unsafe {
// This could result in the task being freed.
raw.read_output(out.as_mut_ptr() as *mut (), state);
self.raw = None;
Poll::Ready(out.assume_init().into_inner())
}
}
}
impl<T, S: 'static> Drop for JoinHandle<T, S> {
fn drop(&mut self) {
if let Some(raw) = self.raw.take() {
if raw.header().state.drop_join_handle_fast() {
return;
}
raw.drop_join_handle_slow();
}
}
}
tokio/src/executor/task/list.rs
+70
View File
@@ -0,0 +1,70 @@
use crate::executor::task::{Header, Task};
use std::fmt;
use std::ptr::NonNull;
pub(crate) struct OwnedList<T: 'static> {
head: Option<NonNull<Header<T>>>,
}
impl<T: 'static> OwnedList<T> {
pub(crate) fn new() -> OwnedList<T> {
OwnedList { head: None }
}
pub(crate) fn insert(&mut self, task: &Task<T>) {
unsafe {
debug_assert!((*task.header().owned_next.get()).is_none());
debug_assert!((*task.header().owned_prev.get()).is_none());
let ptr = Some(task.header().into());
if let Some(next) = self.head {
debug_assert!((*next.as_ref().owned_prev.get()).is_none());
*next.as_ref().owned_prev.get() = ptr;
}
*task.header().owned_next.get() = self.head;
self.head = ptr;
}
}
pub(crate) fn remove(&mut self, task: &Task<T>) {
unsafe {
if let Some(next) = *task.header().owned_next.get() {
*next.as_ref().owned_prev.get() = *task.header().owned_prev.get();
}
if let Some(prev) = *task.header().owned_prev.get() {
*prev.as_ref().owned_next.get() = *task.header().owned_next.get();
} else {
debug_assert_eq!(self.head, Some(task.header().into()));
self.head = *task.header().owned_next.get();
}
}
}
pub(crate) fn is_empty(&self) -> bool {
self.head.is_none()
}
/// Transition all tasks in the list to canceled as part of the shutdown
/// process.
pub(crate) fn shutdown(&self) {
let mut curr = self.head;
while let Some(task) = curr {
unsafe {
let vtable = task.as_ref().vtable;
(vtable.cancel)(task.as_ptr() as *mut (), false);
curr = *task.as_ref().owned_next.get();
}
}
}
}
impl<T: 'static> fmt::Debug for OwnedList<T> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("OwnedList").finish()
}
}
tokio/src/executor/task/mod.rs
+130
View File
@@ -0,0 +1,130 @@
mod core;
pub(crate) use self::core::Header;
mod error;
pub use self::error::Error;
mod harness;
mod join;
pub(crate) use self::join::JoinHandle;
mod list;
pub(crate) use self::list::OwnedList;
mod raw;
mod stack;
pub(crate) use self::stack::TransferStack;
mod state;
mod waker;
/// Unit tests
#[cfg(test)]
mod tests;
use self::raw::RawTask;
use std::future::Future;
use std::ptr::NonNull;
use std::{fmt, mem};
/// An owned handle to the task, tracked by ref count
pub(crate) struct Task<S: 'static> {
raw: RawTask<S>,
}
unsafe impl<S: Send + Sync + 'static> Send for Task<S> {}
/// Task result sent back
pub(crate) type Result<T> = std::result::Result<T, Error>;
pub(crate) trait Schedule: Send + Sync + Sized + 'static {
/// Bind a task to the executor.
///
/// Guaranteed to be called from the thread that called `poll` on the task.
fn bind(&self, task: &Task<Self>);
/// The task has completed work and is ready to be released. The scheduler
/// is free to drop it whenever.
fn release(&self, task: Task<Self>);
/// The has been completed by the executor it was bound to.
fn release_local(&self, task: &Task<Self>);
/// Schedule the task
fn schedule(&self, task: Task<Self>);
}
/// Create a new task without an associated join handle
pub(crate) fn background<T, S>(task: T) -> Task<S>
where
T: Future + Send + 'static,
S: Schedule,
{
let raw = RawTask::new_background(task);
Task { raw }
}
/// Create a new task with an associated join handle
pub(crate) fn joinable<T, S>(task: T) -> (Task<S>, JoinHandle<T::Output, S>)
where
T: Future + Send + 'static,
S: Schedule,
{
let raw = RawTask::new_joinable(task);
let task = Task { raw };
let join = JoinHandle::new(raw);
(task, join)
}
impl<S: 'static> Task<S> {
pub(crate) unsafe fn from_raw(ptr: NonNull<Header<S>>) -> Task<S> {
let raw = RawTask::from_raw(ptr);
Task { raw }
}
pub(crate) fn header(&self) -> &Header<S> {
self.raw.header()
}
pub(crate) fn into_raw(self) -> NonNull<Header<S>> {
let raw = self.raw.into_raw();
mem::forget(self);
raw
}
}
impl<S: Schedule> Task<S> {
/// Returns `self` when the task needs to be immediately re-scheduled
pub(crate) fn run(self, executor: NonNull<S>) -> Option<Self> {
if unsafe { self.raw.poll(executor) } {
Some(self)
} else {
// Cleaning up the `Task` instance is done from within the poll
// function.
mem::forget(self);
None
}
}
/// Pre-emptively cancel the task as part of the shutdown process.
pub(crate) fn shutdown(self) {
self.raw.cancel_from_queue();
mem::forget(self);
}
}
impl<S: 'static> Drop for Task<S> {
fn drop(&mut self) {
self.raw.drop_task();
}
}
impl<S> fmt::Debug for Task<S> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Task").finish()
}
}
tokio/src/executor/task/raw.rs
+190
View File
@@ -0,0 +1,190 @@
use crate::executor::loom::alloc::Track;
use crate::executor::task::core::Cell;
use crate::executor::task::harness::Harness;
use crate::executor::task::state::{Snapshot, State};
use crate::executor::task::{Header, Schedule};
use std::future::Future;
use std::ptr::NonNull;
use std::task::Waker;
/// Raw task handle
pub(super) struct RawTask<S: 'static> {
ptr: NonNull<Header<S>>,
}
pub(super) struct Vtable<S: 'static> {
/// Poll the future
pub(super) poll: unsafe fn(*mut (), NonNull<S>) -> bool,
/// The task handle has been dropped and the join waker needs to be dropped
/// or the task struct needs to be deallocated
pub(super) drop_task: unsafe fn(*mut ()),
/// Read the task output
pub(super) read_output: unsafe fn(*mut (), *mut (), Snapshot),
/// Store the join handle's waker
///
/// Returns a snapshot of the state **after** the transition
pub(super) store_join_waker: unsafe fn(*mut (), &Waker) -> Snapshot,
/// Replace the join handle's waker
///
/// Returns a snapshot of the state **after** the transition
pub(super) swap_join_waker: unsafe fn(*mut (), &Waker, Snapshot) -> Snapshot,
/// The join handle has been dropped
pub(super) drop_join_handle_slow: unsafe fn(*mut ()),
/// The task is being canceled
pub(super) cancel: unsafe fn(*mut (), bool),
}
/// Get the vtable for the requested `T` and `S` generics.
pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable<S> {
&Vtable {
poll: poll::<T, S>,
drop_task: drop_task::<T, S>,
read_output: read_output::<T, S>,
store_join_waker: store_join_waker::<T, S>,
swap_join_waker: swap_join_waker::<T, S>,
drop_join_handle_slow: drop_join_handle_slow::<T, S>,
cancel: cancel::<T, S>,
}
}
impl<S> RawTask<S> {
pub(super) fn new_background<T>(task: T) -> RawTask<S>
where
T: Future + Send + 'static,
S: Schedule,
{
RawTask::new(task, State::new_background())
}
pub(super) fn new_joinable<T>(task: T) -> RawTask<S>
where
T: Future + Send + 'static,
S: Schedule,
{
RawTask::new(task, State::new_joinable())
}
fn new<T>(task: T, state: State) -> RawTask<S>
where
T: Future + Send + 'static,
S: Schedule,
{
let ptr = Box::into_raw(Cell::<T, S>::new(task, state));
let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header<S>) };
RawTask { ptr }
}
pub(super) unsafe fn from_raw(ptr: NonNull<Header<S>>) -> RawTask<S> {
RawTask { ptr }
}
/// Returns a reference to the task's meta structure.
///
/// Safe as `Header` is `Sync`.
pub(super) fn header(&self) -> &Header<S> {
unsafe { self.ptr.as_ref() }
}
/// Returns a raw pointer to the task's meta structure.
pub(super) fn into_raw(self) -> NonNull<Header<S>> {
self.ptr
}
/// Safety: mutual exclusion is required to call this function.
///
/// Returns `true` if the task needs to be scheduled again.
pub(super) unsafe fn poll(self, executor: NonNull<S>) -> bool {
// Get the vtable without holding a ref to the meta struct. This is done
// because a mutable reference to the task is passed into the poll fn.
let vtable = self.header().vtable;
(vtable.poll)(self.ptr.as_ptr() as *mut (), executor)
}
pub(super) fn drop_task(self) {
let vtable = self.header().vtable;
unsafe {
(vtable.drop_task)(self.ptr.as_ptr() as *mut ());
}
}
pub(super) unsafe fn read_output(self, dst: *mut (), state: Snapshot) {
let vtable = self.header().vtable;
(vtable.read_output)(self.ptr.as_ptr() as *mut (), dst, state);
}
pub(super) fn store_join_waker(self, waker: &Waker) -> Snapshot {
let vtable = self.header().vtable;
unsafe { (vtable.store_join_waker)(self.ptr.as_ptr() as *mut (), waker) }
}
pub(super) fn swap_join_waker(self, waker: &Waker, prev: Snapshot) -> Snapshot {
let vtable = self.header().vtable;
unsafe { (vtable.swap_join_waker)(self.ptr.as_ptr() as *mut (), waker, prev) }
}
pub(super) fn drop_join_handle_slow(self) {
let vtable = self.header().vtable;
unsafe { (vtable.drop_join_handle_slow)(self.ptr.as_ptr() as *mut ()) }
}
pub(super) fn cancel_from_queue(self) {
let vtable = self.header().vtable;
unsafe { (vtable.cancel)(self.ptr.as_ptr() as *mut (), true) }
}
}
impl<S: 'static> Clone for RawTask<S> {
fn clone(&self) -> Self {
RawTask { ptr: self.ptr }
}
}
impl<S: 'static> Copy for RawTask<S> {}
unsafe fn poll<T: Future, S: Schedule>(ptr: *mut (), executor: NonNull<S>) -> bool {
let harness = Harness::<T, S>::from_raw(ptr);
harness.poll(executor)
}
unsafe fn drop_task<T: Future, S: Schedule>(ptr: *mut ()) {
let harness = Harness::<T, S>::from_raw(ptr);
harness.drop_task();
}
unsafe fn read_output<T: Future, S: Schedule>(ptr: *mut (), dst: *mut (), state: Snapshot) {
let harness = Harness::<T, S>::from_raw(ptr);
harness.read_output(dst as *mut Track<super::Result<T::Output>>, state);
}
unsafe fn store_join_waker<T: Future, S: Schedule>(ptr: *mut (), waker: &Waker) -> Snapshot {
let harness = Harness::<T, S>::from_raw(ptr);
harness.store_join_waker(waker)
}
unsafe fn swap_join_waker<T: Future, S: Schedule>(
ptr: *mut (),
waker: &Waker,
prev: Snapshot,
) -> Snapshot {
let harness = Harness::<T, S>::from_raw(ptr);
harness.swap_join_waker(waker, prev)
}
unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: *mut ()) {
let harness = Harness::<T, S>::from_raw(ptr);
harness.drop_join_handle_slow()
}
unsafe fn cancel<T: Future, S: Schedule>(ptr: *mut (), from_queue: bool) {
let harness = Harness::<T, S>::from_raw(ptr);
harness.cancel(from_queue)
}
tokio/src/executor/task/stack.rs
+85
View File
@@ -0,0 +1,85 @@
use crate::executor::loom::sync::atomic::AtomicPtr;
use crate::executor::task::{Header, Task};
use std::ptr::{self, NonNull};
use std::sync::atomic::Ordering::{Acquire, Relaxed, Release};
/// Concurrent stack of tasks, used to pass ownership of a task from one worker
/// to another.
pub(crate) struct TransferStack<T: 'static> {
head: AtomicPtr<Header<T>>,
}
impl<T: 'static> TransferStack<T> {
pub(crate) fn new() -> TransferStack<T> {
TransferStack {
head: AtomicPtr::new(ptr::null_mut()),
}
}
pub(crate) fn push(&self, task: Task<T>) {
unsafe {
let task = task.into_raw();
let next = (*task.as_ref().queue_next.get()) as usize;
// At this point, the queue_next field may also be used to track
// whether or not the task must drop the join waker.
debug_assert_eq!(0, next & 1);
// We don't care about any memory associated w/ setting the `head`
// field, just the current value.
let mut curr = self.head.load(Relaxed);
loop {
*task.as_ref().queue_next.get() = (next | curr as usize) as *const _;
let res =
self.head
.compare_exchange(curr, task.as_ptr() as *mut _, Release, Relaxed);
match res {
Ok(_) => return,
Err(actual) => {
curr = actual;
}
}
}
}
}
pub(crate) fn drain(&self) -> impl Iterator<Item = Task<T>> {
struct Iter<T: 'static>(*mut Header<T>);
impl<T: 'static> Iterator for Iter<T> {
type Item = Task<T>;
fn next(&mut self) -> Option<Task<T>> {
let task = NonNull::new(self.0)?;
unsafe {
let next = *task.as_ref().queue_next.get() as usize;
// remove the data bit
self.0 = (next & !1) as *mut _;
Some(Task::from_raw(task))
}
}
}
impl<T: 'static> Drop for Iter<T> {
fn drop(&mut self) {
use std::process;
if !self.0.is_null() {
// we have bugs
process::abort();
}
}
}
let ptr = self.head.swap(ptr::null_mut(), Acquire);
Iter(ptr)
}
}
tokio/src/executor/task/state.rs
+502
View File
@@ -0,0 +1,502 @@
use crate::executor::loom::sync::atomic::AtomicUsize;
use std::fmt;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
use std::usize;
pub(super) struct State {
val: AtomicUsize,
}
/// Current state value
#[derive(Copy, Clone)]
pub(super) struct Snapshot(usize);
/// The task is currently being run.
const RUNNING: usize = 0b00_0001;
/// The task has been notified by a waker.
const NOTIFIED: usize = 0b00_0010;
/// The task is complete.
///
/// Once this bit is set, it is never unset
const COMPLETE: usize = 0b00_0100;
/// The primary task handle has been dropped.
const RELEASED: usize = 0b00_1000;
/// The join handle is still around
const JOIN_INTEREST: usize = 0b01_0000;
/// A join handle waker has been set
const JOIN_WAKER: usize = 0b10_0000;
/// The task has been forcibly canceled.
const CANCELLED: usize = 0b100_0000;
/// All bits
const LIFECYCLE_MASK: usize =
RUNNING | NOTIFIED | COMPLETE | RELEASED | JOIN_INTEREST | JOIN_WAKER | CANCELLED;
/// Bits used by the waker ref count portion of the state.
///
/// Ref counts only cover **wakers**. Other handles are tracked with other state
/// bits.
const WAKER_COUNT_MASK: usize = usize::MAX - LIFECYCLE_MASK;
/// Number of positions to shift the ref count
const WAKER_COUNT_SHIFT: usize = WAKER_COUNT_MASK.count_zeros() as usize;
/// One ref count
const WAKER_ONE: usize = 1 << WAKER_COUNT_SHIFT;
/// Initial state
const INITIAL_STATE: usize = NOTIFIED;
/// All transitions are performed via RMW operations. This establishes an
/// unambiguous modification order.
impl State {
/// Starts with a ref count of 1
pub(super) fn new_background() -> State {
State {
val: AtomicUsize::new(INITIAL_STATE),
}
}
/// Starts with a ref count of 2
pub(super) fn new_joinable() -> State {
State {
val: AtomicUsize::new(INITIAL_STATE | JOIN_INTEREST),
}
}
/// Load the current state, establishes `Acquire` ordering.
pub(super) fn load(&self) -> Snapshot {
Snapshot(self.val.load(Acquire))
}
/// Transition a task to the `Running` state.
///
/// Returns a snapshot of the state **after** the transition.
pub(super) fn transition_to_running(&self) -> Snapshot {
const DELTA: usize = RUNNING | NOTIFIED;
let prev = Snapshot(self.val.fetch_xor(DELTA, Acquire));
debug_assert!(prev.is_notified());
if prev.is_running() {
// We were signalled to cancel
//
// Apply the state
let prev = self.val.fetch_or(CANCELLED, AcqRel);
return Snapshot(prev | CANCELLED);
}
debug_assert!(!prev.is_running());
let next = Snapshot(prev.0 ^ DELTA);
debug_assert!(next.is_running());
debug_assert!(!next.is_notified());
next
}
/// Transition the task from `Running` -> `Idle`.
///
/// Returns a snapshot of the state **after** the transition.
pub(super) fn transition_to_idle(&self) -> Snapshot {
const DELTA: usize = RUNNING;
let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));
if !prev.is_running() {
// We were signaled to cancel.
//
// Apply the state
let prev = self.val.fetch_or(CANCELLED, AcqRel);
return Snapshot(prev | CANCELLED);
}
let next = Snapshot(prev.0 ^ DELTA);
debug_assert!(!next.is_running());
next
}
/// Transition the task from `Running` -> `Complete`.
///
/// Returns a snapshot of the state **after** the transition.
pub(super) fn transition_to_complete(&self) -> Snapshot {
const DELTA: usize = RUNNING | COMPLETE;
let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));
debug_assert!(!prev.is_complete());
let next = Snapshot(prev.0 ^ DELTA);
debug_assert!(next.is_complete());
next
}
/// Transition the task from `Running` -> `Released`.
///
/// Returns a snapshot of the state **after** the transition.
pub(super) fn transition_to_released(&self) -> Snapshot {
const DELTA: usize = RUNNING | COMPLETE | RELEASED;
let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel));
debug_assert!(prev.is_running());
debug_assert!(!prev.is_complete());
debug_assert!(!prev.is_released());
let next = Snapshot(prev.0 ^ DELTA);
debug_assert!(!next.is_running());
debug_assert!(next.is_complete());
debug_assert!(next.is_released());
next
}
/// Transition the task to the canceled state.
///
/// Returns the snapshot of the state **after** the transition **if** the
/// transition was made successfully
///
/// # States
///
/// - Notifed: task may be in a queue, caller must not release.
/// - Running: cannot drop. The poll handle will handle releasing.
/// - Other prior states do not require cancellation.
///
/// If the task has been notified, then it may still be in a queue. The
/// caller must not release the task.
pub(super) fn transition_to_canceled_from_queue(&self) -> Snapshot {
let prev = Snapshot(self.val.fetch_or(CANCELLED, AcqRel));
debug_assert!(!prev.is_complete());
debug_assert!(!prev.is_running() || prev.is_notified());
Snapshot(prev.0 | CANCELLED)
}
pub(super) fn transition_to_canceled_from_list(&self) -> Option<Snapshot> {
let mut prev = self.load();
loop {
if !prev.is_active() {
return None;
}
let mut next = prev;
// Use the running flag to signal cancellation
if prev.is_running() {
next.0 -= RUNNING;
} else if prev.is_notified() {
next.0 += RUNNING;
} else {
next.0 |= CANCELLED;
}
let res = self.val.compare_exchange(prev.0, next.0, AcqRel, Acquire);
match res {
Ok(_) if next.is_canceled() => return Some(next),
Ok(_) => return None,
Err(actual) => {
prev = Snapshot(actual);
}
}
}
}
/// Final transition to `Released`. Called when primary task handle is
/// dropped. This is roughly a "ref decrement" operation.
///
/// Returns a snapshot of the state **after** the transition.
pub(super) fn release_task(&self) -> Snapshot {
use crate::executor::loom::sync::atomic;
const DELTA: usize = RELEASED;
let prev = Snapshot(self.val.fetch_or(DELTA, Release));
debug_assert!(!prev.is_released());
debug_assert!(prev.is_terminal(), "state = {:?}", prev);
let next = Snapshot(prev.0 | DELTA);
debug_assert!(next.is_released());
if next.is_final_ref() || (next.has_join_waker() && !next.is_join_interested()) {
// The final reference to the task was dropped, the caller must free the
// memory. Establish an acquire ordering.
atomic::fence(Acquire);
}
next
}
/// Transition the state to `Scheduled`.
///
/// Returns `true` if the task needs to be submitted to the pool for
/// execution
pub(super) fn transition_to_notified(&self) -> bool {
const MASK: usize = RUNNING | NOTIFIED | COMPLETE | CANCELLED;
let prev = self.val.fetch_or(NOTIFIED, Release);
prev & MASK == 0
}
/// Optimistically try to swap the state assuming the join handle is
/// __immediately__ dropped on spawn
pub(super) fn drop_join_handle_fast(&self) -> bool {
use std::sync::atomic::Ordering::Relaxed;
// Relaxed is acceptable as if this function is called and succeeds,
// then nothing has been done w/ the join handle.
//
// The moment the join handle is used (polled), the `JOIN_WAKER` flag is
// set, at which point the CAS will fail.
//
// Given this, there is no risk if this operation is reordered.
self.val
.compare_exchange_weak(
INITIAL_STATE | JOIN_INTEREST,
INITIAL_STATE,
Relaxed,
Relaxed,
)
.is_ok()
}
/// The join handle has completed by reading the output
///
/// Returns a snapshot of the state **after** the transition.
pub(super) fn complete_join_handle(&self) -> Snapshot {
use crate::executor::loom::sync::atomic;
const DELTA: usize = JOIN_INTEREST;
let prev = Snapshot(self.val.fetch_sub(DELTA, Release));
debug_assert!(prev.is_join_interested());
let next = Snapshot(prev.0 - DELTA);
if !next.is_final_ref() {
return next;
}
atomic::fence(Acquire);
next
}
/// The join handle is being dropped, this fails if the task has been
/// completed and the output must be dropped first then
/// `complete_join_handle` should be called.
///
/// Returns a snapshot of the state **after** the transition.
pub(super) fn drop_join_handle_slow(&self) -> Result<Snapshot, Snapshot> {
const MASK: usize = COMPLETE | CANCELLED;
let mut prev = self.val.load(Acquire);
loop {
// Once the complete bit is set, it is never unset.
if prev & MASK != 0 {
return Err(Snapshot(prev));
}
debug_assert!(prev & JOIN_INTEREST == JOIN_INTEREST);
let next = (prev - JOIN_INTEREST) & !JOIN_WAKER;
let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);
match res {
Ok(_) => {
return Ok(Snapshot(next));
}
Err(actual) => {
prev = actual;
}
}
}
}
/// Store the join waker.
pub(super) fn store_join_waker(&self) -> Snapshot {
use crate::executor::loom::sync::atomic;
const DELTA: usize = JOIN_WAKER;
let prev = Snapshot(self.val.fetch_xor(DELTA, Release));
debug_assert!(!prev.has_join_waker());
let next = Snapshot(prev.0 ^ DELTA);
debug_assert!(next.has_join_waker());
if next.is_complete() {
atomic::fence(Acquire);
}
next
}
pub(super) fn unset_waker(&self) -> Snapshot {
const MASK: usize = COMPLETE | CANCELLED;
let mut prev = self.val.load(Acquire);
loop {
// Once the `COMPLETE` bit is set, it is never unset
if prev & MASK != 0 {
return Snapshot(prev);
}
debug_assert!(Snapshot(prev).has_join_waker());
let next = prev - JOIN_WAKER;
let res = self.val.compare_exchange(prev, next, AcqRel, Acquire);
match res {
Ok(_) => return Snapshot(next),
Err(actual) => {
prev = actual;
}
}
}
}
pub(super) fn ref_inc(&self) {
use std::process;
use std::sync::atomic::Ordering::Relaxed;
// Using a relaxed ordering is alright here, as knowledge of the
// original reference prevents other threads from erroneously deleting
// the object.
//
// As explained in the [Boost documentation][1], Increasing the
// reference counter can always be done with memory_order_relaxed: New
// references to an object can only be formed from an existing
// reference, and passing an existing reference from one thread to
// another must already provide any required synchronization.
//
// [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
let prev = self.val.fetch_add(WAKER_ONE, Relaxed);
// If the reference count overflowed, abort.
if prev > isize::max_value() as usize {
process::abort();
}
}
/// Returns `true` if the task should be released.
pub(super) fn ref_dec(&self) -> bool {
use crate::executor::loom::sync::atomic;
let prev = self.val.fetch_sub(WAKER_ONE, Release);
let next = Snapshot(prev - WAKER_ONE);
if next.is_final_ref() {
atomic::fence(Acquire);
}
next.is_final_ref()
}
}
impl Snapshot {
pub(super) fn is_running(self) -> bool {
self.0 & RUNNING == RUNNING
}
pub(super) fn is_notified(self) -> bool {
self.0 & NOTIFIED == NOTIFIED
}
pub(super) fn is_released(self) -> bool {
self.0 & RELEASED == RELEASED
}
pub(super) fn is_complete(self) -> bool {
self.0 & COMPLETE == COMPLETE
}
pub(super) fn is_canceled(self) -> bool {
self.0 & CANCELLED == CANCELLED
}
/// Used during normal runtime.
pub(super) fn is_active(self) -> bool {
self.0 & (COMPLETE | CANCELLED) == 0
}
/// Used before dropping the task
pub(super) fn is_terminal(self) -> bool {
// When both the notified & running flags are set, the task was canceled
// after being notified, before it was run.
//
// There is a race where:
// - The task state transitions to notified
// - The global queue is shutdown
// - The waker attempts to push into the global queue and fails.
// - The waker holds the last reference to the task, thus drops it.
//
// In this scenario, the cancelled bit will never get set.
!self.is_active() || (self.is_notified() && self.is_running())
}
pub(super) fn is_join_interested(self) -> bool {
self.0 & JOIN_INTEREST == JOIN_INTEREST
}
pub(super) fn has_join_waker(self) -> bool {
self.0 & JOIN_WAKER == JOIN_WAKER
}
pub(super) fn is_final_ref(self) -> bool {
const MASK: usize = WAKER_COUNT_MASK | RELEASED | JOIN_INTEREST;
(self.0 & MASK) == RELEASED
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
use std::sync::atomic::Ordering::SeqCst;
let snapshot = Snapshot(self.val.load(SeqCst));
fmt.debug_struct("State")
.field("snapshot", &snapshot)
.finish()
}
}
impl fmt::Debug for Snapshot {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Snapshot")
.field("is_running", &self.is_running())
.field("is_notified", &self.is_notified())
.field("is_released", &self.is_released())
.field("is_complete", &self.is_complete())
.field("is_canceled", &self.is_canceled())
.field("is_join_interested", &self.is_join_interested())
.field("has_join_waker", &self.has_join_waker())
.field("is_final_ref", &self.is_final_ref())
.finish()
}
}
tokio/src/executor/task/tests/loom.rs
+277
View File
@@ -0,0 +1,277 @@
use crate::executor::task;
use crate::executor::tests::loom_schedule::LoomSchedule;
use tokio_test::{assert_err, assert_ok};
use loom::future::block_on;
use loom::sync::atomic::AtomicBool;
use loom::sync::atomic::Ordering::{Acquire, Release};
use loom::thread;
use std::future::Future;
#[test]
fn create_drop_join_handle() {
loom::model(|| {
let (task, join_handle) = task::joinable(async { "hello" });
let schedule = LoomSchedule::new();
let schedule = From::from(&schedule);
let th = thread::spawn(move || {
drop(join_handle);
});
assert_none!(task.run(schedule));
th.join().unwrap();
});
}
#[test]
fn poll_drop_handle_then_drop() {
use futures_util::future::poll_fn;
use std::pin::Pin;
use std::task::Poll;
loom::model(|| {
let (task, mut join_handle) = task::joinable(async { "hello" });
let schedule = LoomSchedule::new();
let schedule = From::from(&schedule);
let th = thread::spawn(move || {
block_on(poll_fn(|cx| {
let _ = Pin::new(&mut join_handle).poll(cx);
Poll::Ready(())
}));
});
assert_none!(task.run(schedule));
th.join().unwrap();
});
}
#[test]
fn join_output() {
loom::model(|| {
let (task, join_handle) = task::joinable(async { "hello world" });
let schedule = LoomSchedule::new();
let schedule = From::from(&schedule);
let th = thread::spawn(move || {
let out = assert_ok!(block_on(join_handle));
assert_eq!("hello world", out);
});
assert_none!(task.run(schedule));
th.join().unwrap();
});
}
#[test]
fn wake_by_ref() {
loom::model(|| {
let (task, join_handle) = task::joinable(gated(2, true, false));
let schedule = LoomSchedule::new();
let schedule = &schedule;
schedule.push_task(task);
let th = join_one_task(join_handle);
work(schedule);
assert_ok!(th.join().unwrap());
});
}
#[test]
fn wake_by_val() {
loom::model(|| {
let (task, join_handle) = task::joinable(gated(2, true, true));
let schedule = LoomSchedule::new();
let schedule = &schedule;
schedule.push_task(task);
let th = join_one_task(join_handle);
work(schedule);
assert_ok!(th.join().unwrap());
});
}
#[test]
fn release_remote() {
loom::model(|| {
let (task, join_handle) = task::joinable(gated(1, false, true));
let s1 = LoomSchedule::new();
let s2 = LoomSchedule::new();
// Join handle
let th = join_one_task(join_handle);
let task = match task.run(From::from(&s1)) {
Some(task) => task,
None => s1.recv().expect("released!"),
};
assert_none!(task.run(From::from(&s2)));
assert_none!(s1.recv());
assert_ok!(th.join().unwrap());
});
}
#[test]
fn shutdown_task_before_poll() {
loom::model(|| {
let (task, join_handle) = task::joinable::<_, LoomSchedule>(async { "hello" });
let th = join_one_task(join_handle);
task.shutdown();
assert_err!(th.join().unwrap());
});
}
#[test]
fn shutdown_from_list_after_poll() {
loom::model(|| {
let (task, join_handle) = task::joinable(gated(1, false, false));
let s1 = LoomSchedule::new();
let mut list = task::OwnedList::new();
list.insert(&task);
// Join handle
let th = join_two_tasks(join_handle);
match task.run(From::from(&s1)) {
Some(task) => {
// always drain the list before calling shutdown on tasks
list.shutdown();
// The task was scheduled, drain it explicitly.
task.shutdown();
}
None => {
list.shutdown();
}
};
match s1.recv() {
Some(task) => task.shutdown(),
None => {}
}
assert_err!(th.join().unwrap());
});
}
#[test]
fn shutdown_from_queue_after_poll() {
loom::model(|| {
let (task, join_handle) = task::joinable(gated(1, false, false));
let s1 = LoomSchedule::new();
// Join handle
let th = join_two_tasks(join_handle);
let task = match task.run(From::from(&s1)) {
Some(task) => task,
None => assert_some!(s1.recv()),
};
task.shutdown();
assert_err!(th.join().unwrap());
});
}
fn gated(n: usize, complete_first_poll: bool, by_val: bool) -> impl Future<Output = &'static str> {
use futures_util::future::poll_fn;
use std::sync::Arc;
use std::task::Poll;
let gate = Arc::new(AtomicBool::new(false));
let mut fired = false;
poll_fn(move |cx| {
if !fired {
for _ in 0..n {
let gate = gate.clone();
let waker = cx.waker().clone();
thread::spawn(move || {
gate.store(true, Release);
if by_val {
waker.wake()
} else {
waker.wake_by_ref();
}
});
}
fired = true;
if !complete_first_poll {
return Poll::Pending;
}
}
if gate.load(Acquire) {
Poll::Ready("hello world")
} else {
Poll::Pending
}
})
}
fn work(schedule: &LoomSchedule) {
while let Some(task) = schedule.recv() {
let mut task = Some(task);
while let Some(t) = task.take() {
task = t.run(From::from(schedule));
}
}
}
/// Spawn a thread to wait on the join handle. Uses a single task.
fn join_one_task<T: Future + 'static>(join_handle: T) -> loom::thread::JoinHandle<T::Output> {
thread::spawn(move || block_on(join_handle))
}
/// Spawn a thread to wait on the join handle using two tasks. First, poll the
/// join handle on the first task. If the join handle is not ready, then use a
/// second task to wait on it.
fn join_two_tasks<T: Future + Unpin + 'static>(
join_handle: T,
) -> loom::thread::JoinHandle<T::Output> {
use futures_util::future::poll_fn;
use std::task::Poll;
// Join handle
thread::spawn(move || {
let mut join_handle = Some(join_handle);
block_on(poll_fn(move |cx| {
use std::pin::Pin;
let res = Pin::new(join_handle.as_mut().unwrap()).poll(cx);
if res.is_ready() {
return res;
}
// Yes, we are nesting
Poll::Ready(block_on(join_handle.take().unwrap()))
}))
})
}
tokio/src/executor/task/tests/mod.rs
+5
View File
@@ -0,0 +1,5 @@
#[cfg(loom)]
mod loom;
#[cfg(not(loom))]
mod task;
tokio/src/executor/task/tests/task.rs
+643
View File
@@ -0,0 +1,643 @@
use crate::executor::task::{self, Header};
use crate::executor::tests::backoff::*;
use crate::executor::tests::mock_schedule::{mock, Mock};
use crate::executor::tests::track_drop::track_drop;
use crate::sync::oneshot;
use tokio_test::task::spawn;
use tokio_test::{assert_pending, assert_ready_err, assert_ready_ok};
use futures_util::future::poll_fn;
use std::sync::mpsc;
#[test]
fn header_lte_cache_line() {
use std::mem::size_of;
assert!(size_of::<Header<()>>() <= 8 * size_of::<*const ()>());
}
#[test]
fn create_complete_drop() {
let (tx, rx) = mpsc::channel();
let (task, did_drop) = track_drop(async move {
tx.send(1).unwrap();
});
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
// Nothing is returned
assert!(task.run(mock).is_none());
// The message was sent
assert!(rx.try_recv().is_ok());
// The future & output were dropped.
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
fn create_yield_complete_drop() {
let (tx, rx) = mpsc::channel();
let (task, did_drop) = track_drop(async move {
backoff(1).await;
tx.send(1).unwrap();
});
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
// Task is returned
let task = assert_some!(task.run(mock));
// The future was **not** dropped.
assert!(!did_drop.did_drop_future());
assert_none!(task.run(mock));
// The message was sent
assert!(rx.try_recv().is_ok());
// The future was dropped.
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
fn create_clone_yield_complete_drop() {
let (tx, rx) = mpsc::channel();
let (task, did_drop) = track_drop(async move {
backoff_clone(1).await;
tx.send(1).unwrap();
});
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
// Task is returned
let task = assert_some!(task.run(mock));
// The future was **not** dropped.
assert!(!did_drop.did_drop_future());
assert_none!(task.run(mock));
// The message was sent
assert!(rx.try_recv().is_ok());
// The future was dropped.
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
fn create_wake_drop() {
let (tx, rx) = oneshot::channel();
let (task, did_drop) = track_drop(async move { rx.await });
let task = task::background(task);
let mock = mock().bind(&task).schedule().release_local();
assert_none!(task.run(From::from(&mock)));
assert_none!(mock.next_pending_run());
// The future was **not** dropped.
assert!(!did_drop.did_drop_future());
tx.send("hello").unwrap();
let task = assert_some!(mock.next_pending_run());
assert_none!(task.run(From::from(&mock)));
// The future was dropped.
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
fn notify_complete() {
use std::task::Poll::Ready;
let (task, did_drop) = track_drop(async move {
poll_fn(|cx| {
cx.waker().wake_by_ref();
Ready(())
})
.await;
});
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
assert_none!(task.run(mock));
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
fn complete_on_second_schedule_obj() {
let (tx, rx) = mpsc::channel();
let (task, did_drop) = track_drop(async move {
backoff(1).await;
tx.send(1).unwrap();
});
let task = task::background(task);
let mock1 = mock();
let mock2 = mock().bind(&task).release();
// Task is returned
let task = assert_some!(task.run(From::from(&mock2)));
assert_none!(task.run(From::from(&mock1)));
// The message was sent
assert!(rx.try_recv().is_ok());
// The future was dropped.
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
let _ = assert_some!(mock2.next_pending_drop());
}
#[test]
fn join_task_immediate_drop_handle() {
let (task, did_drop) = track_drop(async move { "hello".to_string() });
let (task, _) = task::joinable(task);
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
fn join_task_immediate_complete_1() {
let (task, did_drop) = track_drop(async move { "hello".to_string() });
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
assert!(!handle.is_woken());
let out = assert_ready_ok!(handle.poll());
assert_eq!(out.get_ref(), "hello");
drop(out);
assert!(did_drop.did_drop_output());
}
#[test]
fn join_task_immediate_complete_2() {
let (task, did_drop) = track_drop(async move { "hello".to_string() });
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
let mock = mock().bind(&task).release_local();
assert_pending!(handle.poll());
assert!(task.run(From::from(&mock)).is_none());
assert!(did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
assert!(handle.is_woken());
let out = assert_ready_ok!(handle.poll());
assert_eq!(out.get_ref(), "hello");
drop(out);
assert!(did_drop.did_drop_output());
}
#[test]
fn join_task_complete_later() {
let (task, did_drop) = track_drop(async move {
backoff(1).await;
"hello".to_string()
});
let (task, handle) = task::joinable(task);
let mut handle = spawn(async { handle.await });
let mock = mock().bind(&task).release_local();
let task = assert_some!(task.run(From::from(&mock)));
assert!(!did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock)));
assert!(handle.is_woken());
let out = assert_ready_ok!(handle.poll());
assert_eq!(out.get_ref(), "hello");
drop(out);
assert!(did_drop.did_drop_output());
assert_eq!(1, handle.waker_ref_count());
}
#[test]
fn drop_join_after_poll() {
let (task, did_drop) = track_drop(async move {
backoff(1).await;
"hello".to_string()
});
let (task, handle) = task::joinable(task);
let mut handle = spawn(async { handle.await });
let mock = mock().bind(&task).release_local();
assert_pending!(handle.poll());
drop(handle);
let task = assert_some!(task.run(From::from(&mock)));
assert!(!did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
assert_none!(task.run(From::from(&mock)));
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
fn join_handle_change_task_complete() {
use std::future::Future;
use std::pin::Pin;
let (task, did_drop) = track_drop(async move {
backoff(1).await;
"hello".to_string()
});
let (task, mut handle) = task::joinable(task);
let mut t1 = spawn(poll_fn(|cx| Pin::new(&mut handle).poll(cx)));
let mock = mock().bind(&task).release_local();
assert_pending!(t1.poll());
drop(t1);
let task = assert_some!(task.run(From::from(&mock)));
let mut t2 = spawn(poll_fn(|cx| Pin::new(&mut handle).poll(cx)));
assert_pending!(t2.poll());
assert!(!did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
assert_none!(task.run(From::from(&mock)));
assert!(t2.is_woken());
let out = assert_ready_ok!(t2.poll());
assert_eq!(out.get_ref(), "hello");
drop(out);
assert!(did_drop.did_drop_output());
assert_eq!(1, t2.waker_ref_count());
}
#[test]
fn drop_handle_after_complete() {
let (task, did_drop) = track_drop(async move { "hello".to_string() });
let (task, handle) = task::joinable(task);
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
drop(handle);
assert!(did_drop.did_drop_output());
}
#[test]
fn non_initial_task_state_drop_join_handle_without_polling() {
let (tx, rx) = oneshot::channel::<()>();
let (task, did_drop) = track_drop(async move {
rx.await.unwrap();
"hello".to_string()
});
let (task, handle) = task::joinable(task);
let mock = mock().bind(&task).schedule().release_local();
assert_none!(task.run(From::from(&mock)));
drop(handle);
assert!(!did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
tx.send(()).unwrap();
let task = assert_some!(mock.next_pending_run());
assert!(task.run(From::from(&mock)).is_none());
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
}
#[test]
#[cfg(not(miri))]
fn task_panic_background() {
let (task, did_drop) = track_drop(async move {
if true {
panic!()
}
"hello"
});
let task = task::background(task);
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(did_drop.did_drop_future());
}
#[test]
#[cfg(not(miri))]
fn task_panic_join() {
let (task, did_drop) = track_drop(async move {
if true {
panic!()
}
"hello"
});
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
let mock = mock().bind(&task).release_local();
assert_pending!(handle.poll());
assert!(task.run(From::from(&mock)).is_none());
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
assert_ready_err!(handle.poll());
}
#[test]
fn complete_second_schedule_obj_before_join() {
let (tx, rx) = oneshot::channel();
let (task, did_drop) = track_drop(async move { rx.await.unwrap() });
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
let mock1 = mock();
let mock2 = mock().bind(&task).schedule().release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock2)));
tx.send("hello").unwrap();
let task = assert_some!(mock2.next_pending_run());
assert_none!(task.run(From::from(&mock1)));
assert!(did_drop.did_drop_future());
// The join handle was notified
assert!(handle.is_woken());
// Drop the task
let _ = assert_some!(mock2.next_pending_drop());
// Get the output
let out = assert_ready_ok!(handle.poll());
assert_eq!(*out.get_ref(), "hello");
}
#[test]
fn complete_second_schedule_obj_after_join() {
let (tx, rx) = oneshot::channel();
let (task, did_drop) = track_drop(async move { rx.await.unwrap() });
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
let mock1 = mock();
let mock2 = mock().bind(&task).schedule().release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock2)));
tx.send("hello").unwrap();
let task = assert_some!(mock2.next_pending_run());
assert_none!(task.run(From::from(&mock1)));
assert!(did_drop.did_drop_future());
// The join handle was notified
assert!(handle.is_woken());
// Get the output
let out = assert_ready_ok!(handle.poll());
assert_eq!(*out.get_ref(), "hello");
// Drop the task
let _ = assert_some!(mock2.next_pending_drop());
assert_eq!(1, handle.waker_ref_count());
}
#[test]
fn shutdown_from_list_before_notified() {
let (tx, rx) = oneshot::channel::<()>();
let mut list = task::OwnedList::new();
let (task, did_drop) = track_drop(async move { rx.await });
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
list.insert(&task);
let mock = mock().bind(&task).release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock)));
list.shutdown();
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
let task = assert_some!(mock.next_pending_drop());
drop(task);
assert_ready_err!(handle.poll());
drop(tx);
}
#[test]
fn shutdown_from_list_after_notified() {
let (tx, rx) = oneshot::channel::<()>();
let mut list = task::OwnedList::new();
let (task, did_drop) = track_drop(async move { rx.await });
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
list.insert(&task);
let mock = mock().bind(&task).schedule().release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock)));
tx.send(()).unwrap();
let task = assert_some!(mock.next_pending_run());
list.shutdown();
assert_none!(mock.next_pending_drop());
assert_none!(task.run(From::from(&mock)));
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
let task = assert_some!(mock.next_pending_drop());
drop(task);
assert_ready_err!(handle.poll());
}
#[test]
fn shutdown_from_list_after_complete() {
let mut list = task::OwnedList::new();
let (task, did_drop) = track_drop(async move {
backoff(1).await;
"hello"
});
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
list.insert(&task);
let m1 = mock().bind(&task).release();
let m2 = mock();
assert_pending!(handle.poll());
let task = assert_some!(task.run(From::from(&m1)));
assert_none!(task.run(From::from(&m2)));
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
list.shutdown();
let task = assert_some!(m1.next_pending_drop());
drop(task);
let out = assert_ready_ok!(handle.poll());
assert_eq!(*out.get_ref(), "hello");
}
#[test]
fn shutdown_from_task_before_notified() {
let (tx, rx) = oneshot::channel::<()>();
let (task, did_drop) = track_drop(async move { rx.await });
let (task, handle) = task::joinable::<_, Mock>(task);
let mut handle = spawn(handle);
assert_pending!(handle.poll());
task.shutdown();
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
assert_ready_err!(handle.poll());
drop(tx);
}
#[test]
fn shutdown_from_task_after_notified() {
let (tx, rx) = oneshot::channel::<()>();
let (task, did_drop) = track_drop(async move { rx.await });
let (task, handle) = task::joinable(task);
let mut handle = spawn(handle);
let mock = mock().bind(&task).schedule().release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock)));
tx.send(()).unwrap();
let task = assert_some!(mock.next_pending_run());
task.shutdown();
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
let task = assert_some!(mock.next_pending_drop());
drop(task);
assert_ready_err!(handle.poll());
}
tokio/src/executor/task/waker.rs
+107
View File
@@ -0,0 +1,107 @@
use crate::executor::task::harness::Harness;
use crate::executor::task::{Header, Schedule};
use std::future::Future;
use std::marker::PhantomData;
use std::ops;
use std::task::{RawWaker, RawWakerVTable, Waker};
pub(super) struct WakerRef<'a, S: 'static> {
waker: Waker,
_p: PhantomData<&'a Header<S>>,
}
/// Returns a `WakerRef` which avoids having to pre-emptively increase the
/// refcount if there is no need to do so.
pub(super) fn waker_ref<T, S>(meta: &Header<S>) -> WakerRef<'_, S>
where
T: Future,
S: Schedule,
{
let ptr = meta as *const _ as *const ();
let vtable = &RawWakerVTable::new(
clone_waker::<T, S>,
wake_unreachable,
wake_by_local_ref::<T, S>,
noop,
);
let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, vtable)) };
WakerRef {
waker,
_p: PhantomData,
}
}
impl<S> ops::Deref for WakerRef<'_, S> {
type Target = Waker;
fn deref(&self) -> &Waker {
&self.waker
}
}
unsafe fn clone_waker<T, S>(ptr: *const ()) -> RawWaker
where
T: Future,
S: Schedule,
{
let meta = ptr as *const Header<S>;
(*meta).state.ref_inc();
let vtable = &RawWakerVTable::new(
clone_waker::<T, S>,
wake_by_val::<T, S>,
wake_by_ref::<T, S>,
drop_waker::<T, S>,
);
RawWaker::new(ptr, vtable)
}
unsafe fn drop_waker<T, S>(ptr: *const ())
where
T: Future,
S: Schedule,
{
let harness = Harness::<T, S>::from_raw(ptr as *mut _);
harness.drop_waker();
}
// `wake()` cannot be called on the ref variaant.
unsafe fn wake_unreachable(_data: *const ()) {
unreachable!();
}
unsafe fn wake_by_val<T, S>(ptr: *const ())
where
T: Future,
S: Schedule,
{
let harness = Harness::<T, S>::from_raw(ptr as *mut _);
harness.wake_by_val();
}
// This function can only be called when on the runtime.
unsafe fn wake_by_local_ref<T, S>(ptr: *const ())
where
T: Future,
S: Schedule,
{
let harness = Harness::<T, S>::from_raw(ptr as *mut _);
harness.wake_by_local_ref();
}
// Wake without consuming the waker
unsafe fn wake_by_ref<T, S>(ptr: *const ())
where
T: Future,
S: Schedule,
{
let harness = Harness::<T, S>::from_raw(ptr as *mut _);
harness.wake_by_ref();
}
unsafe fn noop(_ptr: *const ()) {}
tokio/src/executor/tests/backoff.rs
+32
View File
@@ -0,0 +1,32 @@
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
pub(crate) struct Backoff(usize, bool);
pub(crate) fn backoff(n: usize) -> impl Future<Output = ()> {
Backoff(n, false)
}
/// Back off, but clone the waker each time
pub(crate) fn backoff_clone(n: usize) -> impl Future<Output = ()> {
Backoff(n, true)
}
impl Future for Backoff {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if self.0 == 0 {
return Poll::Ready(());
}
self.0 -= 1;
if self.1 {
cx.waker().clone().wake();
} else {
cx.waker().wake_by_ref();
}
Poll::Pending
}
}
tokio/src/executor/tests/loom_oneshot.rs
+49
View File
@@ -0,0 +1,49 @@
use loom::sync::Notify;
use std::sync::{Arc, Mutex};
pub(crate) fn channel<T>() -> (Sender<T>, Receiver<T>) {
let inner = Arc::new(Inner {
notify: Notify::new(),
value: Mutex::new(None),
});
let tx = Sender {
inner: inner.clone(),
};
let rx = Receiver { inner };
(tx, rx)
}
pub(crate) struct Sender<T> {
inner: Arc<Inner<T>>,
}
pub(crate) struct Receiver<T> {
inner: Arc<Inner<T>>,
}
struct Inner<T> {
notify: Notify,
value: Mutex<Option<T>>,
}
impl<T> Sender<T> {
pub(crate) fn send(self, value: T) {
*self.inner.value.lock().unwrap() = Some(value);
self.inner.notify.notify();
}
}
impl<T> Receiver<T> {
pub(crate) fn recv(self) -> T {
loop {
if let Some(v) = self.inner.value.lock().unwrap().take() {
return v;
}
self.inner.notify.wait();
}
}
}
tokio/src/executor/tests/loom_schedule.rs
+51
View File
@@ -0,0 +1,51 @@
use crate::executor::task::{Schedule, Task};
use loom::sync::Notify;
use std::collections::VecDeque;
use std::sync::Mutex;
pub(crate) struct LoomSchedule {
notify: Notify,
pending: Mutex<VecDeque<Option<Task<Self>>>>,
}
impl LoomSchedule {
pub(crate) fn new() -> LoomSchedule {
LoomSchedule {
notify: Notify::new(),
pending: Mutex::new(VecDeque::new()),
}
}
pub(crate) fn push_task(&self, task: Task<Self>) {
self.schedule(task);
}
pub(crate) fn recv(&self) -> Option<Task<Self>> {
loop {
if let Some(task) = self.pending.lock().unwrap().pop_front() {
return task;
}
self.notify.wait();
}
}
}
impl Schedule for LoomSchedule {
fn bind(&self, _task: &Task<Self>) {}
fn release(&self, task: Task<Self>) {
self.release_local(&task);
}
fn release_local(&self, _task: &Task<Self>) {
self.pending.lock().unwrap().push_back(None);
self.notify.notify();
}
fn schedule(&self, task: Task<Self>) {
self.pending.lock().unwrap().push_back(Some(task));
self.notify.notify();
}
}
tokio/src/executor/tests/mock_park.rs
+66
View File
@@ -0,0 +1,66 @@
#![allow(warnings)]
use crate::executor::park::{Park, Unpark};
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering::SeqCst};
use std::sync::Arc;
use std::time::Duration;
pub struct MockPark {
parks: HashMap<usize, Arc<Inner>>,
}
#[derive(Clone)]
struct ParkImpl(Arc<Inner>);
struct Inner {
unparked: AtomicBool,
}
impl MockPark {
pub fn new() -> MockPark {
MockPark {
parks: HashMap::new(),
}
}
pub fn is_unparked(&self, index: usize) -> bool {
self.parks[&index].unparked.load(SeqCst)
}
pub fn clear(&self, index: usize) {
self.parks[&index].unparked.store(false, SeqCst);
}
pub fn mk_park(&mut self, index: usize) -> impl Park {
let inner = Arc::new(Inner {
unparked: AtomicBool::new(false),
});
self.parks.insert(index, inner.clone());
ParkImpl(inner)
}
}
impl Park for ParkImpl {
type Unpark = ParkImpl;
type Error = ();
fn unpark(&self) -> Self::Unpark {
self.clone()
}
fn park(&mut self) -> Result<(), Self::Error> {
unimplemented!();
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
unimplemented!();
}
}
impl Unpark for ParkImpl {
fn unpark(&self) {
self.0.unparked.store(true, SeqCst);
}
}
tokio/src/executor/tests/mock_schedule.rs
+131
View File
@@ -0,0 +1,131 @@
#![allow(warnings)]
use crate::executor::task::{Header, Schedule, Task};
use std::collections::VecDeque;
use std::sync::Mutex;
use std::thread;
pub(crate) struct Mock {
inner: Mutex<Inner>,
}
pub(crate) struct Noop;
pub(crate) static NOOP_SCHEDULE: Noop = Noop;
struct Inner {
calls: VecDeque<Call>,
pending_run: VecDeque<Task<Mock>>,
pending_drop: VecDeque<Task<Mock>>,
}
unsafe impl Send for Inner {}
unsafe impl Sync for Inner {}
#[derive(Debug, Eq, PartialEq)]
enum Call {
Bind(*const Header<Mock>),
Release,
ReleaseLocal,
Schedule,
}
pub(crate) fn mock() -> Mock {
Mock {
inner: Mutex::new(Inner {
calls: VecDeque::new(),
pending_run: VecDeque::new(),
pending_drop: VecDeque::new(),
}),
}
}
impl Mock {
pub(crate) fn bind(self, task: &Task<Mock>) -> Self {
self.push(Call::Bind(task.header() as *const _));
self
}
pub(crate) fn release(self) -> Self {
self.push(Call::Release);
self
}
pub(crate) fn release_local(self) -> Self {
self.push(Call::ReleaseLocal);
self
}
pub(crate) fn schedule(self) -> Self {
self.push(Call::Schedule);
self
}
pub(crate) fn next_pending_run(&self) -> Option<Task<Self>> {
self.inner.lock().unwrap().pending_run.pop_front()
}
pub(crate) fn next_pending_drop(&self) -> Option<Task<Self>> {
self.inner.lock().unwrap().pending_drop.pop_front()
}
fn push(&self, call: Call) {
self.inner.lock().unwrap().calls.push_back(call);
}
fn next(&self, name: &str) -> Call {
self.inner
.lock()
.unwrap()
.calls
.pop_front()
.expect(&format!("received `{}`, but none expected", name))
}
}
impl Schedule for Mock {
fn bind(&self, task: &Task<Self>) {
match self.next("bind") {
Call::Bind(ptr) => {
assert!(ptr.eq(&(task.header() as *const _)));
}
call => panic!("expected `Bind`, was {:?}", call),
}
}
fn release(&self, task: Task<Self>) {
match self.next("release") {
Call::Release => {
self.inner.lock().unwrap().pending_drop.push_back(task);
}
call => panic!("expected `Release`, was {:?}", call),
}
}
fn release_local(&self, _task: &Task<Self>) {
assert_eq!(Call::ReleaseLocal, self.next("release_local"));
}
fn schedule(&self, task: Task<Self>) {
self.inner.lock().unwrap().pending_run.push_back(task);
assert_eq!(Call::Schedule, self.next("schedule"));
}
}
impl Drop for Mock {
fn drop(&mut self) {
if !thread::panicking() {
assert!(self.inner.lock().unwrap().calls.is_empty());
}
}
}
impl Schedule for Noop {
fn bind(&self, _task: &Task<Self>) {}
fn release(&self, _task: Task<Self>) {}
fn release_local(&self, _task: &Task<Self>) {}
fn schedule(&self, _task: Task<Self>) {}
}
tokio/src/executor/tests/mod.rs
+40
View File
@@ -0,0 +1,40 @@
//! Testing utilities
#[cfg(not(loom))]
pub(crate) mod backoff;
#[cfg(loom)]
pub(crate) mod loom_oneshot;
#[cfg(loom)]
pub(crate) mod loom_schedule;
#[cfg(not(loom))]
pub(crate) mod mock_park;
pub(crate) mod mock_schedule;
#[cfg(not(loom))]
pub(crate) mod track_drop;
/// Panic if expression results in `None`.
#[macro_export]
macro_rules! assert_some {
($e:expr) => {{
match $e {
Some(v) => v,
_ => panic!("expected some, was none"),
}
}};
}
/// Panic if expression results in `Some`.
#[macro_export]
macro_rules! assert_none {
($e:expr) => {{
match $e {
Some(v) => panic!("expected none, was {:?}", v),
_ => {}
}
}};
}
tokio/src/executor/tests/track_drop.rs
+57
View File
@@ -0,0 +1,57 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::Arc;
use std::task::{Context, Poll};
#[derive(Debug)]
pub(crate) struct TrackDrop<T>(T, Arc<AtomicBool>);
#[derive(Debug)]
pub(crate) struct DidDrop(Arc<AtomicBool>, Arc<AtomicBool>);
pub(crate) fn track_drop<T: Future>(
future: T,
) -> (impl Future<Output = TrackDrop<T::Output>>, DidDrop) {
let did_drop_future = Arc::new(AtomicBool::new(false));
let did_drop_output = Arc::new(AtomicBool::new(false));
let did_drop = DidDrop(did_drop_future.clone(), did_drop_output.clone());
let future = async move { TrackDrop(future.await, did_drop_output) };
let future = TrackDrop(future, did_drop_future);
(future, did_drop)
}
impl<T> TrackDrop<T> {
pub(crate) fn get_ref(&self) -> &T {
&self.0
}
}
impl<T: Future> Future for TrackDrop<T> {
type Output = T::Output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = unsafe { Pin::map_unchecked_mut(self, |x| &mut x.0) };
me.poll(cx)
}
}
impl<T> Drop for TrackDrop<T> {
fn drop(&mut self) {
self.1.store(true, SeqCst);
}
}
impl DidDrop {
pub(crate) fn did_drop_future(&self) -> bool {
self.0.load(SeqCst)
}
pub(crate) fn did_drop_output(&self) -> bool {
self.1.load(SeqCst)
}
}
tokio/src/executor/thread_pool/builder.rs
+259
View File
@@ -0,0 +1,259 @@
use crate::executor::loom::sync::Arc;
use crate::executor::loom::sys::num_cpus;
use crate::executor::loom::thread;
use crate::executor::park::Park;
use crate::executor::thread_pool::park::DefaultPark;
use crate::executor::thread_pool::{shutdown, worker, worker::Worker, Spawner, ThreadPool};
use std::{fmt, usize};
/// Builds a thread pool with custom configuration values.
pub struct Builder {
/// Number of worker threads to spawn
pool_size: usize,
/// Thread name
name: String,
/// Thread stack size
stack_size: Option<usize>,
/// Around worker callback
around_worker: Option<Callback>,
}
// The Arc<Box<_>> is needed because loom doesn't support Arc<T> where T: !Sized
// loom doesn't support that because it requires CoerceUnsized, which is unstable
type Callback = Arc<Box<dyn Fn(usize, &mut dyn FnMut()) + Send + Sync>>;
impl Builder {
/// Returns a new thread pool builder initialized with default configuration
/// values.
pub fn new() -> Builder {
Builder {
pool_size: num_cpus(),
name: "tokio-runtime-worker".to_string(),
stack_size: None,
around_worker: None,
}
}
/// Set the number of threads running async tasks.
///
/// This must be a number between 1 and 2,048 though it is advised to keep
/// this value on the smaller side.
///
/// The default value is the number of cores available to the system.
///
/// # Examples
///
/// ```
/// use tokio::executor::thread_pool::Builder;
///
/// let thread_pool = Builder::new()
/// .num_threads(4)
/// .build();
/// ```
pub fn num_threads(&mut self, value: usize) -> &mut Self {
self.pool_size = value;
self
}
/// Set name of threads spawned by the scheduler
///
/// If this configuration is not set, then the thread will use the system
/// default naming scheme.
///
/// # Examples
///
/// ```
/// use tokio::executor::thread_pool::Builder;
///
/// let thread_pool = Builder::new()
/// .name("my-pool")
/// .build();
/// ```
pub fn name<S: Into<String>>(&mut self, val: S) -> &mut Self {
self.name = val.into();
self
}
/// Set the stack size (in bytes) for worker threads.
///
/// The actual stack size may be greater than this value if the platform
/// specifies minimal stack size.
///
/// The default stack size for spawned threads is 2 MiB, though this
/// particular stack size is subject to change in the future.
///
/// # Examples
///
/// ```
/// use tokio::executor::thread_pool::Builder;
///
/// let thread_pool = Builder::new()
/// .stack_size(32 * 1024)
/// .build();
/// ```
pub fn stack_size(&mut self, val: usize) -> &mut Self {
self.stack_size = Some(val);
self
}
/// Execute function `f` on each worker thread.
///
/// This function is provided a function that executes the worker and is
/// expected to call it, otherwise the worker thread will shutdown without
/// doing any work.
///
/// # Examples
///
/// ```
/// use tokio::executor::thread_pool::Builder;
///
/// let thread_pool = Builder::new()
/// .around_worker(|index, work| {
/// println!("worker {} is starting up", index);
/// work();
/// println!("worker {} is shutting down", index);
/// })
/// .build();
/// ```
pub fn around_worker<F>(&mut self, f: F) -> &mut Self
where
F: Fn(usize, &mut dyn FnMut()) + Send + Sync + 'static,
{
self.around_worker = Some(Arc::new(Box::new(f)));
self
}
/// Create the configured `ThreadPool`.
///
/// The returned `ThreadPool` instance is ready to spawn tasks.
///
/// # Examples
///
/// ```
/// use tokio::executor::thread_pool::Builder;
///
/// let thread_pool = Builder::new()
/// .build();
/// ```
pub fn build(&self) -> ThreadPool {
self.build_with_park(|_| DefaultPark::new())
}
/// Create the configured `ThreadPool` with a custom `park` instances.
///
/// The provided closure `build_park` is called once per worker and returns
/// a `Park` instance that is used by the worker to put itself to sleep.
pub fn build_with_park<F, P>(&self, mut build_park: F) -> ThreadPool
where
F: FnMut(usize) -> P,
P: Park + Send + 'static,
{
let (shutdown_tx, shutdown_rx) = shutdown::channel();
let around_worker = self.around_worker.as_ref().map(Arc::clone);
let launch_worker = move |worker: Worker<BoxedPark<P>>| {
let shutdown_tx = shutdown_tx.clone();
let around_worker = around_worker.as_ref().map(Arc::clone);
Box::new(move || {
struct AbortOnPanic;
impl Drop for AbortOnPanic {
fn drop(&mut self) {
if thread::panicking() {
eprintln!("[ERROR] unhandled panic in Tokio scheduler. This is a bug and should be reported.");
std::process::abort();
}
}
}
let _abort_on_panic = AbortOnPanic;
if let Some(cb) = around_worker.as_ref() {
let idx = worker.id();
let mut f = Some(move || worker.run());
cb(idx, &mut || {
(f.take()
.expect("around_thread callback called closure twice"))(
)
})
} else {
worker.run()
}
// Dropping the handle must happen __after__ the callback
drop(shutdown_tx);
}) as Box<dyn FnOnce() + Send + 'static>
};
let mut blocking = crate::executor::blocking::Builder::default();
blocking.name(self.name.clone());
if let Some(ss) = self.stack_size {
blocking.stack_size(ss);
}
let blocking = Arc::new(blocking.build());
let (pool, workers) = worker::create_set::<_, BoxedPark<P>>(
self.pool_size,
|i| BoxedPark::new(build_park(i)),
blocking.clone(),
);
// Spawn threads for each worker
for worker in workers {
crate::executor::blocking::Pool::spawn(&blocking, launch_worker(worker))
}
let spawner = Spawner::new(pool);
let blocking = crate::executor::blocking::PoolWaiter::from(blocking);
ThreadPool::from_parts(spawner, shutdown_rx, blocking)
}
}
impl Default for Builder {
fn default() -> Builder {
Builder::new()
}
}
impl fmt::Debug for Builder {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Builder")
.field("pool_size", &self.pool_size)
.field("name", &self.name)
.field("stack_size", &self.stack_size)
.finish()
}
}
pub(crate) struct BoxedPark<P> {
inner: P,
}
impl<P> BoxedPark<P> {
pub(crate) fn new(inner: P) -> Self {
BoxedPark { inner }
}
}
impl<P> Park for BoxedPark<P>
where
P: Park,
{
type Unpark = Box<dyn crate::executor::park::Unpark>;
type Error = P::Error;
fn unpark(&self) -> Self::Unpark {
Box::new(self.inner.unpark())
}
fn park(&mut self) -> Result<(), Self::Error> {
self.inner.park()
}
fn park_timeout(&mut self, duration: std::time::Duration) -> Result<(), Self::Error> {
self.inner.park_timeout(duration)
}
}
tokio/src/executor/thread_pool/current.rs
+85
View File
@@ -0,0 +1,85 @@
use crate::executor::loom::sync::Arc;
use crate::executor::park::Unpark;
use crate::executor::thread_pool::{worker, Owned};
use std::cell::Cell;
use std::ptr;
/// Tracks the current worker
#[derive(Debug)]
pub(super) struct Current {
inner: Inner,
}
#[derive(Debug, Copy, Clone)]
struct Inner {
// thread-local variables cannot track generics. However, the current worker
// is only checked when `P` is already known, so the type can be figured out
// on demand.
workers: *const (),
idx: usize,
}
// Pointer to the current worker info
thread_local!(static CURRENT_WORKER: Cell<Inner> = Cell::new(Inner::new()));
pub(super) fn set<F, R, P>(pool: &Arc<worker::Set<P>>, index: usize, f: F) -> R
where
F: FnOnce() -> R,
P: Unpark,
{
CURRENT_WORKER.with(|cell| {
assert!(cell.get().workers.is_null());
struct Guard<'a>(&'a Cell<Inner>);
impl Drop for Guard<'_> {
fn drop(&mut self) {
self.0.set(Inner::new());
}
}
cell.set(Inner {
workers: pool.shared() as *const _ as *const (),
idx: index,
});
let _g = Guard(cell);
f()
})
}
pub(super) fn get<F, R>(f: F) -> R
where
F: FnOnce(&Current) -> R,
{
CURRENT_WORKER.with(|cell| {
let current = Current { inner: cell.get() };
f(&current)
})
}
impl Current {
pub(super) fn as_member<'a, P>(&self, set: &'a worker::Set<P>) -> Option<&'a Owned<P>>
where
P: Unpark,
{
let inner = CURRENT_WORKER.with(|cell| cell.get());
if ptr::eq(inner.workers as *const _, set.shared().as_ptr()) {
Some(unsafe { &*set.owned()[inner.idx].get() })
} else {
None
}
}
}
impl Inner {
fn new() -> Inner {
Inner {
workers: ptr::null(),
idx: 0,
}
}
}
tokio/src/executor/thread_pool/idle.rs
+229
View File
@@ -0,0 +1,229 @@
//! Coordinates idling workers
use crate::executor::loom::sync::atomic::AtomicUsize;
use crate::executor::loom::sync::Mutex;
use std::fmt;
use std::sync::atomic::Ordering::{self, AcqRel, Relaxed, SeqCst};
pub(super) struct Idle {
/// Tracks both the number of searching workers and the number of unparked
/// workers.
///
/// Used as a fast-path to avoid acquiring the lock when needed.
state: AtomicUsize,
/// Sleeping workers
sleepers: Mutex<Vec<usize>>,
/// Total number of workers.
num_workers: usize,
}
const UNPARK_SHIFT: usize = 16;
const UNPARK_MASK: usize = !SEARCH_MASK;
const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1;
#[derive(Copy, Clone)]
struct State(usize);
impl Idle {
pub(super) fn new(num_workers: usize) -> Idle {
let init = State::new(num_workers);
Idle {
state: AtomicUsize::new(init.into()),
sleepers: Mutex::new(Vec::with_capacity(num_workers)),
num_workers,
}
}
/// If there are no workers actively searching, returns the index of a
/// worker currently sleeping.
pub(super) fn worker_to_notify(&self) -> Option<usize> {
// If at least one worker is spinning, work being notified will
// eventully be found. A searching thread will find **some** work and
// notify another worker, eventually leading to our work being found.
//
// For this to happen, this load must happen before the thread
// transitioning `num_searching` to zero. Acquire / Relese does not
// provide sufficient guarantees, so this load is done with `SeqCst` and
// will pair with the `fetch_sub(1)` when transitioning out of
// searching.
if !self.notify_should_wakeup() {
return None;
}
// Acquire the lock
let mut sleepers = self.sleepers.lock().unwrap();
// Check again, now that the lock is acquired
if !self.notify_should_wakeup() {
return None;
}
// A worker should be woken up, atomically increment the number of
// searching workers as well as the number of unparked workers.
State::unpark_one(&self.state);
// Get the worker to unpark
let ret = sleepers.pop();
debug_assert!(ret.is_some());
ret
}
/// Returns `true` if the worker needs to do a final check for submitted
/// work.
pub(super) fn transition_worker_to_parked(&self, worker: usize, is_searching: bool) -> bool {
// Acquire the lock
let mut sleepers = self.sleepers.lock().unwrap();
// Decrement the number of unparked threads
let ret = State::dec_num_unparked(&self.state, is_searching);
// Track the sleeping worker
sleepers.push(worker);
ret
}
pub(super) fn transition_worker_to_searching(&self) -> bool {
// Using `Relaxed` ordering is acceptable here as it is just an
// optimization. This load has does not need to synchronize with
// anything, and the algorithm is correct no matter what the load
// returns (as in, it could return absolutely any `usize` value and the
// pool would be correct.
let state = State::load(&self.state, Relaxed);
if 2 * state.num_searching() >= self.num_workers {
return false;
}
// It is possible for this routine to allow more than 50% of the workers
// to search. That is OK. Limiting searchers is only an optimization to
// prevent too much contention.
//
// At this point, we do not need a hard synchronization with `notify_work`, so `AcqRel` is sufficient.
State::inc_num_searching(&self.state, AcqRel);
true
}
/// A lightweight transition from searching -> running.
///
/// Returns `true` if this is the final searching worker. The caller
/// **must** notify a new worker.
pub(super) fn transition_worker_from_searching(&self) -> bool {
State::dec_num_searching(&self.state)
}
/// Unpark a specific worker. This happens if tasks are submitted from
/// within the worker's park routine.
pub(super) fn unpark_worker_by_id(&self, worker_id: usize) {
let mut sleepers = self.sleepers.lock().unwrap();
for index in 0..sleepers.len() {
if sleepers[index] == worker_id {
sleepers.swap_remove(index);
// Update the state accordingly whle the lock is held.
State::unpark_one(&self.state);
return;
}
}
}
/// Returns `true` if `worker_id` is contained in the sleep set
pub(super) fn is_parked(&self, worker_id: usize) -> bool {
let sleepers = self.sleepers.lock().unwrap();
sleepers.contains(&worker_id)
}
fn notify_should_wakeup(&self) -> bool {
let state = State::load(&self.state, SeqCst);
state.num_searching() == 0 && state.num_unparked() < self.num_workers
}
}
impl State {
fn new(num_workers: usize) -> State {
// All workers start in the unparked state
let ret = State(num_workers << UNPARK_SHIFT);
debug_assert_eq!(num_workers, ret.num_unparked());
debug_assert_eq!(0, ret.num_searching());
ret
}
fn load(cell: &AtomicUsize, ordering: Ordering) -> State {
State(cell.load(ordering))
}
fn unpark_one(cell: &AtomicUsize) {
cell.fetch_add(1 | (1 << UNPARK_SHIFT), SeqCst);
}
fn inc_num_searching(cell: &AtomicUsize, ordering: Ordering) {
cell.fetch_add(1, ordering);
}
/// Returns `true` if this is the final searching worker
fn dec_num_searching(cell: &AtomicUsize) -> bool {
let state = State(cell.fetch_sub(1, SeqCst));
state.num_searching() == 1
}
/// Track a sleeping worker
///
/// Returns `true` if this is the final searching worker.
fn dec_num_unparked(cell: &AtomicUsize, is_searching: bool) -> bool {
let mut dec = 1 << UNPARK_SHIFT;
if is_searching {
dec += 1;
}
let prev = State(cell.fetch_sub(dec, SeqCst));
is_searching && prev.num_searching() == 1
}
/// Number of workers currently searching
fn num_searching(self) -> usize {
self.0 & SEARCH_MASK
}
/// Number of workers currently unparked
fn num_unparked(self) -> usize {
(self.0 & UNPARK_MASK) >> UNPARK_SHIFT
}
}
impl From<usize> for State {
fn from(src: usize) -> State {
State(src)
}
}
impl From<State> for usize {
fn from(src: State) -> usize {
src.0
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("worker::State")
.field("num_unparked", &self.num_unparked())
.field("num_searching", &self.num_searching())
.finish()
}
}
#[test]
fn test_state() {
assert_eq!(0, UNPARK_MASK & SEARCH_MASK);
assert_eq!(0, !(UNPARK_MASK | SEARCH_MASK));
let state = State::new(10);
assert_eq!(10, state.num_unparked());
assert_eq!(0, state.num_searching());
}
tokio/src/executor/thread_pool/join.rs
+42
View File
@@ -0,0 +1,42 @@
use crate::executor::park::Unpark;
use crate::executor::task;
use crate::executor::thread_pool::Shared;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
/// An owned permission to join on a task (await its termination).
pub struct JoinHandle<T> {
task: task::JoinHandle<T, Shared<Box<dyn Unpark>>>,
}
impl<T> JoinHandle<T>
where
T: Send + 'static,
{
pub(super) fn new(task: task::JoinHandle<T, Shared<Box<dyn Unpark>>>) -> JoinHandle<T> {
JoinHandle { task }
}
}
impl<T> Future for JoinHandle<T>
where
T: Send + 'static,
{
type Output = task::Result<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.task).poll(cx)
}
}
impl<T> fmt::Debug for JoinHandle<T>
where
T: fmt::Debug,
{
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("JoinHandle").finish()
}
}
tokio/src/executor/thread_pool/mod.rs
+58
View File
@@ -0,0 +1,58 @@
//! Threadpool
mod builder;
pub use self::builder::Builder;
mod current;
mod idle;
use self::idle::Idle;
mod join;
pub use self::join::JoinHandle;
mod owned;
use self::owned::Owned;
mod park;
mod pool;
pub use self::pool::ThreadPool;
mod queue;
mod spawner;
pub use self::spawner::Spawner;
mod set;
mod shared;
use self::shared::Shared;
mod shutdown;
mod worker;
/// Unit tests
#[cfg(test)]
mod tests;
// Re-export `task::Error`
pub use crate::executor::task::Error;
// These exports are used in tests
#[cfg(test)]
#[allow(warnings)]
pub(crate) use self::worker::create_set as create_pool;
pub(crate) type BoxFuture =
std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send + 'static>>;
#[cfg(not(loom))]
const LOCAL_QUEUE_CAPACITY: usize = 256;
// Shrink the size of the local queue when using loom. This shouldn't impact
// logic, but allows loom to test more edge cases in a reasonable a mount of
// time.
#[cfg(loom)]
const LOCAL_QUEUE_CAPACITY: usize = 2;
tokio/src/executor/thread_pool/owned.rs
+77
View File
@@ -0,0 +1,77 @@
use crate::executor::task::{self, Task};
use crate::executor::thread_pool::{queue, Shared};
use crate::executor::util::FastRand;
use std::cell::Cell;
/// Per-worker data accessible only by the thread driving the worker.
#[derive(Debug)]
pub(super) struct Owned<P: 'static> {
/// Worker tick number. Used to schedule bookkeeping tasks every so often.
pub(super) tick: Cell<u16>,
/// Caches the pool run state.
pub(super) is_running: Cell<bool>,
/// `true` if the worker is currently searching for more work.
pub(super) is_searching: Cell<bool>,
/// `true` when worker notification should be delayed.
///
/// This is used to batch notifications triggered by the parker.
pub(super) defer_notification: Cell<bool>,
/// `true` if a task was submitted while `defer_notification` was set
pub(super) did_submit_task: Cell<bool>,
/// Fast random number generator
pub(super) rand: FastRand,
/// Work queue
pub(super) work_queue: queue::Worker<Shared<P>>,
/// List of tasks owned by the worker
pub(super) owned_tasks: task::OwnedList<Shared<P>>,
}
impl<P> Owned<P>
where
P: 'static,
{
pub(super) fn new(work_queue: queue::Worker<Shared<P>>, rand: FastRand) -> Owned<P> {
Owned {
tick: Cell::new(1),
is_running: Cell::new(true),
is_searching: Cell::new(false),
defer_notification: Cell::new(false),
did_submit_task: Cell::new(false),
rand,
work_queue,
owned_tasks: task::OwnedList::new(),
}
}
/// Returns `true` if a worker should be notified
pub(super) fn submit_local(&self, task: Task<Shared<P>>) -> bool {
let ret = self.work_queue.push(task);
if self.defer_notification.get() {
self.did_submit_task.set(true);
false
} else {
ret
}
}
pub(super) fn submit_local_yield(&self, task: Task<Shared<P>>) {
self.work_queue.push_yield(task);
}
pub(super) fn bind_task(&mut self, task: &Task<Shared<P>>) {
self.owned_tasks.insert(task);
}
pub(super) fn release_task(&mut self, task: &Task<Shared<P>>) {
self.owned_tasks.remove(task);
}
}
tokio/src/executor/thread_pool/park.rs
+182
View File
@@ -0,0 +1,182 @@
use crate::executor::loom::sync::atomic::AtomicUsize;
use crate::executor::loom::sync::{Arc, Condvar, Mutex};
use crate::executor::park::{Park, Unpark};
use std::error::Error;
use std::fmt;
use std::sync::atomic::Ordering::SeqCst;
use std::time::Duration;
/// Parks the thread.
#[derive(Debug)]
pub(crate) struct DefaultPark {
inner: Arc<Inner>,
}
/// Unparks threads that were parked by `DefaultPark`.
#[derive(Debug)]
pub(crate) struct DefaultUnpark {
inner: Arc<Inner>,
}
/// Error returned by [`ParkThread`]
///
/// This currently is never returned, but might at some point in the future.
///
/// [`ParkThread`]: struct.ParkThread.html
#[derive(Debug)]
pub(crate) struct ParkError {
_p: (),
}
const EMPTY: usize = 0;
const PARKED: usize = 1;
const NOTIFIED: usize = 2;
#[derive(Debug)]
struct Inner {
state: AtomicUsize,
lock: Mutex<()>,
cvar: Condvar,
}
impl DefaultPark {
/// Creates a new `DefaultPark` instance.
pub(crate) fn new() -> DefaultPark {
DefaultPark {
inner: Arc::new(Inner {
state: AtomicUsize::new(EMPTY),
lock: Mutex::new(()),
cvar: Condvar::new(),
}),
}
}
}
impl Park for DefaultPark {
type Unpark = DefaultUnpark;
type Error = ParkError;
fn unpark(&self) -> Self::Unpark {
let inner = self.inner.clone();
DefaultUnpark { inner }
}
fn park(&mut self) -> Result<(), Self::Error> {
self.inner.park(None);
Ok(())
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
self.inner.park(Some(duration));
Ok(())
}
}
impl Unpark for DefaultUnpark {
fn unpark(&self) {
self.inner.unpark();
}
}
impl fmt::Display for ParkError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "unknown park error")
}
}
impl Error for ParkError {}
impl Inner {
fn park(&self, timeout: Option<Duration>) {
// If we were previously notified then we consume this notification and return quickly.
if self
.state
.compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst)
.is_ok()
{
return;
}
// If the timeout is zero, then there is no need to actually block.
if let Some(ref dur) = timeout {
if *dur == Duration::from_millis(0) {
return;
}
}
// Otherwise we need to coordinate going to sleep.
let mut _m = self.lock.lock().unwrap();
match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) {
Ok(_) => {}
// Consume this notification to avoid spurious wakeups in the next park.
Err(NOTIFIED) => {
// We must read `state` here, even though we know it will be `NOTIFIED`. This is
// because `unpark` may have been called again since we read `NOTIFIED` in the
// `compare_exchange` above. We must perform an acquire operation that synchronizes
// with that `unpark` to observe any writes it made before the call to `unpark`. To
// do that we must read from the write it made to `state`.
let old = self.state.swap(EMPTY, SeqCst);
assert_eq!(old, NOTIFIED, "park state changed unexpectedly");
return;
}
Err(n) => panic!("inconsistent park_timeout state: {}", n),
}
match timeout {
None => {
loop {
// Block the current thread on the conditional variable.
_m = self.cvar.wait(_m).unwrap();
if self
.state
.compare_exchange(NOTIFIED, EMPTY, SeqCst, SeqCst)
.is_ok()
{
return; // got a notification
}
// spurious wakeup, go back to sleep
}
}
Some(timeout) => {
// Wait with a timeout, and if we spuriously wake up or otherwise wake up from a
// notification we just want to unconditionally set `state` back to `EMPTY`, either
// consuming a notification or un-flagging ourselves as parked.
_m = self.cvar.wait_timeout(_m, timeout).unwrap().0;
match self.state.swap(EMPTY, SeqCst) {
NOTIFIED => {} // got a notification
PARKED => {} // no notification
n => panic!("inconsistent park_timeout state: {}", n),
}
}
}
}
fn unpark(&self) {
// To ensure the unparked thread will observe any writes we made before this call, we must
// perform a release operation that `park` can synchronize with. To do that we must write
// `NOTIFIED` even if `state` is already `NOTIFIED`. That is why this must be a swap rather
// than a compare-and-swap that returns if it reads `NOTIFIED` on failure.
match self.state.swap(NOTIFIED, SeqCst) {
EMPTY => return, // no one was waiting
NOTIFIED => return, // already unparked
PARKED => {} // gotta go wake someone up
n => panic!("inconsistent state in unpark: {}", n),
}
// There is a period between when the parked thread sets `state` to `PARKED` (or last
// checked `state` in the case of a spurious wakeup) and when it actually waits on `cvar`.
// If we were to notify during this period it would be ignored and then when the parked
// thread went to sleep it would never wake up. Fortunately, it has `lock` locked at this
// stage so we can acquire `lock` to wait until it is ready to receive the notification.
//
// Releasing `lock` before the call to `notify_one` means that when the parked thread wakes
// it doesn't get woken only to have to wait for us to release `lock`.
drop(self.lock.lock());
self.cvar.notify_one();
}
}
tokio/src/executor/thread_pool/pool.rs
+111
View File
@@ -0,0 +1,111 @@
use crate::executor::blocking::PoolWaiter;
use crate::executor::thread_pool::{shutdown, Builder, JoinHandle, Spawner};
use crate::executor::Executor;
use std::fmt;
use std::future::Future;
/// Work-stealing based thread pool for executing futures.
pub struct ThreadPool {
spawner: Spawner,
/// Shutdown waiter
shutdown_rx: shutdown::Receiver,
/// Shutdown valve for Pool
blocking: PoolWaiter,
}
impl ThreadPool {
/// Create a new ThreadPool with default configuration
pub fn new() -> ThreadPool {
Builder::new().build()
}
pub(super) fn from_parts(
spawner: Spawner,
shutdown_rx: shutdown::Receiver,
blocking: PoolWaiter,
) -> ThreadPool {
ThreadPool {
spawner,
shutdown_rx,
blocking,
}
}
/// Returns reference to `Spawner`.
///
/// The `Spawner` handle can be cloned and enables spawning tasks from other
/// threads.
pub fn spawner(&self) -> &Spawner {
&self.spawner
}
/// Spawn a task
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.spawner.spawn(future)
}
/// Spawn a task in the background
pub(crate) fn spawn_background<F>(&self, future: F)
where
F: Future<Output = ()> + Send + 'static,
{
self.spawner.spawn_background(future);
}
/// Block the current thread waiting for the future to complete.
///
/// The future will execute on the current thread, but all spawned tasks
/// will be executed on the thread pool.
pub fn block_on<F>(&self, future: F) -> F::Output
where
F: Future,
{
crate::executor::global::with_threadpool(self, || {
let mut enter = crate::executor::enter().expect("attempting to block while on a Tokio executor");
crate::executor::blocking::with_pool(self.spawner.blocking_pool(), || enter.block_on(future))
})
}
/// Shutdown the thread pool.
pub fn shutdown_now(&mut self) {
if self.spawner.workers().close() {
self.shutdown_rx.wait();
}
self.blocking.shutdown();
}
}
impl Default for ThreadPool {
fn default() -> ThreadPool {
ThreadPool::new()
}
}
impl Executor for &ThreadPool {
fn spawn(
&mut self,
future: std::pin::Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), crate::executor::SpawnError> {
ThreadPool::spawn_background(self, future);
Ok(())
}
}
impl fmt::Debug for ThreadPool {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("ThreadPool").finish()
}
}
impl Drop for ThreadPool {
fn drop(&mut self) {
self.shutdown_now();
}
}
tokio/src/executor/thread_pool/queue/global.rs
+195
View File
@@ -0,0 +1,195 @@
use crate::executor::loom::sync::atomic::AtomicUsize;
use crate::executor::loom::sync::Mutex;
use crate::executor::task::{Header, Task};
use std::ptr::{self, NonNull};
use std::sync::atomic::Ordering::{Acquire, Release};
use std::usize;
pub(super) struct Queue<T: 'static> {
/// Pointers to the head and tail of the queue
pointers: Mutex<Pointers<T>>,
/// Number of pending tasks in the queue. This helps prevent unnecessary
/// locking in the hot path.
///
/// The LSB is a flag tracking whether or not the queue is open or not.
len: AtomicUsize,
}
struct Pointers<T: 'static> {
head: *const Header<T>,
tail: *const Header<T>,
}
const CLOSED: usize = 1;
const MAX_LEN: usize = usize::MAX >> 1;
impl<T: 'static> Queue<T> {
pub(super) fn new() -> Queue<T> {
Queue {
pointers: Mutex::new(Pointers {
head: ptr::null(),
tail: ptr::null(),
}),
len: AtomicUsize::new(0),
}
}
pub(super) fn is_empty(&self) -> bool {
self.len() == 0
}
pub(super) fn is_closed(&self) -> bool {
self.len.load(Acquire) & CLOSED == CLOSED
}
/// Close the worker queue
pub(super) fn close(&self) -> bool {
// Acquire the lock
let _p = self.pointers.lock().unwrap();
let len = unsafe {
// Set the queue as closed. Because all mutations are synchronized by
// the mutex, a read followed by a write is acceptable.
self.len.unsync_load()
};
let ret = len & CLOSED == 0;
self.len.store(len | CLOSED, Release);
ret
}
fn len(&self) -> usize {
self.len.load(Acquire) >> 1
}
pub(super) fn wait_for_unlocked(&self) {
// Acquire and release the lock immediately. This synchronizes the
// caller **after** all external waiters are done w/ the scheduler
// struct.
drop(self.pointers.lock().unwrap());
}
/// Push a value into the queue and call the closure **while still holding
/// the push lock**
pub(super) fn push<F>(&self, task: Task<T>, f: F)
where
F: FnOnce(Result<(), Task<T>>),
{
unsafe {
// Acquire queue lock
let mut p = self.pointers.lock().unwrap();
// Check if the queue is closed. This must happen in the lock.
let len = self.len.unsync_load();
if len & CLOSED == CLOSED {
f(Err(task));
return;
}
let task = task.into_raw();
// The next pointer should already be null
debug_assert!(get_next(task).is_null());
if let Some(tail) = NonNull::new(p.tail as *mut _) {
set_next(tail, task.as_ptr());
} else {
p.head = task.as_ptr();
}
p.tail = task.as_ptr();
// Increment the count.
//
// All updates to the len atomic are guarded by the mutex. As such,
// a non-atomic load followed by a store is safe.
//
// We increment by 2 to avoid touching the shutdown flag
if (len >> 1) == MAX_LEN {
eprintln!("[ERROR] overflowed task counter. This is a bug and should be reported.");
std::process::abort();
}
self.len.store(len + 2, Release);
f(Ok(()));
}
}
pub(super) fn push_batch(&self, batch_head: Task<T>, batch_tail: Task<T>, num: usize) {
unsafe {
let batch_head = batch_head.into_raw().as_ptr();
let batch_tail = batch_tail.into_raw();
debug_assert!(get_next(batch_tail).is_null());
let mut p = self.pointers.lock().unwrap();
if let Some(tail) = NonNull::new(p.tail as *mut _) {
set_next(tail, batch_head);
} else {
p.head = batch_head;
}
p.tail = batch_tail.as_ptr();
// Increment the count.
//
// All updates to the len atomic are guarded by the mutex. As such,
// a non-atomic load followed by a store is safe.
//
// Left shift by 1 to avoid touching the shutdown flag.
let len = self.len.unsync_load();
if (len >> 1) >= (MAX_LEN - num) {
std::process::abort();
}
self.len.store(len + (num << 1), Release);
}
}
pub(super) fn pop(&self) -> Option<Task<T>> {
// Fast path, if len == 0, then there are no values
if self.is_empty() {
return None;
}
unsafe {
let mut p = self.pointers.lock().unwrap();
// It is possible to hit null here if another thread poped the last
// task between us checking `len` and acquiring the lock.
let task = NonNull::new(p.head as *mut _)?;
p.head = get_next(task);
if p.head.is_null() {
p.tail = ptr::null();
}
set_next(task, ptr::null());
// Decrement the count.
//
// All updates to the len atomic are guarded by the mutex. As such,
// a non-atomic load followed by a store is safe.
//
// Decrement by 2 to avoid touching the shutdown flag
self.len.store(self.len.unsync_load() - 2, Release);
Some(Task::from_raw(task))
}
}
}
unsafe fn get_next<T>(meta: NonNull<Header<T>>) -> *const Header<T> {
*meta.as_ref().queue_next.get()
}
unsafe fn set_next<T>(meta: NonNull<Header<T>>, val: *const Header<T>) {
*meta.as_ref().queue_next.get() = val;
}
tokio/src/executor/thread_pool/queue/inject.rs
+36
View File
@@ -0,0 +1,36 @@
use crate::executor::loom::sync::Arc;
use crate::executor::task::Task;
use crate::executor::thread_pool::queue::Cluster;
pub(crate) struct Inject<T: 'static> {
cluster: Arc<Cluster<T>>,
}
impl<T: 'static> Inject<T> {
pub(super) fn new(cluster: Arc<Cluster<T>>) -> Inject<T> {
Inject { cluster }
}
/// Push a value onto the queue
pub(crate) fn push<F>(&self, task: Task<T>, f: F)
where
F: FnOnce(Result<(), Task<T>>),
{
self.cluster.global.push(task, f)
}
/// Close the queue
///
/// Returns `true` if the channel was closed. `false` indicates the pool was
/// previously closed.
pub(crate) fn close(&self) -> bool {
self.cluster.global.close()
}
/// Wait for all locks on the queue to drop.
///
/// This is done by locking w/o doing anything.
pub(crate) fn wait_for_unlocked(&self) {
self.cluster.global.wait_for_unlocked();
}
}
tokio/src/executor/thread_pool/queue/local.rs
+298
View File
@@ -0,0 +1,298 @@
use crate::executor::loom::cell::{CausalCell, CausalCheck};
use crate::executor::loom::sync::atomic::{self, AtomicU32};
use crate::executor::task::Task;
use crate::executor::thread_pool::queue::global;
use crate::executor::thread_pool::LOCAL_QUEUE_CAPACITY;
use std::fmt;
use std::mem::MaybeUninit;
use std::ptr;
use std::sync::atomic::Ordering::{Acquire, Release};
pub(super) struct Queue<T: 'static> {
/// Concurrently updated by many threads.
head: AtomicU32,
/// Only updated by producer thread but read by many threads.
tail: AtomicU32,
/// Elements
buffer: Box<[CausalCell<MaybeUninit<Task<T>>>]>,
}
const MASK: usize = LOCAL_QUEUE_CAPACITY - 1;
impl<T: 'static> Queue<T> {
pub(super) fn new() -> Queue<T> {
debug_assert!(LOCAL_QUEUE_CAPACITY >= 2 && LOCAL_QUEUE_CAPACITY.is_power_of_two());
let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY);
for _ in 0..LOCAL_QUEUE_CAPACITY {
buffer.push(CausalCell::new(MaybeUninit::uninit()));
}
Queue {
head: AtomicU32::new(0),
tail: AtomicU32::new(0),
buffer: buffer.into(),
}
}
}
impl<T> Queue<T> {
/// Push a task onto the local queue.
///
/// This **must** be called by the producer thread.
pub(super) unsafe fn push(&self, mut task: Task<T>, global: &global::Queue<T>) {
loop {
let head = self.head.load(Acquire);
// safety: this is the **only** thread that updates this cell.
let tail = self.tail.unsync_load();
if tail.wrapping_sub(head) < LOCAL_QUEUE_CAPACITY as u32 {
// Map the position to a slot index.
let idx = tail as usize & MASK;
self.buffer[idx].with_mut(|ptr| {
// Write the task to the slot
ptr::write((*ptr).as_mut_ptr(), task);
});
// Make the task available
self.tail.store(tail.wrapping_add(1), Release);
return;
}
// The local buffer is full. Push a batch of work to the global
// queue.
match self.push_overflow(task, head, tail, global) {
Ok(_) => return,
// Lost the race, try again
Err(v) => task = v,
}
atomic::spin_loop_hint();
}
}
/// Move a batch of tasks into the global queue.
///
/// This will temporarily make some of the tasks unavailable to stealers.
/// Once `push_overflow` is done, a notification is sent out, so if other
/// workers "missed" some of the tasks during a steal, they will get
/// another opportunity.
#[inline(never)]
unsafe fn push_overflow(
&self,
task: Task<T>,
head: u32,
tail: u32,
global: &global::Queue<T>,
) -> Result<(), Task<T>> {
const BATCH_LEN: usize = LOCAL_QUEUE_CAPACITY / 2 + 1;
let n = tail.wrapping_sub(head) / 2;
assert_eq!(n as usize, LOCAL_QUEUE_CAPACITY / 2, "queue is not full");
// Claim a bunch of tasks
//
// We are claiming the tasks **before** reading them out of the buffer.
// This is safe because only the **current** thread is able to push new
// tasks.
//
// There isn't really any need for memory ordering... Relaxed would
// work. This is because all tasks are pushed into the queue from the
// current thread (or memory has been acquired if the local queue handle
// moved).
let actual = self.head.compare_and_swap(head, head + n, Release);
if actual != head {
// We failed to claim the tasks, losing the race. Return out of
// this function and try the full `push` routine again. The queue
// may not be full anymore.
return Err(task);
}
// link the tasks
for i in 0..n {
let j = i + 1;
let i_idx = (i + head) as usize & MASK;
let j_idx = (j + head) as usize & MASK;
// Get the next pointer
let next = if j == n {
// The last task in the local queue being moved
task.header() as *const _
} else {
self.buffer[j_idx].with(|ptr| {
let value = (*ptr).as_ptr();
(*value).header() as *const _
})
};
self.buffer[i_idx].with_mut(|ptr| {
let ptr = (*ptr).as_ptr();
debug_assert!((*(*ptr).header().queue_next.get()).is_null());
*(*ptr).header().queue_next.get() = next;
});
}
let head = self.buffer[head as usize & MASK].with(|ptr| ptr::read((*ptr).as_ptr()));
// Push the tasks onto the global queue
global.push_batch(head, task, BATCH_LEN);
Ok(())
}
/// Pop a task from the local queue.
///
/// This **must** be called by the producer thread
pub(super) unsafe fn pop(&self) -> Option<Task<T>> {
loop {
let head = self.head.load(Acquire);
// safety: this is the **only** thread that updates this cell.
let tail = self.tail.unsync_load();
if head == tail {
// queue is empty
return None;
}
// Map the head position to a slot index.
let idx = head as usize & MASK;
let task = self.buffer[idx].with(|ptr| {
// Tentatively read the task at the head position. Note that we
// have not yet claimed the task.
//
ptr::read(ptr)
});
// Attempt to claim the task read above.
let actual = self
.head
.compare_and_swap(head, head.wrapping_add(1), Release);
if actual == head {
return Some(task.assume_init());
}
atomic::spin_loop_hint();
}
}
pub(super) fn is_empty(&self) -> bool {
let head = self.head.load(Acquire);
let tail = self.tail.load(Acquire);
head == tail
}
/// Steal half the tasks from self and place them into `dst`.
pub(super) unsafe fn steal(&self, dst: &Queue<T>) -> Option<Task<T>> {
let dst_tail = dst.tail.unsync_load();
// Steal the tasks into `dst`'s buffer. This does not yet expose the
// tasks in `dst`.
let mut n = self.steal2(dst, dst_tail);
if n == 0 {
// No tasks were stolen
return None;
}
// We are returning a task here
n -= 1;
let ret_pos = dst_tail.wrapping_add(n);
let ret_idx = ret_pos as usize & MASK;
let ret = dst.buffer[ret_idx].with(|ptr| ptr::read((*ptr).as_ptr()));
if n == 0 {
// The `dst` queue is empty, but a single task was stolen
return Some(ret);
}
// Synchronize with stealers
let dst_head = dst.head.load(Acquire);
assert!(dst_tail.wrapping_sub(dst_head) + n <= LOCAL_QUEUE_CAPACITY as u32);
// Make the stolen items available to consumers
dst.tail.store(dst_tail.wrapping_add(n), Release);
Some(ret)
}
unsafe fn steal2(&self, dst: &Queue<T>, dst_tail: u32) -> u32 {
loop {
let src_head = self.head.load(Acquire);
let src_tail = self.tail.load(Acquire);
// Number of available tasks to steal
let n = src_tail.wrapping_sub(src_head);
let n = n - n / 2;
if n == 0 {
return 0;
}
if n > LOCAL_QUEUE_CAPACITY as u32 / 2 {
atomic::spin_loop_hint();
// inconsistent, try again
continue;
}
// Track CausalCell causality checks. The check is deferred until
// the compare_and_swap claims ownership of the tasks.
let mut check = CausalCheck::default();
for i in 0..n {
// Compute the positions
let src_pos = src_head.wrapping_add(i);
let dst_pos = dst_tail.wrapping_add(i);
// Map to slots
let src_idx = src_pos as usize & MASK;
let dst_idx = dst_pos as usize & MASK;
// Read the task
let (task, ch) =
self.buffer[src_idx].with_deferred(|ptr| ptr::read((*ptr).as_ptr()));
check.join(ch);
// Write the task to the new slot
dst.buffer[dst_idx].with_mut(|ptr| ptr::write((*ptr).as_mut_ptr(), task));
}
// Claim all of those tasks!
let actual = self
.head
.compare_and_swap(src_head, src_head.wrapping_add(n), Release);
if actual == src_head {
check.check();
return n;
}
atomic::spin_loop_hint();
}
}
}
impl<T> fmt::Debug for Queue<T> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("local::Queue")
.field("head", &self.head)
.field("tail", &self.tail)
.field("buffer", &"[...]")
.finish()
}
}
tokio/src/executor/thread_pool/queue/mod.rs
+41
View File
@@ -0,0 +1,41 @@
//! The threadpool's task queue system.
mod global;
mod inject;
mod local;
mod worker;
pub(crate) use self::inject::Inject;
pub(crate) use self::worker::Worker;
use crate::executor::loom::sync::Arc;
pub(crate) fn build<T: 'static>(workers: usize) -> Vec<Worker<T>> {
let local: Vec<_> = (0..workers).map(|_| local::Queue::new()).collect();
let cluster = Arc::new(Cluster {
local: local.into_boxed_slice(),
global: global::Queue::new(),
});
(0..workers)
.map(|index| Worker::new(cluster.clone(), index))
.collect()
}
struct Cluster<T: 'static> {
/// per-worker local queues
local: Box<[local::Queue<T>]>,
global: global::Queue<T>,
}
impl<T: 'static> Drop for Cluster<T> {
fn drop(&mut self) {
// Drain all the queues
for queue in &self.local[..] {
while let Some(_) = unsafe { queue.pop() } {}
}
while let Some(_) = self.global.pop() {}
}
}
tokio/src/executor/thread_pool/queue/worker.rs
+127
View File
@@ -0,0 +1,127 @@
use crate::executor::loom::sync::Arc;
use crate::executor::task::Task;
use crate::executor::thread_pool::queue::{local, Cluster, Inject};
use std::cell::Cell;
use std::fmt;
pub(crate) struct Worker<T: 'static> {
cluster: Arc<Cluster<T>>,
index: u16,
/// Task to pop next
next: Cell<Option<Task<T>>>,
}
impl<T: 'static> Worker<T> {
pub(super) fn new(cluster: Arc<Cluster<T>>, index: usize) -> Worker<T> {
Worker {
cluster,
index: index as u16,
next: Cell::new(None),
}
}
pub(crate) fn injector(&self) -> Inject<T> {
Inject::new(self.cluster.clone())
}
/// Returns `true` if the queue is closed
pub(crate) fn is_closed(&self) -> bool {
self.cluster.global.is_closed()
}
/// Push to the local queue.
///
/// If the local queue is full, the task is pushed onto the global queue.
///
/// # Return
///
/// Returns `true` if the pushed task can be stolen by another worker.
pub(crate) fn push(&self, task: Task<T>) -> bool {
let prev = self.next.take();
let ret = prev.is_some();
if let Some(prev) = prev {
// safety: we guarantee that only one thread pushes to this local
// queue at a time.
unsafe {
self.local().push(prev, &self.cluster.global);
}
}
self.next.set(Some(task));
ret
}
pub(crate) fn push_yield(&self, task: Task<T>) {
unsafe { self.local().push(task, &self.cluster.global) }
}
/// Pop a task checking the local queue first.
pub(crate) fn pop_local_first(&self) -> Option<Task<T>> {
self.local_pop().or_else(|| self.cluster.global.pop())
}
/// Pop a task checking the global queue first.
pub(crate) fn pop_global_first(&self) -> Option<Task<T>> {
self.cluster.global.pop().or_else(|| self.local_pop())
}
/// Steal from other local queues.
///
/// `start` specifies the queue from which to start stealing.
pub(crate) fn steal(&self, start: usize) -> Option<Task<T>> {
let num_queues = self.cluster.local.len();
for i in 0..num_queues {
let i = (start + i) % num_queues;
if i == self.index as usize {
continue;
}
// safety: we own the dst queue
let ret = unsafe { self.cluster.local[i].steal(self.local()) };
if ret.is_some() {
return ret;
}
}
None
}
/// An approximation of whether or not the queue is empty.
pub(crate) fn is_empty(&self) -> bool {
for local_queue in &self.cluster.local[..] {
if !local_queue.is_empty() {
return false;
}
}
self.cluster.global.is_empty()
}
fn local_pop(&self) -> Option<Task<T>> {
if let Some(task) = self.next.take() {
return Some(task);
}
// safety: we guarantee that only one thread pushes to this local queue
// at a time.
unsafe { self.local().pop() }
}
fn local(&self) -> &local::Queue<T> {
&self.cluster.local[self.index as usize]
}
}
impl<T: 'static> fmt::Debug for Worker<T> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("queue::Worker")
.field("cluster", &"...")
.field("index", &self.index)
.finish()
}
}
tokio/src/executor/thread_pool/set.rs
+209
View File
@@ -0,0 +1,209 @@
//! Putting a worker to sleep.
//!
//! - Attempt to spin.
use crate::executor::loom::rand::seed;
use crate::executor::loom::sync::Arc;
use crate::executor::park::Unpark;
use crate::executor::task::{self, Task};
use crate::executor::thread_pool::{current, queue, BoxFuture, Idle, JoinHandle, Owned, Shared};
use crate::executor::util::{CachePadded, FastRand};
use crate::executor::{Executor, SpawnError};
use std::cell::UnsafeCell;
use std::future::Future;
pub(crate) struct Set<P>
where
P: 'static,
{
/// Data accessible from all workers.
shared: Box<[Shared<P>]>,
/// Data owned by the worker.
owned: Box<[UnsafeCell<CachePadded<Owned<P>>>]>,
/// Submit work to the pool while *not* currently on a worker thread.
inject: queue::Inject<Shared<P>>,
/// Coordinates idle workers
idle: Idle,
/// Pool where blocking tasks should be spawned.
pub(crate) blocking: Arc<crate::executor::blocking::Pool>,
}
unsafe impl<P: Unpark> Send for Set<P> {}
unsafe impl<P: Unpark> Sync for Set<P> {}
impl<P> Set<P>
where
P: Unpark,
{
/// Create a new worker set using the provided queues.
pub(crate) fn new<F>(
num_workers: usize,
mut mk_unpark: F,
blocking: Arc<crate::executor::blocking::Pool>,
) -> Self
where
F: FnMut(usize) -> P,
{
assert!(num_workers > 0);
let queues = queue::build(num_workers);
let inject = queues[0].injector();
let mut shared = Vec::with_capacity(queues.len());
let mut owned = Vec::with_capacity(queues.len());
for (i, queue) in queues.into_iter().enumerate() {
let unpark = mk_unpark(i);
let rand = FastRand::new(seed());
shared.push(Shared::new(unpark));
owned.push(UnsafeCell::new(CachePadded::new(Owned::new(queue, rand))));
}
Set {
shared: shared.into_boxed_slice(),
owned: owned.into_boxed_slice(),
inject,
idle: Idle::new(num_workers),
blocking,
}
}
fn inject_task(&self, task: Task<Shared<P>>) {
self.inject.push(task, |res| {
if let Err(task) = res {
task.shutdown();
// There may be a worker, in the process of being shutdown, that is
// waiting for this task to be released, so we notify all workers
// just in case.
//
// Over aggressive, but the runtime is in the process of shutting
// down, so efficiency is not critical.
self.notify_all();
} else {
self.notify_work();
}
});
}
pub(super) fn notify_work(&self) {
if let Some(index) = self.idle.worker_to_notify() {
self.shared[index].unpark();
}
}
pub(super) fn notify_all(&self) {
for shared in &self.shared[..] {
shared.unpark();
}
}
pub(crate) fn spawn_background<F>(&self, future: F)
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let task = task::background(future);
self.schedule(task);
}
pub(super) fn blocking_pool(&self) -> &Arc<crate::executor::blocking::Pool> {
&self.blocking
}
pub(crate) fn schedule(&self, task: Task<Shared<P>>) {
current::get(|current_worker| match current_worker.as_member(self) {
Some(worker) => {
if worker.submit_local(task) {
self.notify_work();
}
}
None => {
self.inject_task(task);
}
})
}
pub(crate) fn set_container_ptr(&mut self) {
let ptr = self as *const _;
for shared in &mut self.shared[..] {
shared.set_container_ptr(ptr);
}
}
/// Signal the pool is closed
///
/// Returns `true` if the transition to closed is successful. `false`
/// indicates the pool was already closed.
pub(crate) fn close(&self) -> bool {
if self.inject.close() {
self.notify_all();
true
} else {
false
}
}
pub(crate) fn len(&self) -> usize {
self.shared.len()
}
pub(super) fn index_of(&self, shared: &Shared<P>) -> usize {
use std::mem;
let size = mem::size_of::<Shared<P>>();
((shared as *const _ as usize) - (&self.shared[0] as *const _ as usize)) / size
}
pub(super) fn shared(&self) -> &[Shared<P>] {
&self.shared
}
pub(super) fn owned(&self) -> &[UnsafeCell<CachePadded<Owned<P>>>] {
&self.owned
}
pub(super) fn idle(&self) -> &Idle {
&self.idle
}
}
impl<P: 'static> Drop for Set<P> {
fn drop(&mut self) {
// Before proceeding, wait for all concurrent wakers to exit
self.inject.wait_for_unlocked();
}
}
impl Set<Box<dyn Unpark>> {
pub(crate) fn spawn_typed<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let (task, handle) = task::joinable(future);
self.schedule(task);
JoinHandle::new(handle)
}
}
impl<P> Executor for &Set<P>
where
P: Unpark,
{
fn spawn(&mut self, future: BoxFuture) -> Result<(), SpawnError> {
self.spawn_background(future);
Ok(())
}
fn status(&self) -> Result<(), SpawnError> {
Ok(())
}
}
tokio/src/executor/thread_pool/shared.rs
+104
View File
@@ -0,0 +1,104 @@
use crate::executor::park::Unpark;
use crate::executor::task::{self, Schedule, Task};
use crate::executor::thread_pool::worker;
use std::ptr;
/// Per-worker data accessible from any thread.
///
/// Accessed by:
///
/// - other workers
/// - tasks
///
pub(crate) struct Shared<P>
where
P: 'static,
{
/// Thread unparker
unpark: P,
/// Tasks pending drop. Any worker pushes tasks, only the "owning" worker
/// pops.
pub(super) pending_drop: task::TransferStack<Self>,
/// Untracked pointer to the pool.
///
/// The pool itself is tracked by an `Arc`, but this pointer is not included
/// in the ref count.
///
/// # Safety
///
/// `Worker` instances are stored in the `Pool` and are never removed.
set: *const worker::Set<P>,
}
unsafe impl<P: Unpark> Send for Shared<P> {}
unsafe impl<P: Unpark> Sync for Shared<P> {}
impl<P> Shared<P>
where
P: Unpark,
{
pub(super) fn new(unpark: P) -> Shared<P> {
Shared {
unpark,
pending_drop: task::TransferStack::new(),
set: ptr::null(),
}
}
pub(crate) fn schedule(&self, task: Task<Self>) {
self.set().schedule(task);
}
pub(super) fn unpark(&self) {
self.unpark.unpark();
}
pub(super) fn set_container_ptr(&mut self, set: *const worker::Set<P>) {
self.set = set;
}
fn set(&self) -> &worker::Set<P> {
unsafe { &*self.set }
}
}
impl<P> Schedule for Shared<P>
where
P: Unpark,
{
fn bind(&self, task: &Task<Self>) {
// Get access to the Owned component. This function can only be called
// when on the worker.
unsafe {
let index = self.set().index_of(self);
let owned = &mut *self.set().owned()[index].get();
owned.bind_task(task);
}
}
fn release(&self, task: Task<Self>) {
// This stores the task with the owning worker. The worker is not
// notified. Instead, the worker will clean up the tasks "eventually".
//
self.pending_drop.push(task);
}
fn release_local(&self, task: &Task<Self>) {
// Get access to the Owned component. This function can only be called
// when on the worker.
unsafe {
let index = self.set().index_of(self);
let owned = &mut *self.set().owned()[index].get();
owned.release_task(task);
}
}
fn schedule(&self, task: Task<Self>) {
Self::schedule(self, task);
}
}
tokio/src/executor/thread_pool/shutdown.rs
+48
View File
@@ -0,0 +1,48 @@
//! A shutdown channel.
//!
//! Each worker holds the `Sender` half. When all the `Sender` halves are
//! dropped, the `Receiver` receives a notification.
use crate::executor::loom::sync::Arc;
use tokio_sync::oneshot;
#[derive(Debug, Clone)]
pub(super) struct Sender {
tx: Arc<oneshot::Sender<()>>,
}
#[derive(Debug)]
pub(super) struct Receiver {
rx: oneshot::Receiver<()>,
}
pub(super) fn channel() -> (Sender, Receiver) {
let (tx, rx) = oneshot::channel();
let tx = Sender { tx: Arc::new(tx) };
let rx = Receiver { rx };
(tx, rx)
}
impl Receiver {
/// Block the current thread until all `Sender` handles drop.
pub(crate) fn wait(&mut self) {
use crate::executor::enter;
let mut e = match enter() {
Ok(e) => e,
Err(_) => {
if std::thread::panicking() {
// Already panicking, avoid a double panic
return;
} else {
panic!("cannot block on shutdown from the Tokio runtime");
}
}
};
// The oneshot completes with an Err
let _ = e.block_on(&mut self.rx);
}
}
tokio/src/executor/thread_pool/spawner.rs
+61
View File
@@ -0,0 +1,61 @@
use crate::executor::loom::sync::Arc;
use crate::executor::park::Unpark;
use crate::executor::thread_pool::{worker, JoinHandle};
use std::fmt;
use std::future::Future;
/// Submit futures to the associated thread pool for execution.
///
/// A `Spawner` instance is a handle to a single thread pool, allowing the owner
/// of the handle to spawn futures onto the thread pool.
///
/// The `Spawner` handle is *only* used for spawning new futures. It does not
/// impact the lifecycle of the thread pool in any way. The thread pool may
/// shutdown while there are outstanding `Spawner` instances.
///
/// `Spawner` instances are obtained by calling [`ThreadPool::spawner`].
///
/// [`ThreadPool::spawner`]: struct.ThreadPool.html#method.spawner
#[derive(Clone)]
pub struct Spawner {
workers: Arc<worker::Set<Box<dyn Unpark>>>,
}
impl Spawner {
pub(super) fn new(workers: Arc<worker::Set<Box<dyn Unpark>>>) -> Spawner {
Spawner { workers }
}
/// Spawn a future onto the thread pool
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
self.workers.spawn_typed(future)
}
/// Spawn a task in the background
pub(super) fn spawn_background<F>(&self, future: F)
where
F: Future<Output = ()> + Send + 'static,
{
self.workers.spawn_background(future);
}
pub(super) fn blocking_pool(&self) -> &Arc<crate::executor::blocking::Pool> {
self.workers.blocking_pool()
}
/// Reference to the worker set. Used by `ThreadPool` to initiate shutdown.
pub(super) fn workers(&self) -> &worker::Set<Box<dyn Unpark>> {
&*self.workers
}
}
impl fmt::Debug for Spawner {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Spawner").finish()
}
}
tokio/src/executor/thread_pool/tests/loom_pool.rs
+138
View File
@@ -0,0 +1,138 @@
use crate::spawn;
use crate::executor::loom::sync::atomic::Ordering::{Acquire, Relaxed, Release};
use crate::executor::loom::sync::atomic::{AtomicBool, AtomicUsize};
use crate::executor::loom::sync::{Arc, Mutex};
use crate::executor::tests::loom_oneshot as oneshot;
use crate::executor::thread_pool::ThreadPool;
use std::future::Future;
#[test]
fn pool_multi_spawn() {
loom::model(|| {
let pool = ThreadPool::new();
let c1 = Arc::new(AtomicUsize::new(0));
let (tx, rx) = oneshot::channel();
let tx1 = Arc::new(Mutex::new(Some(tx)));
// Spawn a task
let c2 = c1.clone();
let tx2 = tx1.clone();
pool.spawn(async move {
spawn(async move {
if 1 == c1.fetch_add(1, Relaxed) {
tx1.lock().unwrap().take().unwrap().send(());
}
});
});
// Spawn a second task
pool.spawn(async move {
spawn(async move {
if 1 == c2.fetch_add(1, Relaxed) {
tx2.lock().unwrap().take().unwrap().send(());
}
});
});
rx.recv();
});
}
#[test]
fn pool_multi_notify() {
loom::model(|| {
let pool = ThreadPool::new();
let c1 = Arc::new(AtomicUsize::new(0));
let (done_tx, done_rx) = oneshot::channel();
let done_tx1 = Arc::new(Mutex::new(Some(done_tx)));
// Spawn a task
let c2 = c1.clone();
let done_tx2 = done_tx1.clone();
pool.spawn(async move {
gated().await;
gated().await;
if 1 == c1.fetch_add(1, Relaxed) {
done_tx1.lock().unwrap().take().unwrap().send(());
}
});
// Spawn a second task
pool.spawn(async move {
gated().await;
gated().await;
if 1 == c2.fetch_add(1, Relaxed) {
done_tx2.lock().unwrap().take().unwrap().send(());
}
});
done_rx.recv();
});
}
#[test]
fn pool_shutdown() {
loom::model(|| {
let pool = ThreadPool::new();
pool.spawn(async move {
gated2(true).await;
});
pool.spawn(async move {
gated2(false).await;
});
drop(pool);
});
}
fn gated() -> impl Future<Output = &'static str> {
gated2(false)
}
fn gated2(thread: bool) -> impl Future<Output = &'static str> {
use crate::executor::loom::thread;
use futures_util::future::poll_fn;
use std::sync::Arc;
use std::task::Poll;
let gate = Arc::new(AtomicBool::new(false));
let mut fired = false;
poll_fn(move |cx| {
if !fired {
let gate = gate.clone();
let waker = cx.waker().clone();
if thread {
thread::spawn(move || {
gate.store(true, Release);
waker.wake_by_ref();
});
} else {
spawn(async move {
gate.store(true, Release);
waker.wake_by_ref();
});
}
fired = true;
return Poll::Pending;
}
if gate.load(Acquire) {
Poll::Ready("hello world")
} else {
Poll::Pending
}
})
}
tokio/src/executor/thread_pool/tests/loom_queue.rs
+68
View File
@@ -0,0 +1,68 @@
use crate::executor::task::{self, Task};
use crate::executor::tests::mock_schedule::{Noop, NOOP_SCHEDULE};
use crate::executor::thread_pool::queue;
use loom::thread;
use std::cell::Cell;
use std::rc::Rc;
#[test]
fn multi_worker() {
const THREADS: usize = 2;
const PER_THREAD: usize = 7;
fn work(_i: usize, q: queue::Worker<Noop>, rem: Rc<Cell<usize>>) {
let mut rem_local = PER_THREAD;
while rem.get() != 0 {
for _ in 0..3 {
if rem_local > 0 {
q.push(val(0));
rem_local -= 1;
}
}
// Try to work
while let Some(task) = q.pop_local_first() {
assert!(task.run(From::from(&NOOP_SCHEDULE)).is_none());
let r = rem.get();
assert!(r > 0);
rem.set(r - 1);
}
// Try to steal
if let Some(task) = q.steal(0) {
assert!(task.run(From::from(&NOOP_SCHEDULE)).is_none());
let r = rem.get();
assert!(r > 0);
rem.set(r - 1);
}
thread::yield_now();
}
}
loom::model(|| {
let rem = Rc::new(Cell::new(THREADS * PER_THREAD));
let mut qs = queue::build(THREADS);
let q1 = qs.remove(0);
for i in 1..THREADS {
let q = qs.remove(0);
let rem = rem.clone();
thread::spawn(move || {
work(i, q, rem);
});
}
work(0, q1, rem);
// th.join().unwrap();
});
}
fn val(num: u32) -> Task<Noop> {
task::background(async move { num })
}
tokio/src/executor/thread_pool/tests/mod.rs
+11
View File
@@ -0,0 +1,11 @@
#[cfg(loom)]
mod loom_pool;
#[cfg(loom)]
mod loom_queue;
#[cfg(not(loom))]
mod queue;
#[cfg(not(loom))]
mod worker;
tokio/src/executor/thread_pool/tests/queue.rs
+281
View File
@@ -0,0 +1,281 @@
use crate::executor::task::{self, Task};
use crate::executor::tests::mock_schedule::{Noop, NOOP_SCHEDULE};
use crate::executor::thread_pool::{queue, LOCAL_QUEUE_CAPACITY};
macro_rules! assert_pop {
($q:expr, $expect:expr) => {
assert_eq!(
match $q.pop_local_first() {
Some(v) => num(v),
None => panic!("queue empty"),
},
$expect
)
};
}
macro_rules! assert_pop_global {
($q:expr, $expect:expr) => {
assert_eq!(
match $q.pop_global_first() {
Some(v) => num(v),
None => panic!("queue empty"),
},
$expect
)
};
}
macro_rules! assert_steal {
($q:expr, $n:expr, $expect:expr) => {
assert_eq!(
match $q.steal($n) {
Some(v) => num(v),
None => panic!("queue empty"),
},
$expect
)
};
}
macro_rules! assert_empty {
($q:expr) => {{
let q: &mut queue::Worker<Noop> = &mut $q;
match q.pop_local_first() {
Some(v) => panic!("expected emtpy queue; got {}", num(v)),
None => {}
}
}};
}
#[test]
fn single_worker_push_pop() {
let mut q = queue::build(1).remove(0);
// Queue is empty
assert_empty!(q);
// Push a value
q.push(val(0));
// Pop the value
assert_pop!(q, 0);
// Push two values
q.push(val(1));
q.push(val(2));
q.push(val(3));
// Pop the value
assert_pop!(q, 3);
assert_pop!(q, 1);
assert_pop!(q, 2);
assert_empty!(q);
}
#[test]
fn multi_worker_push_pop() {
let (mut q1, mut q2) = queues_2();
// Queue is empty
assert_empty!(q1);
assert_empty!(q2);
// Push a value
q1.push(val(0));
// Not available on other queue
assert_empty!(q2);
assert_pop!(q1, 0);
q2.push(val(1));
assert_pop!(q2, 1);
assert_empty!(q1);
}
#[test]
fn multi_worker_inject_pop() {
let (mut q1, mut q2) = queues_2();
let i = q1.injector();
// Push a value
i.push(val(0), is_ok);
assert_pop!(q1, 0);
assert_empty!(q2);
// Push another value
i.push(val(1), is_ok);
assert_pop!(q2, 1);
assert_empty!(q1);
i.push(val(2), is_ok);
i.push(val(3), is_ok);
i.push(val(4), is_ok);
assert_pop!(q2, 2);
assert_pop!(q1, 3);
assert_pop!(q1, 4);
}
#[test]
fn overflow_local_queue() {
let (mut q1, mut q2) = queues_2();
for i in 0..LOCAL_QUEUE_CAPACITY {
q1.push(val(i as u32));
}
assert_empty!(q2);
// Fill `next` slot
q1.push(val(999));
// overflow
q1.push(val(1000));
assert_pop!(q2, 0);
assert_pop!(q1, 1000);
// Half the values were moved to the global queue
for i in 128..LOCAL_QUEUE_CAPACITY {
assert_pop!(q1, i as u32);
}
for i in 1..128 {
assert_pop!(q2, i);
}
assert_pop!(q2, 999);
assert_empty!(q2);
assert_empty!(q1);
}
#[test]
fn polling_global_first() {
let (q, _) = queues_2();
let i = q.injector();
i.push(val(1000), is_ok);
i.push(val(1001), is_ok);
for n in 0..5 {
q.push(val(n));
}
assert_pop_global!(q, 1000);
assert_pop!(q, 4);
assert_pop_global!(q, 1001);
assert_pop_global!(q, 0);
assert_pop!(q, 1);
assert_pop_global!(q, 2);
assert_pop_global!(q, 3);
assert!(q.pop_global_first().is_none());
}
#[test]
fn steal() {
let mut qs = queue::build(3);
let (mut q1, mut q2, mut q3) = (qs.remove(0), qs.remove(0), qs.remove(0));
assert!(q1.steal(0).is_none());
assert!(q2.steal(0).is_none());
assert!(q3.steal(0).is_none());
// Steal one value, but not the first one
q1.push(val(0));
q1.push(val(999));
assert_steal!(q2, 0, 0);
assert!(q2.steal(0).is_none());
assert_pop!(q1, 999);
// Steals half the queue
for i in 0..4 {
q1.push(val(i));
}
q1.push(val(999));
assert_steal!(q2, 0, 1);
assert_pop!(q2, 0);
assert_empty!(q2);
assert_pop!(q1, 999);
assert_pop!(q1, 2);
assert_pop!(q1, 3);
assert_empty!(q1);
// Searches multiple queues
q3.push(val(0));
q3.push(val(999));
assert_steal!(q2, 0, 0);
assert_pop!(q3, 999);
assert_empty!(q3);
// Steals from one queue at a time
q1.push(val(0));
q1.push(val(998));
q2.push(val(1));
q2.push(val(999));
assert_steal!(q3, 0, 0);
assert_pop!(q2, 999);
assert_pop!(q2, 1);
assert_empty!(q2);
assert_pop!(q1, 998);
assert_empty!(q1);
}
fn queues_2() -> (queue::Worker<Noop>, queue::Worker<Noop>) {
let mut qs = queue::build(2);
(qs.remove(0), qs.remove(0))
}
// pretty big hack to track tasks
use std::cell::RefCell;
use std::collections::HashMap;
thread_local! {
static TASKS: RefCell<HashMap<u32, task::JoinHandle<u32, Noop>>> = RefCell::new(HashMap::new())
}
fn val(num: u32) -> Task<Noop> {
let (task, join) = task::joinable(async move { num });
let prev = TASKS.with(|t| t.borrow_mut().insert(num, join));
assert!(prev.is_none());
task
}
fn num(task: Task<Noop>) -> u32 {
use futures_util::task::noop_waker_ref;
use std::future::Future;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll::*;
assert!(task.run(From::from(&NOOP_SCHEDULE)).is_none());
// Find the task that completed
TASKS.with(|c| {
let mut map = c.borrow_mut();
let mut num = None;
for (_, join) in map.iter_mut() {
let mut cx = Context::from_waker(noop_waker_ref());
match Pin::new(join).poll(&mut cx) {
Ready(n) => {
num = Some(n.unwrap());
break;
}
_ => {}
}
}
let num = num.expect("no task completed");
map.remove(&num);
num
})
}
fn is_ok<T, E>(r: Result<T, E>) {
assert!(r.is_ok())
}
tokio/src/executor/thread_pool/tests/worker.rs
+68
View File
@@ -0,0 +1,68 @@
use crate::executor::tests::track_drop::track_drop;
use crate::executor::thread_pool;
use tokio_test::assert_ok;
macro_rules! pool {
(2) => {{
let (pool, mut w, mock_park) = pool!(!2);
(pool, w.remove(0), w.remove(0), mock_park)
}};
(! $n:expr) => {{
let mut mock_park = crate::executor::tests::mock_park::MockPark::new();
let blocking = std::sync::Arc::new(crate::executor::blocking::Pool::default());
let (pool, workers) =
thread_pool::create_pool($n, |index| mock_park.mk_park(index), blocking);
(pool, workers, mock_park)
}};
}
macro_rules! enter {
($w:expr, $expr:expr) => {{
$w.enter(move || $expr);
}};
}
#[test]
fn execute_single_task() {
use std::sync::mpsc;
let (p, mut w0, _w1, ..) = pool!(2);
let (tx, rx) = mpsc::channel();
enter!(w0, p.spawn_background(async move { tx.send(1).unwrap() }));
w0.tick();
assert_ok!(rx.try_recv());
}
#[test]
fn task_migrates() {
use std::sync::mpsc;
use crate::sync::oneshot;
let (p, mut w0, mut w1, ..) = pool!(2);
let (tx1, rx1) = oneshot::channel();
let (tx2, rx2) = mpsc::channel();
let (task, did_drop) = track_drop(async move {
let msg = rx1.await.unwrap();
tx2.send(msg).unwrap();
});
enter!(w0, p.spawn_background(task));
w0.tick();
w1.enter(|| tx1.send("hello").unwrap());
w1.tick();
assert_ok!(rx2.try_recv());
// Future drops immediately even though the underlying task is not freed
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
// Tick the spawning worker in order to free memory
w0.tick();
}
tokio/src/executor/thread_pool/worker.rs
+415
View File
@@ -0,0 +1,415 @@
use crate::executor::loom::sync::Arc;
use crate::executor::park::{Park, Unpark};
use crate::executor::task::Task;
use crate::executor::thread_pool::{current, Owned, Shared};
use std::time::Duration;
// TODO: remove this re-export
pub(super) use crate::executor::thread_pool::set::Set;
pub(crate) struct Worker<P: Park + 'static> {
/// Entry in the set of workers.
entry: Entry<P::Unpark>,
/// Park the thread
park: P,
}
pub(crate) fn create_set<F, P>(
pool_size: usize,
mk_park: F,
blocking: Arc<crate::executor::blocking::Pool>,
) -> (Arc<Set<P::Unpark>>, Vec<Worker<P>>)
where
P: Send + Park,
F: FnMut(usize) -> P,
{
// Create the parks...
let parks: Vec<_> = (0..pool_size).map(mk_park).collect();
let mut pool = Arc::new(Set::new(pool_size, |i| parks[i].unpark(), blocking));
// Establish the circular link between the individual worker state
// structure and the container.
Arc::get_mut(&mut pool).unwrap().set_container_ptr();
// This will contain each worker.
let workers = parks
.into_iter()
.enumerate()
.map(|(index, park)| {
// unsafe is safe because we call Worker::new only once with each index in the pool
unsafe { Worker::new(pool.clone(), index, park) }
})
.collect();
(pool, workers)
}
/// After how many ticks is the global queue polled. This helps to ensure
/// fairness.
///
/// The number is fairly arbitrary. I believe this value was copied from golang.
const GLOBAL_POLL_INTERVAL: u16 = 61;
impl<P> Worker<P>
where
P: Send + Park,
{
// unsafe because new may only be called once for each index in pool's set
pub(super) unsafe fn new(pool: Arc<Set<P::Unpark>>, index: usize, park: P) -> Self {
Worker {
entry: Entry::new(pool, index),
park,
}
}
pub(super) fn run(mut self) {
let pool = Arc::clone(&self.entry.pool);
let pool = &pool;
let index = self.entry.index;
let mut executor = &**pool;
let entry = &mut self.entry;
let park = &mut self.park;
let blocking = &executor.blocking;
// Track the current worker
current::set(&pool, index, || {
let _enter = crate::executor::enter().expect("executor already running on thread");
crate::executor::with_default(&mut executor, || {
crate::executor::blocking::with_pool(blocking, || entry.run(park))
})
});
}
pub(super) fn id(&self) -> usize {
self.entry.index
}
#[cfg(test)]
#[allow(warnings)]
pub(crate) fn enter<F, R>(&self, f: F) -> R
where
F: FnOnce() -> R,
{
current::set(&self.entry.pool, self.entry.index, f)
}
#[cfg(test)]
#[allow(warnings)]
pub(crate) fn tick(&mut self) {
self.entry.tick(&mut self.park);
}
}
struct Entry<P: 'static> {
pool: Arc<Set<P>>,
index: usize,
}
impl<P> Entry<P>
where
P: Unpark,
{
// unsafe because Entry::owned assumes there is only one instance of the Entry
unsafe fn new(pool: Arc<Set<P>>, index: usize) -> Self {
Entry { pool, index }
}
fn run(&mut self, park: &mut impl Park<Unpark = P>) {
while self.is_running() {
if self.tick(park) {
self.park(park);
}
}
self.shutdown(park);
}
fn is_running(&mut self) -> bool {
self.owned().is_running.get()
}
/// Returns `true` if the worker needs to park
fn tick(&mut self, park: &mut impl Park<Unpark = P>) -> bool {
// Process all pending tasks in the local queue.
if !self.process_local_queue(park) {
return false;
}
// No more **local** work to process, try transitioning to searching
// in order to attempt to steal work from other workers.
//
// On `false`, the worker has entered the parked state
if self.transition_to_searching() {
// If `true` then work was found
if self.search_for_work() {
return false;
}
}
true
}
/// Process all pending tasks in the local queue, occasionally checking the
/// global queue, but never other worker local queues.
///
/// Returns `false` if processing was interrupted due to the pool shutting
/// down.
fn process_local_queue(&mut self, park: &mut impl Park<Unpark = P>) -> bool {
debug_assert!(self.is_running());
loop {
let tick = self.tick_fetch_inc();
let task = if tick % GLOBAL_POLL_INTERVAL == 0 {
// Sleep light...
self.park_light(park);
// Perform regularly scheduled maintenance work.
self.maintenance();
if !self.is_running() {
return false;
}
// Check the global queue
self.owned().work_queue.pop_global_first()
} else {
self.owned().work_queue.pop_local_first()
};
if let Some(task) = task {
self.run_task(task);
} else {
return true;
}
}
}
fn steal_work(&mut self) -> Option<Task<Shared<P>>> {
let num_workers = self.pool.len();
let start = self.owned().rand.fastrand_n(num_workers as u32);
self.owned()
.work_queue
.steal(start as usize)
// Fallback on checking the local queue, which will also check the
// injector.
.or_else(|| self.owned().work_queue.pop_global_first())
}
/// Runs maintenance work such as free pending tasks and check the pool's
/// state.
fn maintenance(&mut self) {
// Free any completed tasks
self.drain_tasks_pending_drop();
// Update the pool state cache
let closed = self.owned().work_queue.is_closed();
self.owned().is_running.set(!closed)
}
fn search_for_work(&mut self) -> bool {
debug_assert!(self.is_searching());
if let Some(task) = self.steal_work() {
self.run_task(task);
true
} else {
// Perform some routine work
self.drain_tasks_pending_drop();
false
}
}
fn transition_to_searching(&mut self) -> bool {
if self.is_searching() {
return true;
}
let ret = self.set().idle().transition_worker_to_searching();
self.owned().is_searching.set(ret);
ret
}
fn transition_from_searching(&mut self) {
debug_assert!(self.is_searching());
self.owned().is_searching.set(false);
if self.set().idle().transition_worker_from_searching() {
// We are the final searching worker. Because work was found, we
// need to notify another worker.
self.set().notify_work();
}
}
/// Returns `true` if the worker must check for any work.
fn transition_to_parked(&mut self) -> bool {
let idx = self.index;
let is_searching = self.is_searching();
let ret = self
.set()
.idle()
.transition_worker_to_parked(idx, is_searching);
// The worker is no longer searching. Setting this is the local cache
// only.
self.owned().is_searching.set(false);
// When tasks are submitted locally (from the parker), defer any
// notifications in hopes that the curent worker will grab those tasks.
self.owned().defer_notification.set(true);
ret
}
/// Returns `true` if the transition happened.
fn transition_from_parked(&mut self) -> bool {
if self.owned().did_submit_task.get() || !self.is_running() {
// Remove the worker from the sleep set.
self.set().idle().unpark_worker_by_id(self.index);
self.owned().is_searching.set(true);
self.owned().defer_notification.set(false);
true
} else {
let ret = !self.set().idle().is_parked(self.index);
if ret {
self.owned().is_searching.set(true);
self.owned().defer_notification.set(false);
}
ret
}
}
fn run_task(&mut self, task: Task<Shared<P>>) {
if self.is_searching() {
self.transition_from_searching();
}
if let Some(task) = task.run(self.shared().into()) {
self.owned().submit_local_yield(task);
self.set().notify_work();
}
}
fn final_work_sweep(&mut self) {
if !self.owned().work_queue.is_empty() {
self.set().notify_work();
}
}
fn park(&mut self, park: &mut impl Park<Unpark = P>) {
if self.transition_to_parked() {
// We are the final searching worker, check if any work arrived
// before parking
self.final_work_sweep();
}
// The state has been transitioned to parked, we can now wait by
// calling the parker. This is done in a loop as spurious wakeups are
// permitted.
loop {
park.park().ok().expect("park failed");
// We might have been woken to clean up a dropped task
self.maintenance();
if self.transition_from_parked() {
return;
}
}
}
fn park_light(&mut self, park: &mut impl Park<Unpark = P>) {
// When tasks are submitted locally (from the parker), defer any
// notifications in hopes that the curent worker will grab those tasks.
self.owned().defer_notification.set(true);
park.park_timeout(Duration::from_millis(0))
.ok()
.expect("park failed");
self.owned().defer_notification.set(false);
if self.owned().did_submit_task.get() {
self.set().notify_work();
self.owned().did_submit_task.set(false)
}
}
fn drain_tasks_pending_drop(&mut self) {
for task in self.shared().pending_drop.drain() {
unsafe {
let owned = &mut *self.set().owned()[self.index].get();
owned.release_task(&task);
}
drop(task);
}
}
/// Shutdown the worker.
///
/// Once the shutdown flag has been observed, it is guaranteed that no
/// further tasks may be pushed into the global queue.
fn shutdown(&mut self, park: &mut impl Park<Unpark = P>) {
// Transition all tasks owned by the worker to canceled.
self.owned().owned_tasks.shutdown();
// First, drain all tasks from both the local & global queue.
while let Some(task) = self.owned().work_queue.pop_local_first() {
task.shutdown();
}
// Notify all workers in case they have pending tasks to drop
//
// Not super efficient, but we are also shutting down.
self.pool.notify_all();
// The worker can only shutdown once there are no further owned tasks.
while !self.owned().owned_tasks.is_empty() {
// Wait until task that this worker owns are released.
//
// `transition_to_parked` is not called as we are not working
// anymore. When a task is released, the owning worker is unparked
// directly.
park.park().ok().expect("park failed");
// Try draining more tasks
self.drain_tasks_pending_drop();
}
}
/// Increment the tick, returning the value from before the increment.
fn tick_fetch_inc(&mut self) -> u16 {
let tick = self.owned().tick.get();
self.owned().tick.set(tick.wrapping_add(1));
tick
}
fn is_searching(&mut self) -> bool {
self.owned().is_searching.get()
}
fn set(&self) -> &Set<P> {
&self.pool
}
fn shared(&self) -> &Shared<P> {
&self.set().shared()[self.index]
}
fn owned(&mut self) -> &Owned<P> {
// safety: we own the slot
unsafe { &*self.set().owned()[self.index].get() }
}
}
tokio/src/executor/typed.rs
+178
View File
@@ -0,0 +1,178 @@
use crate::executor::SpawnError;
/// A value that spawns futures of a specific type.
///
/// The trait is generic over `T`: the type of future that can be spawened. This
/// is useful for implementing an executor that is only able to spawn a specific
/// type of future.
///
/// The [`spawn`] function is used to submit the future to the executor. Once
/// submitted, the executor takes ownership of the future and becomes
/// responsible for driving the future to completion.
///
/// This trait is useful as a bound for applications and libraries in order to
/// be generic over futures that are `Send` vs. `!Send`.
///
/// # Examples
///
/// Consider a function that provides an API for draining a `Stream` in the
/// background. To do this, a task must be spawned to perform the draining. As
/// such, the function takes a stream and an executor on which the background
/// task is spawned.
///
/// [`spawn`]: TypedExecutor::spawn
/// ```
/// use tokio::executor::TypedExecutor;
/// use tokio::sync::oneshot;
///
/// use futures_core::{ready, Stream};
/// use std::future::Future;
/// use std::pin::Pin;
/// use std::task::{Context, Poll};
///
/// async fn drain<T, E>(stream: T, executor: &mut E)
/// where
/// T: Stream + Unpin,
/// E: TypedExecutor<Drain<T>>
/// {
/// let (tx, rx) = oneshot::channel();
///
/// executor.spawn(Drain {
/// stream,
/// tx: Some(tx),
/// }).unwrap();
///
/// rx.await.unwrap()
/// }
///
/// // The background task
/// pub struct Drain<T> {
/// stream: T,
/// tx: Option<oneshot::Sender<()>>,
/// }
///
/// impl<T: Stream + Unpin> Future for Drain<T> {
/// type Output = ();
///
/// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
/// loop {
/// let item = ready!(
/// Pin::new(&mut self.stream).poll_next(cx)
/// );
///
/// if item.is_none() { break; }
/// }
///
/// let _ = self.tx.take().unwrap().send(()).map_err(|_| ());
/// Poll::Ready(())
/// }
/// }
/// ```
///
/// By doing this, the `drain` fn can accept a stream that is `!Send` as long as
/// the supplied executor is able to spawn `!Send` types.
pub trait TypedExecutor<T> {
/// Spawns a future to run on this executor.
///
/// `future` is passed to the executor, which will begin running it. The
/// executor takes ownership of the future and becomes responsible for
/// driving the future to completion.
///
/// # Panics
///
/// Implementations are encouraged to avoid panics. However, panics are
/// permitted and the caller should check the implementation specific
/// documentation for more details on possible panics.
///
/// # Examples
///
/// ```rust
/// use tokio::executor::TypedExecutor;
///
/// use std::future::Future;
/// use std::pin::Pin;
/// use std::task::{Context, Poll};
///
/// fn example<T>(my_executor: &mut T)
/// where
/// T: TypedExecutor<MyFuture>,
/// {
/// my_executor.spawn(MyFuture).unwrap();
/// }
///
/// struct MyFuture;
///
/// impl Future for MyFuture {
/// type Output = ();
///
/// fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
/// println!("running on the executor");
/// Poll::Ready(())
/// }
/// }
/// ```
fn spawn(&mut self, future: T) -> Result<(), SpawnError>;
/// Provides a best effort **hint** to whether or not `spawn` will succeed.
///
/// This function may return both false positives **and** false negatives.
/// If `status` returns `Ok`, then a call to `spawn` will *probably*
/// succeed, but may fail. If `status` returns `Err`, a call to `spawn` will
/// *probably* fail, but may succeed.
///
/// This allows a caller to avoid creating the task if the call to `spawn`
/// has a high likelihood of failing.
///
/// # Panics
///
/// This function must not panic. Implementers must ensure that panics do
/// not happen.
///
/// # Examples
///
/// ```rust
/// use tokio::executor::TypedExecutor;
///
/// use std::future::Future;
/// use std::pin::Pin;
/// use std::task::{Context, Poll};
///
/// fn example<T>(my_executor: &mut T)
/// where
/// T: TypedExecutor<MyFuture>,
/// {
/// if my_executor.status().is_ok() {
/// my_executor.spawn(MyFuture).unwrap();
/// } else {
/// println!("the executor is not in a good state");
/// }
/// }
///
/// struct MyFuture;
///
/// impl Future for MyFuture {
/// type Output = ();
///
/// fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
/// println!("running on the executor");
/// Poll::Ready(())
/// }
/// }
/// ```
fn status(&self) -> Result<(), SpawnError> {
Ok(())
}
}
impl<E, T> TypedExecutor<T> for Box<E>
where
E: TypedExecutor<T>,
{
fn spawn(&mut self, future: T) -> Result<(), SpawnError> {
(**self).spawn(future)
}
fn status(&self) -> Result<(), SpawnError> {
(**self).status()
}
}
tokio/src/executor/util/mod.rs
+5
View File
@@ -0,0 +1,5 @@
mod pad;
mod rand;
pub(crate) use self::pad::CachePadded;
pub(crate) use self::rand::FastRand;
tokio/src/executor/util/pad.rs
+52
View File
@@ -0,0 +1,52 @@
use core::fmt;
use core::ops::{Deref, DerefMut};
#[derive(Clone, Copy, Default, Hash, PartialEq, Eq)]
// Starting from Intel's Sandy Bridge, spatial prefetcher is now pulling pairs of 64-byte cache
// lines at a time, so we have to align to 128 bytes rather than 64.
//
// Sources:
// - https://www.intel.com/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-optimization-manual.pdf
// - https://github.com/facebook/folly/blob/1b5288e6eea6df074758f877c849b6e73bbb9fbb/folly/lang/Align.h#L107
#[cfg_attr(target_arch = "x86_64", repr(align(128)))]
#[cfg_attr(not(target_arch = "x86_64"), repr(align(64)))]
pub(crate) struct CachePadded<T> {
value: T,
}
unsafe impl<T: Send> Send for CachePadded<T> {}
unsafe impl<T: Sync> Sync for CachePadded<T> {}
impl<T> CachePadded<T> {
pub(crate) fn new(t: T) -> CachePadded<T> {
CachePadded::<T> { value: t }
}
}
impl<T> Deref for CachePadded<T> {
type Target = T;
fn deref(&self) -> &T {
&self.value
}
}
impl<T> DerefMut for CachePadded<T> {
fn deref_mut(&mut self) -> &mut T {
&mut self.value
}
}
impl<T: fmt::Debug> fmt::Debug for CachePadded<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CachePadded")
.field("value", &self.value)
.finish()
}
}
impl<T> From<T> for CachePadded<T> {
fn from(t: T) -> Self {
CachePadded::new(t)
}
}
tokio/src/executor/util/rand.rs
+52
View File
@@ -0,0 +1,52 @@
use std::cell::Cell;
/// Fast random number generate
///
/// Implement xorshift64+: 2 32-bit xorshift sequences added together.
/// Shift triplet [17,7,16] was calculated as indicated in Marsaglia's
/// Xorshift paper: https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf
/// This generator passes the SmallCrush suite, part of TestU01 framework:
/// http://simul.iro.umontreal.ca/testu01/tu01.html
#[derive(Debug)]
pub(crate) struct FastRand {
one: Cell<u32>,
two: Cell<u32>,
}
impl FastRand {
/// Initialize a new, thread-local, fast random number generator.
pub(crate) fn new(seed: u64) -> FastRand {
let one = (seed >> 32) as u32;
let mut two = seed as u32;
if two == 0 {
// This value cannot be zero
two = 1;
}
FastRand {
one: Cell::new(one),
two: Cell::new(two),
}
}
pub(crate) fn fastrand_n(&self, n: u32) -> u32 {
// This is similar to fastrand() % n, but faster.
// See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
let mul = (self.fastrand() as u64).wrapping_mul(n as u64);
(mul >> 32) as u32
}
fn fastrand(&self) -> u32 {
let mut s1 = self.one.get();
let s0 = self.two.get();
s1 ^= s1 << 17;
s1 = s1 ^ s0 ^ s1 >> 7 ^ s0 >> 16;
self.one.set(s0);
self.two.set(s1);
s0.wrapping_add(s1)
}
}
tokio/src/fs/mod.rs
+2 -5
View File
@@ -18,12 +18,9 @@
//! type. Adaptions also extend to traits like `std::io::Read` where methods
//! return `std::io::Result`. Be warned that these adapted methods may return
//! `std::io::ErrorKind::WouldBlock` if a *worker* thread can not be converted
//! to a *backup* thread immediately. See [tokio-executor] for more details
//! of the threading model and [`blocking`].
//! to a *backup* thread immediately.
//!
//! [`blocking`]: https://docs.rs/tokio-executor/0.2.0-alpha.2/tokio_executor/threadpool/fn.blocking.html
//! [`AsyncRead`]: https://docs.rs/tokio-io/0.1/tokio_io/trait.AsyncRead.html
//! [tokio-executor]: https://docs.rs/tokio-executor/0.2.0-alpha.2/tokio_executor/threadpool/index.html
pub(crate) mod blocking;
@@ -94,5 +91,5 @@ where
mod sys {
pub(crate) use std::fs::File;
pub(crate) use tokio_executor::blocking::{run, Blocking};
pub(crate) use crate::executor::blocking::{run, Blocking};
}
tokio/src/lib.rs
+7 -1
View File
@@ -79,6 +79,11 @@ macro_rules! if_runtime {
)*)
}
#[cfg(all(loom, test))]
macro_rules! thread_local {
($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } }
}
#[cfg(feature = "timer")]
pub mod clock;
@@ -101,10 +106,11 @@ mod loom;
pub mod prelude;
#[cfg(feature = "process")]
#[cfg(all(feature = "process", not(loom)))]
pub mod process;
#[cfg(feature = "signal")]
#[cfg(not(loom))]
pub mod signal;
pub mod stream;
tokio/src/net/addr.rs
+2 -2
View File
@@ -1,4 +1,4 @@
use tokio_executor::blocking;
use crate::executor::blocking;
use futures_util::future;
use std::io;
@@ -143,7 +143,7 @@ pub(crate) mod sealed {
//! part of the `ToSocketAddrs` public API. The details will change over
//! time.
use tokio_executor::blocking::Blocking;
use crate::executor::blocking::Blocking;
use futures_core::ready;
use std::future::Future;
tokio/src/net/driver/mod.rs
+2 -2
View File
@@ -124,12 +124,12 @@
//! [`PollEvented`]: struct.PollEvented.html
//! [`std::io::Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
//! [`std::io::Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
#[cfg(loom)]
#[cfg(all(loom, test))]
macro_rules! loom_thread_local {
($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } }
}
#[cfg(not(loom))]
#[cfg(any(not(loom), not(test)))]
macro_rules! loom_thread_local {
($($tts:tt)+) => { std::thread_local!{ $($tts)+ } }
}
tokio/src/net/driver/reactor/mod.rs
+1 -1
View File
@@ -1,4 +1,5 @@
use super::platform;
use crate::executor::park::{Park, Unpark};
use crate::loom::atomic::{AtomicUsize, Ordering::SeqCst};
mod dispatch;
@@ -15,7 +16,6 @@ use std::sync::{Arc, Weak};
use std::task::Waker;
use std::time::Duration;
use std::{fmt, usize};
use tokio_executor::park::{Park, Unpark};
/// The core reactor, or event loop.
///
tokio/src/runtime/current_thread/builder.rs
+1 -2
View File
@@ -1,10 +1,9 @@
use crate::executor::current_thread::CurrentThread;
use crate::net::driver::Reactor;
use crate::runtime::current_thread::Runtime;
use crate::timer::clock::Clock;
use crate::timer::timer::Timer;
use tokio_executor::current_thread::CurrentThread;
use std::io;
/// Builds a Single-threaded runtime with custom configuration values.
tokio/src/runtime/current_thread/mod.rs
+2 -2
View File
@@ -63,5 +63,5 @@ mod runtime;
pub use self::builder::Builder;
pub use self::runtime::{Handle, Runtime, RunError};
pub use tokio_executor::current_thread::spawn;
pub use tokio_executor::current_thread::TaskExecutor;
pub use crate::executor::current_thread::spawn;
pub use crate::executor::current_thread::TaskExecutor;
tokio/src/runtime/current_thread/runtime.rs
+5 -6
View File
@@ -1,11 +1,10 @@
use crate::executor::current_thread::Handle as ExecutorHandle;
use crate::executor::current_thread::{self, CurrentThread};
use crate::net::driver::{self, Reactor};
use crate::runtime::current_thread::Builder;
use crate::timer::clock::{self, Clock};
use crate::timer::timer::{self, Timer};
use tokio_executor::current_thread::Handle as ExecutorHandle;
use tokio_executor::current_thread::{self, CurrentThread};
use std::error::Error;
use std::fmt;
use std::future::Future;
@@ -38,7 +37,7 @@ impl Handle {
///
/// This function panics if the spawn fails. Failure occurs if the `CurrentThread`
/// instance of the `Handle` does not exist anymore.
pub fn spawn<F>(&self, future: F) -> Result<(), tokio_executor::SpawnError>
pub fn spawn<F>(&self, future: F) -> Result<(), crate::executor::SpawnError>
where
F: Future<Output = ()> + Send + 'static,
{
@@ -54,7 +53,7 @@ impl Handle {
///
/// This allows a caller to avoid creating the task if the call to `spawn`
/// has a high likelihood of failing.
pub fn status(&self) -> Result<(), tokio_executor::SpawnError> {
pub fn status(&self) -> Result<(), crate::executor::SpawnError> {
self.0.status()
}
}
@@ -201,7 +200,7 @@ impl Runtime {
// to run the provided future, another to install as the default
// one). We use the fake one here as the default one.
let mut default_executor = current_thread::TaskExecutor::current();
tokio_executor::with_default(&mut default_executor, || f(executor))
crate::executor::with_default(&mut default_executor, || f(executor))
})
}
}
tokio/src/runtime/mod.rs
+2 -2
View File
@@ -19,7 +19,7 @@
//! Creating a [`Runtime`] does the following:
//!
//! * Spawn a background thread running a [`Reactor`] instance.
//! * Start a [`ThreadPool`] for executing futures.
//! * Start a thread pool for executing futures.
//! * Run an instance of `Timer` **per** thread pool worker thread.
//!
//! The thread pool uses a work-stealing strategy and is configured to start a
@@ -124,12 +124,12 @@
//! [timer]: ../timer/index.html
//! [`Runtime`]: struct.Runtime.html
//! [`Reactor`]: ../reactor/struct.Reactor.html
//! [`ThreadPool`]: https://docs.rs/tokio-executor/0.2.0-alpha.2/tokio_executor/threadpool/struct.ThreadPool.html
//! [`run`]: fn.run.html
//! [`tokio::spawn`]: ../executor/fn.spawn.html
//! [`tokio::main`]: ../../tokio_macros/attr.main.html
pub mod current_thread;
#[cfg(feature = "rt-full")]
mod threadpool;
tokio/src/runtime/threadpool/builder.rs
+1 -2
View File
@@ -1,10 +1,9 @@
use crate::executor::thread_pool;
use crate::net::driver::{self, Reactor};
use crate::runtime::threadpool::{Inner, Runtime};
use crate::timer::clock::{self, Clock};
use crate::timer::timer::{self, Timer};
use tokio_executor::thread_pool;
use std::sync::{Arc, Mutex};
use std::{fmt, io};
tokio/src/runtime/threadpool/mod.rs
+2 -2
View File
@@ -7,12 +7,12 @@ mod spawner;
pub use self::spawner::Spawner;
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
pub use tokio_executor::thread_pool::JoinHandle;
pub use crate::executor::thread_pool::JoinHandle;
use crate::net::driver;
use crate::timer::timer;
use tokio_executor::thread_pool::ThreadPool;
use crate::executor::thread_pool::ThreadPool;
use std::future::Future;
use std::io;
tokio/src/runtime/threadpool/spawner.rs
+1 -2
View File
@@ -1,7 +1,6 @@
use crate::executor::thread_pool;
use crate::runtime::JoinHandle;
use tokio_executor::thread_pool;
use std::future::Future;
/// Spawns futures on the runtime
tokio/src/timer/timer/mod.rs
+1 -2
View File
@@ -41,13 +41,12 @@ pub(crate) use self::registration::Registration;
mod stack;
use self::stack::Stack;
use crate::executor::park::{Park, ParkThread, Unpark};
use crate::timer::atomic::AtomicU64;
use crate::timer::clock::Clock;
use crate::timer::wheel;
use crate::timer::Error;
use tokio_executor::park::{Park, ParkThread, Unpark};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::Arc;
tokio/tests/current_thread.rs
+781
View File
@@ -0,0 +1,781 @@
#![warn(rust_2018_idioms)]
#![cfg(not(miri))]
use tokio::executor::current_thread::{self, block_on_all, CurrentThread, TaskExecutor};
use tokio::executor::TypedExecutor;
use tokio::sync::oneshot;
use std::any::Any;
use std::cell::{Cell, RefCell};
use std::future::Future;
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Context, Poll};
use std::thread;
use std::time::Duration;
mod from_block_on_all {
use super::*;
fn test<F: Fn(Pin<Box<dyn Future<Output = ()>>>) + 'static>(spawn: F) {
let cnt = Rc::new(Cell::new(0));
let c = cnt.clone();
let msg = block_on_all(async move {
c.set(1 + c.get());
// Spawn!
spawn(Box::pin(async move {
c.set(1 + c.get());
}));
"hello"
});
assert_eq!(2, cnt.get());
assert_eq!(msg, "hello");
}
#[test]
fn spawn() {
test(current_thread::spawn)
}
#[test]
fn execute() {
test(|f| {
TaskExecutor::current().spawn(f).unwrap();
});
}
}
#[test]
fn block_waits() {
let (tx, rx) = oneshot::channel();
thread::spawn(|| {
thread::sleep(Duration::from_millis(1000));
tx.send(()).unwrap();
});
let cnt = Rc::new(Cell::new(0));
let cnt2 = cnt.clone();
block_on_all(async move {
rx.await.unwrap();
cnt.set(1 + cnt.get());
});
assert_eq!(1, cnt2.get());
}
#[test]
fn spawn_many() {
const ITER: usize = 200;
let cnt = Rc::new(Cell::new(0));
let mut tokio_current_thread = CurrentThread::new();
for _ in 0..ITER {
let cnt = cnt.clone();
tokio_current_thread.spawn(async move {
cnt.set(1 + cnt.get());
});
}
tokio_current_thread.run().unwrap();
assert_eq!(cnt.get(), ITER);
}
mod does_not_set_global_executor_by_default {
use super::*;
fn test<F: Fn(Pin<Box<dyn Future<Output = ()> + Send>>) -> Result<(), E> + 'static, E>(
spawn: F,
) {
block_on_all(async {
spawn(Box::pin(async {})).unwrap_err();
});
}
#[test]
fn spawn() {
test(|f| tokio::executor::DefaultExecutor::current().spawn(f))
}
}
mod from_block_on_future {
use super::*;
fn test<F: Fn(Pin<Box<dyn Future<Output = ()>>>)>(spawn: F) {
let cnt = Rc::new(Cell::new(0));
let cnt2 = cnt.clone();
let mut tokio_current_thread = CurrentThread::new();
tokio_current_thread.block_on(async move {
let cnt3 = cnt2.clone();
spawn(Box::pin(async move {
cnt3.set(1 + cnt3.get());
}));
});
tokio_current_thread.run().unwrap();
assert_eq!(1, cnt.get());
}
#[test]
fn spawn() {
test(current_thread::spawn);
}
#[test]
fn execute() {
test(|f| {
current_thread::TaskExecutor::current().spawn(f).unwrap();
});
}
}
mod outstanding_tasks_are_dropped_when_executor_is_dropped {
use super::*;
#[allow(unreachable_code)] // TODO: remove this when https://github.com/rust-lang/rust/issues/64636 fixed.
async fn never(_rc: Rc<()>) {
loop {
yield_once().await;
}
}
fn test<F, G>(spawn: F, dotspawn: G)
where
F: Fn(Pin<Box<dyn Future<Output = ()>>>) + 'static,
G: Fn(&mut CurrentThread, Pin<Box<dyn Future<Output = ()>>>),
{
let mut rc = Rc::new(());
let mut tokio_current_thread = CurrentThread::new();
dotspawn(&mut tokio_current_thread, Box::pin(never(rc.clone())));
drop(tokio_current_thread);
// Ensure the daemon is dropped
assert!(Rc::get_mut(&mut rc).is_some());
// Using the global spawn fn
let mut rc = Rc::new(());
let rc2 = rc.clone();
let mut tokio_current_thread = CurrentThread::new();
tokio_current_thread.block_on(async move {
spawn(Box::pin(never(rc2)));
});
drop(tokio_current_thread);
// Ensure the daemon is dropped
assert!(Rc::get_mut(&mut rc).is_some());
}
#[test]
fn spawn() {
test(current_thread::spawn, |rt, f| {
rt.spawn(f);
})
}
#[test]
fn execute() {
test(
|f| {
current_thread::TaskExecutor::current().spawn(f).unwrap();
},
// Note: `CurrentThread` doesn't currently implement
// `futures::Executor`, so we'll call `.spawn(...)` rather than
// `.execute(...)` for now. If `CurrentThread` is changed to
// implement Executor, change this to `.execute(...).unwrap()`.
|rt, f| {
rt.spawn(f);
},
);
}
}
#[test]
#[should_panic]
fn nesting_run() {
block_on_all(async {
block_on_all(async {});
});
}
mod run_in_future {
use super::*;
#[test]
#[should_panic]
fn spawn() {
block_on_all(async {
current_thread::spawn(async {
block_on_all(async {});
});
});
}
#[test]
#[should_panic]
fn execute() {
block_on_all(async {
current_thread::TaskExecutor::current()
.spawn(async {
block_on_all(async {});
})
.unwrap();
});
}
}
#[test]
fn tick_on_infini_future() {
let num = Rc::new(Cell::new(0));
#[allow(unreachable_code)] // TODO: remove this when https://github.com/rust-lang/rust/issues/64636 fixed.
async fn infini(num: Rc<Cell<usize>>) {
loop {
num.set(1 + num.get());
yield_once().await
}
}
CurrentThread::new()
.spawn(infini(num.clone()))
.turn(None)
.unwrap();
assert_eq!(1, num.get());
}
mod tasks_are_scheduled_fairly {
use super::*;
#[allow(unreachable_code)] // TODO: remove this when https://github.com/rust-lang/rust/issues/64636 fixed.
async fn spin(state: Rc<RefCell<[i32; 2]>>, idx: usize) {
loop {
// borrow_mut scope
{
let mut state = state.borrow_mut();
if idx == 0 {
let diff = state[0] - state[1];
assert!(diff.abs() <= 1);
if state[0] >= 50 {
return;
}
}
state[idx] += 1;
if state[idx] >= 100 {
return;
}
}
yield_once().await;
}
}
fn test<F: Fn(Pin<Box<dyn Future<Output = ()>>>)>(spawn: F) {
let state = Rc::new(RefCell::new([0, 0]));
block_on_all(async move {
spawn(Box::pin(spin(state.clone(), 0)));
spawn(Box::pin(spin(state, 1)));
});
}
#[test]
fn spawn() {
test(current_thread::spawn)
}
#[test]
fn execute() {
test(|f| {
current_thread::TaskExecutor::current().spawn(f).unwrap();
})
}
}
mod and_turn {
use super::*;
fn test<F, G>(spawn: F, dotspawn: G)
where
F: Fn(Pin<Box<dyn Future<Output = ()>>>) + 'static,
G: Fn(&mut CurrentThread, Pin<Box<dyn Future<Output = ()>>>),
{
let cnt = Rc::new(Cell::new(0));
let c = cnt.clone();
let mut tokio_current_thread = CurrentThread::new();
// Spawn a basic task to get the executor to turn
dotspawn(&mut tokio_current_thread, Box::pin(async {}));
// Turn once...
tokio_current_thread.turn(None).unwrap();
dotspawn(
&mut tokio_current_thread,
Box::pin(async move {
c.set(1 + c.get());
// Spawn!
spawn(Box::pin(async move {
c.set(1 + c.get());
}));
}),
);
// This does not run the newly spawned thread
tokio_current_thread.turn(None).unwrap();
assert_eq!(1, cnt.get());
// This runs the newly spawned thread
tokio_current_thread.turn(None).unwrap();
assert_eq!(2, cnt.get());
}
#[test]
fn spawn() {
test(current_thread::spawn, |rt, f| {
rt.spawn(f);
})
}
#[test]
fn execute() {
test(
|f| {
current_thread::TaskExecutor::current().spawn(f).unwrap();
},
// Note: `CurrentThread` doesn't currently implement
// `futures::Executor`, so we'll call `.spawn(...)` rather than
// `.execute(...)` for now. If `CurrentThread` is changed to
// implement Executor, change this to `.execute(...).unwrap()`.
|rt, f| {
rt.spawn(f);
},
);
}
}
mod in_drop {
use super::*;
struct OnDrop<F: FnOnce()>(Option<F>);
impl<F: FnOnce()> Drop for OnDrop<F> {
fn drop(&mut self) {
(self.0.take().unwrap())();
}
}
async fn noop(_data: Box<dyn Any>) {}
fn test<F, G>(spawn: F, dotspawn: G)
where
F: Fn(Pin<Box<dyn Future<Output = ()>>>) + 'static,
G: Fn(&mut CurrentThread, Pin<Box<dyn Future<Output = ()>>>),
{
let mut tokio_current_thread = CurrentThread::new();
let (tx, rx) = oneshot::channel();
dotspawn(
&mut tokio_current_thread,
Box::pin(noop(Box::new(OnDrop(Some(move || {
spawn(Box::pin(async move {
tx.send(()).unwrap();
}));
}))))),
);
tokio_current_thread.block_on(rx).unwrap();
tokio_current_thread.run().unwrap();
}
#[test]
fn spawn() {
test(current_thread::spawn, |rt, f| {
rt.spawn(f);
})
}
#[test]
fn execute() {
test(
|f| {
current_thread::TaskExecutor::current().spawn(f).unwrap();
},
// Note: `CurrentThread` doesn't currently implement
// `futures::Executor`, so we'll call `.spawn(...)` rather than
// `.execute(...)` for now. If `CurrentThread` is changed to
// implement Executor, change this to `.execute(...).unwrap()`.
|rt, f| {
rt.spawn(f);
},
);
}
}
/*
#[test]
fn hammer_turn() {
use futures::sync::mpsc;
const ITER: usize = 100;
const N: usize = 100;
const THREADS: usize = 4;
for _ in 0..ITER {
let mut ths = vec![];
// Add some jitter
for _ in 0..THREADS {
let th = thread::spawn(|| {
let mut tokio_current_thread = CurrentThread::new();
let (tx, rx) = mpsc::unbounded();
tokio_current_thread.spawn({
let cnt = Rc::new(Cell::new(0));
let c = cnt.clone();
rx.for_each(move |_| {
c.set(1 + c.get());
Ok(())
})
.map_err(|e| panic!("err={:?}", e))
.map(move |v| {
assert_eq!(N, cnt.get());
v
})
});
thread::spawn(move || {
for _ in 0..N {
tx.unbounded_send(()).unwrap();
thread::yield_now();
}
});
while !tokio_current_thread.is_idle() {
tokio_current_thread.turn(None).unwrap();
}
});
ths.push(th);
}
for th in ths {
th.join().unwrap();
}
}
}
*/
#[test]
fn turn_has_polled() {
let mut tokio_current_thread = CurrentThread::new();
// Spawn oneshot receiver
let (sender, receiver) = oneshot::channel::<()>();
tokio_current_thread.spawn(async move {
let _ = receiver.await;
});
// Turn once...
let res = tokio_current_thread
.turn(Some(Duration::from_millis(0)))
.unwrap();
// Should've polled the receiver once, but considered it not ready
assert!(res.has_polled());
// Turn another time
let res = tokio_current_thread
.turn(Some(Duration::from_millis(0)))
.unwrap();
// Should've polled nothing, the receiver is not ready yet
assert!(!res.has_polled());
// Make the receiver ready
sender.send(()).unwrap();
// Turn another time
let res = tokio_current_thread
.turn(Some(Duration::from_millis(0)))
.unwrap();
// Should've polled the receiver, it's ready now
assert!(res.has_polled());
// Now the executor should be empty
assert!(tokio_current_thread.is_idle());
let res = tokio_current_thread
.turn(Some(Duration::from_millis(0)))
.unwrap();
// So should've polled nothing
assert!(!res.has_polled());
}
// Our own mock Park that is never really waiting and the only
// thing it does is to send, on request, something (once) to a oneshot
// channel
struct MyPark {
sender: Option<oneshot::Sender<()>>,
send_now: Rc<Cell<bool>>,
}
struct MyUnpark;
impl tokio::executor::park::Park for MyPark {
type Unpark = MyUnpark;
type Error = ();
fn unpark(&self) -> Self::Unpark {
MyUnpark
}
fn park(&mut self) -> Result<(), Self::Error> {
// If called twice with send_now, this will intentionally panic
if self.send_now.get() {
self.sender.take().unwrap().send(()).unwrap();
}
Ok(())
}
fn park_timeout(&mut self, _duration: Duration) -> Result<(), Self::Error> {
self.park()
}
}
impl tokio::executor::park::Unpark for MyUnpark {
fn unpark(&self) {}
}
#[test]
fn turn_fair() {
let send_now = Rc::new(Cell::new(false));
let (sender, receiver) = oneshot::channel::<()>();
let (sender_2, receiver_2) = oneshot::channel::<()>();
let (sender_3, receiver_3) = oneshot::channel::<()>();
let my_park = MyPark {
sender: Some(sender_3),
send_now: send_now.clone(),
};
let mut tokio_current_thread = CurrentThread::new_with_park(my_park);
let receiver_1_done = Rc::new(Cell::new(false));
let receiver_1_done_clone = receiver_1_done.clone();
// Once an item is received on the oneshot channel, it will immediately
// immediately make the second oneshot channel ready
tokio_current_thread.spawn(async move {
receiver.await.unwrap();
sender_2.send(()).unwrap();
receiver_1_done_clone.set(true);
});
let receiver_2_done = Rc::new(Cell::new(false));
let receiver_2_done_clone = receiver_2_done.clone();
tokio_current_thread.spawn(async move {
receiver_2.await.unwrap();
receiver_2_done_clone.set(true);
});
// The third receiver is only woken up from our Park implementation, it simulates
// e.g. a socket that first has to be polled to know if it is ready now
let receiver_3_done = Rc::new(Cell::new(false));
let receiver_3_done_clone = receiver_3_done.clone();
tokio_current_thread.spawn(async move {
receiver_3.await.unwrap();
receiver_3_done_clone.set(true);
});
// First turn should've polled both and considered them not ready
let res = tokio_current_thread
.turn(Some(Duration::from_millis(0)))
.unwrap();
assert!(res.has_polled());
// Next turn should've polled nothing
let res = tokio_current_thread
.turn(Some(Duration::from_millis(0)))
.unwrap();
assert!(!res.has_polled());
assert!(!receiver_1_done.get());
assert!(!receiver_2_done.get());
assert!(!receiver_3_done.get());
// After this the receiver future will wake up the second receiver future,
// so there are pending futures again
sender.send(()).unwrap();
// Now the first receiver should be done, the second receiver should be ready
// to be polled again and the socket not yet
let res = tokio_current_thread.turn(None).unwrap();
assert!(res.has_polled());
assert!(receiver_1_done.get());
assert!(!receiver_2_done.get());
assert!(!receiver_3_done.get());
// Now let our park implementation know that it should send something to sender 3
send_now.set(true);
// This should resolve the second receiver directly, but also poll the socket
// and read the packet from it. If it didn't do both here, we would handle
// futures that are woken up from the reactor and directly unfairly and would
// favour the ones that are woken up directly.
let res = tokio_current_thread.turn(None).unwrap();
assert!(res.has_polled());
assert!(receiver_1_done.get());
assert!(receiver_2_done.get());
assert!(receiver_3_done.get());
// Don't send again
send_now.set(false);
// Now we should be idle and turning should not poll anything
assert!(tokio_current_thread.is_idle());
let res = tokio_current_thread.turn(None).unwrap();
assert!(!res.has_polled());
}
#[test]
fn spawn_from_other_thread() {
let mut current_thread = CurrentThread::new();
let handle = current_thread.handle();
let (sender, receiver) = oneshot::channel::<()>();
thread::spawn(move || {
handle
.spawn(async move {
sender.send(()).unwrap();
})
.unwrap();
});
let _ = current_thread.block_on(receiver).unwrap();
}
#[test]
fn spawn_from_other_thread_unpark() {
use std::sync::mpsc::channel as mpsc_channel;
let mut current_thread = CurrentThread::new();
let handle = current_thread.handle();
let (sender_1, receiver_1) = oneshot::channel::<()>();
let (sender_2, receiver_2) = mpsc_channel::<()>();
thread::spawn(move || {
let _ = receiver_2.recv().unwrap();
handle
.spawn(async move {
sender_1.send(()).unwrap();
})
.unwrap();
});
// Ensure that unparking the executor works correctly. It will first
// check if there are new futures (there are none), then execute the
// lazy future below which will cause the future to be spawned from
// the other thread. Then the executor will park but should be woken
// up because *now* we have a new future to schedule
let _ = current_thread.block_on(async move {
// inlined 'lazy'
async move {
sender_2.send(()).unwrap();
}
.await;
receiver_1.await.unwrap();
});
}
#[test]
fn spawn_from_executor_with_handle() {
let mut current_thread = CurrentThread::new();
let handle = current_thread.handle();
let (tx, rx) = oneshot::channel();
current_thread.spawn(async move {
handle
.spawn(async move {
tx.send(()).unwrap();
})
.unwrap();
});
current_thread.block_on(rx).unwrap();
}
#[test]
fn handle_status() {
let current_thread = CurrentThread::new();
let handle = current_thread.handle();
assert!(handle.status().is_ok());
drop(current_thread);
assert!(handle.spawn(async { () }).is_err());
assert!(handle.status().is_err());
}
#[test]
fn handle_is_sync() {
let current_thread = CurrentThread::new();
let handle = current_thread.handle();
let _box: Box<dyn Sync> = Box::new(handle);
}
async fn yield_once() {
YieldOnce(false).await
}
struct YieldOnce(bool);
impl Future for YieldOnce {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if self.0 {
Poll::Ready(())
} else {
self.0 = true;
// Push to the back of the executor's queue
cx.waker().wake_by_ref();
Poll::Pending
}
}
}
tokio/tests/executor.rs
+24
View File
@@ -0,0 +1,24 @@
#![warn(rust_2018_idioms)]
use tokio::executor::DefaultExecutor;
use std::future::Future;
use std::pin::Pin;
mod out_of_executor_context {
use super::*;
use tokio::executor::Executor;
fn test<F, E>(spawn: F)
where
F: Fn(Pin<Box<dyn Future<Output = ()> + Send>>) -> Result<(), E>,
{
let res = spawn(Box::pin(async {}));
assert!(res.is_err());
}
#[test]
fn spawn() {
test(|f| DefaultExecutor::current().spawn(f));
}
}
tokio/tests/executor_enter.rs
+17
View File
@@ -0,0 +1,17 @@
#![warn(rust_2018_idioms)]
#[test]
fn block_on_ready() {
let mut enter = tokio::executor::enter().unwrap();
let val = enter.block_on(async { 123 });
assert_eq!(val, 123);
}
#[test]
fn block_on_pending() {
let mut enter = tokio::executor::enter().unwrap();
let val = enter.block_on(async { 123 });
assert_eq!(val, 123);
}
tokio/tests/executor_global.rs
+17
View File
@@ -0,0 +1,17 @@
use tokio::executor::{with_default, DefaultExecutor};
#[test]
fn default_executor_is_send_and_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<DefaultExecutor>();
}
#[test]
#[should_panic]
fn nested_default_executor_status() {
let _enter = tokio::executor::enter().unwrap();
let mut executor = DefaultExecutor::current();
let _result = with_default(&mut executor, || ());
}
tokio/tests/net_driver.rs
+1 -1
View File
@@ -61,7 +61,7 @@ fn test_drop_on_notify() {
}
}));
let _enter = tokio_executor::enter().unwrap();
let _enter = tokio::executor::enter().unwrap();
{
let handle = reactor.handle();
tokio/tests/runtime_current_thread.rs
+1 -2
View File
@@ -1,5 +1,4 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "default")]
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
@@ -133,6 +132,6 @@ fn racy() {
// wait for runtime thread to exit
jh.join().unwrap();
let mut e = tokio_executor::enter().unwrap();
let mut e = tokio::executor::enter().unwrap();
e.block_on(rx).unwrap();
}
tokio/tests/runtime_threaded.rs
+1 -4
View File
@@ -1,7 +1,5 @@
#![warn(rust_2018_idioms)]
#![cfg(feature = "default")]
use tokio;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::runtime::Runtime;
@@ -9,7 +7,6 @@ use tokio::sync::oneshot;
use tokio::timer::delay;
use tokio_test::{assert_err, assert_ok};
use env_logger;
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};
@@ -146,7 +143,7 @@ fn nested_enter() {
let rt = Runtime::new().unwrap();
rt.block_on(async {
assert_err!(tokio_executor::enter());
assert_err!(tokio::executor::enter());
let res = panic::catch_unwind(move || {
let rt = Runtime::new().unwrap();
tokio/tests/thread_pool.rs
+478
View File
@@ -0,0 +1,478 @@
#![warn(rust_2018_idioms)]
use tokio::executor::park::{Park, Unpark};
use tokio::executor::thread_pool::*;
use futures_util::future::poll_fn;
use std::cell::Cell;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::atomic::*;
use std::sync::{mpsc, Arc};
use std::task::{Context, Poll, Waker};
use std::time::Duration;
thread_local!(static FOO: Cell<u32> = Cell::new(0));
#[test]
fn shutdown_drops_futures() {
for _ in 0..1_000 {
let num_inc = Arc::new(AtomicUsize::new(0));
let num_dec = Arc::new(AtomicUsize::new(0));
let num_drop = Arc::new(AtomicUsize::new(0));
struct Never(Arc<AtomicUsize>);
impl Future for Never {
type Output = ();
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
Poll::Pending
}
}
impl Drop for Never {
fn drop(&mut self) {
self.0.fetch_add(1, Relaxed);
}
}
let a = num_inc.clone();
let b = num_dec.clone();
let mut pool = Builder::new()
.around_worker(move |_, work| {
a.fetch_add(1, Relaxed);
work();
b.fetch_add(1, Relaxed);
})
.build();
// let tx = pool.sender().clone();
pool.spawn(Never(num_drop.clone()));
// Wait for the pool to shutdown
pool.shutdown_now();
// Assert that only a single thread was spawned.
let a = num_inc.load(Relaxed);
assert!(a >= 1);
// Assert that all threads shutdown
let b = num_dec.load(Relaxed);
assert_eq!(a, b);
// Assert that the future was dropped
let c = num_drop.load(Relaxed);
assert_eq!(c, 1);
}
}
#[test]
fn drop_threadpool_drops_futures() {
const NUM_THREADS: usize = 10;
for _ in 0..1_000 {
let num_inc = Arc::new(AtomicUsize::new(0));
let num_dec = Arc::new(AtomicUsize::new(0));
let num_drop = Arc::new(AtomicUsize::new(0));
struct Never(Arc<AtomicUsize>);
impl Future for Never {
type Output = ();
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
Poll::Pending
}
}
impl Drop for Never {
fn drop(&mut self) {
self.0.fetch_add(1, Relaxed);
}
}
let a = num_inc.clone();
let b = num_dec.clone();
let pool = Builder::new()
.num_threads(NUM_THREADS)
.around_worker(move |_, work| {
a.fetch_add(1, Relaxed);
work();
b.fetch_add(1, Relaxed);
})
.build();
pool.spawn(Never(num_drop.clone()));
// Wait for the pool to shutdown
drop(pool);
// Assert that all the threads spawned
let a = num_inc.load(Relaxed);
assert_eq!(a, NUM_THREADS);
// Assert that all threads shutdown
let b = num_dec.load(Relaxed);
assert_eq!(a, b);
// Assert that the future was dropped
let c = num_drop.load(Relaxed);
assert_eq!(c, 1);
}
}
#[test]
fn many_oneshot_futures() {
// used for notifying the main thread
const NUM: usize = 10_000;
for _ in 0..50 {
let (tx, rx) = mpsc::channel();
let mut pool = new_pool();
let cnt = Arc::new(AtomicUsize::new(0));
for _ in 0..NUM {
let cnt = cnt.clone();
let tx = tx.clone();
pool.spawn(async move {
let num = cnt.fetch_add(1, Relaxed) + 1;
if num == NUM {
tx.send(()).unwrap();
}
});
}
rx.recv().unwrap();
// Wait for the pool to shutdown
pool.shutdown_now();
}
}
#[test]
fn many_multishot_futures() {
use tokio::sync::mpsc;
const CHAIN: usize = 200;
const CYCLES: usize = 5;
const TRACKS: usize = 50;
for _ in 0..50 {
let pool = new_pool();
let mut start_txs = Vec::with_capacity(TRACKS);
let mut final_rxs = Vec::with_capacity(TRACKS);
for _ in 0..TRACKS {
let (start_tx, mut chain_rx) = mpsc::channel(10);
for _ in 0..CHAIN {
let (mut next_tx, next_rx) = mpsc::channel(10);
// Forward all the messages
pool.spawn(async move {
while let Some(v) = chain_rx.recv().await {
next_tx.send(v).await.unwrap();
}
});
chain_rx = next_rx;
}
// This final task cycles if needed
let (mut final_tx, final_rx) = mpsc::channel(10);
let mut cycle_tx = start_tx.clone();
let mut rem = CYCLES;
pool.spawn(async move {
for _ in 0..CYCLES {
let msg = chain_rx.recv().await.unwrap();
rem -= 1;
if rem == 0 {
final_tx.send(msg).await.unwrap();
} else {
cycle_tx.send(msg).await.unwrap();
}
}
});
start_txs.push(start_tx);
final_rxs.push(final_rx);
}
{
let mut e = tokio::executor::enter().unwrap();
e.block_on(async move {
for mut start_tx in start_txs {
start_tx.send("ping").await.unwrap();
}
for mut final_rx in final_rxs {
final_rx.recv().await.unwrap();
}
});
}
}
}
#[test]
fn global_executor_is_configured() {
let pool = new_pool();
let (signal_tx, signal_rx) = mpsc::channel();
pool.spawn(async move {
tokio::executor::spawn(async move {
signal_tx.send(()).unwrap();
});
});
signal_rx.recv().unwrap();
}
#[test]
fn new_threadpool_is_idle() {
let mut pool = new_pool();
pool.shutdown_now();
}
#[test]
fn panic_in_task() {
let pool = new_pool();
let (tx, rx) = mpsc::channel();
struct Boom(mpsc::Sender<()>);
impl Future for Boom {
type Output = ();
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
panic!();
}
}
impl Drop for Boom {
fn drop(&mut self) {
assert!(::std::thread::panicking());
self.0.send(()).unwrap();
}
}
pool.spawn(Boom(tx));
rx.recv().unwrap();
}
#[test]
fn multi_threadpool() {
use tokio_sync::oneshot;
let pool1 = new_pool();
let pool2 = new_pool();
let (tx, rx) = oneshot::channel();
let (done_tx, done_rx) = mpsc::channel();
pool2.spawn(async move {
rx.await.unwrap();
done_tx.send(()).unwrap();
});
pool1.spawn(async move {
tx.send(()).unwrap();
});
done_rx.recv().unwrap();
}
#[test]
fn eagerly_drops_futures() {
use std::sync::{mpsc, Mutex};
struct MyPark {
rx: mpsc::Receiver<()>,
tx: Mutex<mpsc::Sender<()>>,
#[allow(dead_code)]
park_tx: mpsc::SyncSender<()>,
unpark_tx: mpsc::SyncSender<()>,
}
impl Park for MyPark {
type Unpark = MyUnpark;
type Error = ();
fn unpark(&self) -> Self::Unpark {
MyUnpark {
tx: Mutex::new(self.tx.lock().unwrap().clone()),
unpark_tx: self.unpark_tx.clone(),
}
}
fn park(&mut self) -> Result<(), Self::Error> {
let _ = self.rx.recv();
Ok(())
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
let _ = self.rx.recv_timeout(duration);
Ok(())
}
}
struct MyUnpark {
tx: Mutex<mpsc::Sender<()>>,
#[allow(dead_code)]
unpark_tx: mpsc::SyncSender<()>,
}
impl Unpark for MyUnpark {
fn unpark(&self) {
let _ = self.tx.lock().unwrap().send(());
}
}
let (task_tx, task_rx) = mpsc::channel();
let (drop_tx, drop_rx) = mpsc::channel();
let (park_tx, park_rx) = mpsc::sync_channel(0);
let (unpark_tx, unpark_rx) = mpsc::sync_channel(0);
let pool = Builder::new().num_threads(4).build_with_park(move |_| {
let (tx, rx) = mpsc::channel();
MyPark {
tx: Mutex::new(tx),
rx,
park_tx: park_tx.clone(),
unpark_tx: unpark_tx.clone(),
}
});
struct MyTask {
task_tx: Option<mpsc::Sender<Waker>>,
drop_tx: mpsc::Sender<()>,
}
impl Future for MyTask {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
if let Some(tx) = self.get_mut().task_tx.take() {
tx.send(cx.waker().clone()).unwrap();
}
Poll::Pending
}
}
impl Drop for MyTask {
fn drop(&mut self) {
self.drop_tx.send(()).unwrap();
}
}
pool.spawn(MyTask {
task_tx: Some(task_tx),
drop_tx,
});
// Wait until we get the task handle.
let task = task_rx.recv().unwrap();
// Drop the pool, this should result in futures being forcefully dropped.
drop(pool);
// Make sure `MyPark` and `MyUnpark` were dropped during shutdown.
assert_eq!(park_rx.try_recv(), Err(mpsc::TryRecvError::Disconnected));
assert_eq!(unpark_rx.try_recv(), Err(mpsc::TryRecvError::Disconnected));
// If the future is forcefully dropped, then we will get a signal here.
drop_rx.recv().unwrap();
// Ensure `task` lives until after the test completes.
drop(task);
}
#[test]
fn park_called_at_interval() {
struct MyPark {
park_light: Arc<AtomicBool>,
}
struct MyUnpark {}
impl Park for MyPark {
type Unpark = MyUnpark;
type Error = ();
fn unpark(&self) -> Self::Unpark {
MyUnpark {}
}
fn park(&mut self) -> Result<(), Self::Error> {
use std::thread;
use std::time::Duration;
thread::sleep(Duration::from_millis(1));
Ok(())
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
if duration == Duration::from_millis(0) {
self.park_light.store(true, Relaxed);
Ok(())
} else {
self.park()
}
}
}
impl Unpark for MyUnpark {
fn unpark(&self) {}
}
let park_light_1 = Arc::new(AtomicBool::new(false));
let park_light_2 = park_light_1.clone();
let (done_tx, done_rx) = mpsc::channel();
// Use 1 thread to ensure the worker stays busy.
let pool = Builder::new().num_threads(1).build_with_park(move |idx| {
assert_eq!(idx, 0);
MyPark {
park_light: park_light_2.clone(),
}
});
let mut cnt = 0;
pool.spawn(poll_fn(move |cx| {
let did_park_light = park_light_1.load(Relaxed);
if did_park_light {
// There is a bit of a race where the worker can tick a few times
// before seeing the task
assert!(cnt > 50);
done_tx.send(()).unwrap();
return Poll::Ready(());
}
cnt += 1;
cx.waker().wake_by_ref();
Poll::Pending
}));
done_rx.recv().unwrap();
}
fn new_pool() -> ThreadPool {
Builder::new().num_threads(4).build()
}
tokio/tests/timer_hammer.rs
+2 -4
View File
@@ -1,11 +1,9 @@
#![warn(rust_2018_idioms)]
use tokio::executor::current_thread::CurrentThread;
use tokio::executor::park::{Park, Unpark, UnparkThread};
use tokio::timer::{Delay, Timer};
use tokio_executor::current_thread::CurrentThread;
use tokio_executor::park::{Park, Unpark, UnparkThread};
use rand;
use rand::Rng;
use std::cmp;
use std::future::Future;