executor: rewrite the work-stealing thread pool (#1657)

This patch is a ground up rewrite of the existing work-stealing thread
pool. The goal is to reduce overhead while simplifying code when
possible.

At a high level, the following architectural changes were made:

- The local run queues were switched for bounded circle buffer queues.
- Reduce cross-thread synchronization.
- Refactor task constructs to use a single allocation and always include
  a join handle (#887).
- Simplify logic around putting workers to sleep and waking them up.

**Local run queues**

Move away from crossbeam's implementation of the Chase-Lev deque. This
implementation included unnecessary overhead as it supported
capabilities that are not needed for the work-stealing thread pool.
Instead, a fixed size circle buffer is used for the local queue. When
the local queue is full, half of the tasks contained in it are moved to
the global run queue.

**Reduce cross-thread synchronization**

This is done via many small improvements. Primarily, an upper bound is
placed on the number of concurrent stealers. Limiting the number of
stealers results in lower contention. Secondly, the rate at which
workers are notified and woken up is throttled. This also reduces
contention by preventing many threads from racing to steal work.

**Refactor task structure**

Now that Tokio is able to target a rust version that supports
`std::alloc` as well as `std::task`, the pool is able to optimize how
the task structure is laid out. Now, a single allocation per task is
required and a join handle is always provided enabling the spawner to
retrieve the result of the task (#887).

**Simplifying logic**

When possible, complexity is reduced in the implementation. This is done
by using locks and other simpler constructs in cold paths. The set of
sleeping workers is now represented as a `Mutex<VecDeque<usize>>`.
Instead of optimizing access to this structure, we reduce the amount the
pool must access this structure.

Secondly, we have (temporarily) removed `threadpool::blocking`. This
capability will come back later, but the original implementation was way
more complicated than necessary.

**Results**

The thread pool benchmarks have improved significantly:

Old thread pool:

```
test chained_spawn ... bench:   2,019,796 ns/iter (+/- 302,168)
test ping_pong     ... bench:   1,279,948 ns/iter (+/- 154,365)
test spawn_many    ... bench:  10,283,608 ns/iter (+/- 1,284,275)
test yield_many    ... bench:  21,450,748 ns/iter (+/- 1,201,337)
```

New thread pool:

```
test chained_spawn ... bench:     147,943 ns/iter (+/- 6,673)
test ping_pong     ... bench:     537,744 ns/iter (+/- 20,928)
test spawn_many    ... bench:   7,454,898 ns/iter (+/- 283,449)
test yield_many    ... bench:  16,771,113 ns/iter (+/- 733,424)
```

Real-world benchmarks improve significantly as well. This is testing the hyper hello
world server using: `wrk -t1 -c50 -d10`:

Old scheduler:

```
Running 10s test @ http://127.0.0.1:3000
  1 threads and 50 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   371.53us   99.05us   1.97ms   60.53%
    Req/Sec   114.61k     8.45k  133.85k    67.00%
  1139307 requests in 10.00s, 95.61MB read
Requests/sec: 113923.19
Transfer/sec:      9.56MB
```

New scheduler:

```
Running 10s test @ http://127.0.0.1:3000
  1 threads and 50 connections
  Thread Stats   Avg      Stdev     Max   +/- Stdev
    Latency   275.05us   69.81us   1.09ms   73.57%
    Req/Sec   153.17k    10.68k  171.51k    71.00%
  1522671 requests in 10.00s, 127.79MB read
Requests/sec: 152258.70
Transfer/sec:     12.78MB
```
This commit is contained in:
Carl Lerche
2019-10-19 11:09:40 -07:00
committed by GitHub
parent 2a181320b7
commit ed5a94eb2d
100 changed files with 7409 additions and 6796 deletions
+8 -1
View File
@@ -65,7 +65,7 @@ jobs:
tokio-codec: []
tokio-executor:
- current-thread
- threadpool
- thread-pool
tokio-io:
- util
tokio-sync:
@@ -94,6 +94,13 @@ jobs:
- tokio-no-features
- tokio-with-net
# Run loom tests
- template: ci/azure-loom.yml
parameters:
rust: beta
crates:
- tokio-executor
# Try cross compiling
- template: ci/azure-cross-compile.yml
parameters:
+18
View File
@@ -0,0 +1,18 @@
jobs:
- job: loom
displayName: Loom tests
pool:
vmImage: ubuntu-16.04
steps:
- template: azure-install-rust.yml
parameters:
rust_version: ${{ parameters.rust }}
- ${{ each crate in parameters.crates }}:
- script: RUSTFLAGS="--cfg loom" cargo test --lib --release
env:
LOOM_MAX_PREEMPTIONS: 2
CI: 'True'
displayName: test ${{ crate }}
workingDirectory: $(Build.SourcesDirectory)/${{ crate }}
+7 -16
View File
@@ -23,40 +23,31 @@ categories = ["concurrency", "asynchronous"]
[features]
blocking = ["tokio-sync", "lazy_static"]
current-thread = ["crossbeam-channel"]
threadpool = [
"tokio-sync",
"crossbeam-deque",
"crossbeam-queue",
"crossbeam-utils",
"futures-core-preview",
"num_cpus",
"lazy_static",
"slab",
]
thread-pool = ["num_cpus"]
[dependencies]
futures-util-preview = { version = "=0.3.0-alpha.19", features = ["channel"] }
tokio-sync = { version = "=0.2.0-alpha.6", optional = true, path = "../tokio-sync" }
tracing = { version = "0.1.5", optional = true }
futures-util-preview = { version = "=0.3.0-alpha.19", features = ["channel"] }
# current-thread dependencies
crossbeam-channel = { version = "0.3.8", optional = true }
# threadpool dependencies
crossbeam-deque = { version = "0.7.0", optional = true }
crossbeam-queue = { version = "0.1.0", optional = true }
crossbeam-utils = { version = "0.6.4", optional = true }
futures-core-preview = { version = "=0.3.0-alpha.19", optional = true }
num_cpus = { version = "1.2", optional = true }
# blocking
futures-core-preview = { version = "=0.3.0-alpha.19", optional = true }
lazy_static = { version = "1", optional = true }
slab = { version = "0.4.1", optional = true }
[dev-dependencies]
tokio = { version = "=0.2.0-alpha.6", path = "../tokio" }
tokio-sync = { version = "=0.2.0-alpha.6", path = "../tokio-sync" }
tokio-test = { version = "=0.2.0-alpha.6", path = "../tokio-test" }
futures-core-preview = "=0.3.0-alpha.19"
loom = { version = "0.2.9", features = ["futures", "checkpoint"] }
rand = "0.7"
[package.metadata.docs.rs]
-133
View File
@@ -1,133 +0,0 @@
#![cfg(feature = "broken")]
#![feature(test)]
#![warn(rust_2018_idioms)]
extern crate test;
const ITER: usize = 1_000;
mod blocking {
use super::*;
use futures::future::*;
use tokio_executor::threadpool::{blocking, Builder};
#[bench]
fn cpu_bound(b: &mut test::Bencher) {
let pool = Builder::new().pool_size(2).max_blocking(20).build();
b.iter(|| {
let count_down = Arc::new(CountDown::new(ITER));
for _ in 0..ITER {
let count_down = count_down.clone();
pool.spawn(lazy(move || {
poll_fn(|| blocking(|| perform_complex_computation()).map_err(|_| panic!()))
.and_then(move |_| {
// Do something with the value
count_down.dec();
Ok(())
})
}));
}
count_down.wait();
})
}
}
mod message_passing {
use super::*;
use futures::future::*;
use futures::sync::oneshot;
use tokio_executor::threadpool::Builder;
#[bench]
fn cpu_bound(b: &mut test::Bencher) {
let pool = Builder::new().pool_size(2).max_blocking(20).build();
let blocking = threadpool::ThreadPool::new(20);
b.iter(|| {
let count_down = Arc::new(CountDown::new(ITER));
for _ in 0..ITER {
let count_down = count_down.clone();
let blocking = blocking.clone();
pool.spawn(lazy(move || {
// Create a channel to receive the return value.
let (tx, rx) = oneshot::channel();
// Spawn a task on the blocking thread pool to process the
// computation.
blocking.execute(move || {
let res = perform_complex_computation();
tx.send(res).unwrap();
});
rx.and_then(move |_| {
count_down.dec();
Ok(())
})
.map_err(|_| panic!())
}));
}
count_down.wait();
})
}
}
fn perform_complex_computation() -> usize {
use rand::*;
// Simulate a CPU heavy computation
let mut rng = rand::thread_rng();
rng.gen()
}
// Util for waiting until the tasks complete
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::*;
use std::sync::*;
struct CountDown {
rem: AtomicUsize,
mutex: Mutex<()>,
condvar: Condvar,
}
impl CountDown {
fn new(rem: usize) -> Self {
CountDown {
rem: AtomicUsize::new(rem),
mutex: Mutex::new(()),
condvar: Condvar::new(),
}
}
fn dec(&self) {
let prev = self.rem.fetch_sub(1, AcqRel);
if prev != 1 {
return;
}
let _lock = self.mutex.lock().unwrap();
self.condvar.notify_all();
}
fn wait(&self) {
let mut lock = self.mutex.lock().unwrap();
loop {
if self.rem.load(Acquire) == 0 {
return;
}
lock = self.condvar.wait(lock).unwrap();
}
}
}
+161
View File
@@ -0,0 +1,161 @@
#![feature(test)]
extern crate test;
use tokio_executor::thread_pool::{Builder, Spawner, ThreadPool};
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();
});
}
-161
View File
@@ -1,161 +0,0 @@
#![cfg(feature = "broken")]
#![feature(test)]
#![warn(rust_2018_idioms)]
extern crate test;
const NUM_SPAWN: usize = 10_000;
const NUM_YIELD: usize = 1_000;
const TASKS_PER_CPU: usize = 50;
mod threadpool {
use futures::{future, task, Async};
use num_cpus;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::{mpsc, Arc};
use tokio_executor::threadpool::*;
#[bench]
fn spawn_many(b: &mut test::Bencher) {
let threadpool = ThreadPool::new();
let (tx, rx) = mpsc::sync_channel(10);
let rem = Arc::new(AtomicUsize::new(0));
b.iter(move || {
rem.store(super::NUM_SPAWN, SeqCst);
for _ in 0..super::NUM_SPAWN {
let tx = tx.clone();
let rem = rem.clone();
threadpool.spawn(future::lazy(move || {
if 1 == rem.fetch_sub(1, SeqCst) {
tx.send(()).unwrap();
}
Ok(())
}));
}
let _ = rx.recv().unwrap();
});
}
#[bench]
fn yield_many(b: &mut test::Bencher) {
let threadpool = ThreadPool::new();
let tasks = super::TASKS_PER_CPU * num_cpus::get();
let (tx, rx) = mpsc::sync_channel(tasks);
b.iter(move || {
for _ in 0..tasks {
let mut rem = super::NUM_YIELD;
let tx = tx.clone();
threadpool.spawn(future::poll_fn(move || {
rem -= 1;
if rem == 0 {
tx.send(()).unwrap();
Ok(Async::Ready(()))
} else {
// Notify the current task
task::current().notify();
// Not ready
Ok(Async::NotReady)
}
}));
}
for _ in 0..tasks {
let _ = rx.recv().unwrap();
}
});
}
}
// In this case, CPU pool completes the benchmark faster, but this is due to how
// CpuPool currently behaves, starving other futures. This completes the
// benchmark quickly but results in poor runtime characteristics for a thread
// pool.
//
// See rust-lang-nursery/futures-rs#617
//
mod cpupool {
use futures::future::{self, Executor};
use futures::{task, Async};
use futures_cpupool::*;
use num_cpus;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::SeqCst;
use std::sync::{mpsc, Arc};
#[bench]
fn spawn_many(b: &mut test::Bencher) {
let pool = CpuPool::new(num_cpus::get());
let (tx, rx) = mpsc::sync_channel(10);
let rem = Arc::new(AtomicUsize::new(0));
b.iter(move || {
rem.store(super::NUM_SPAWN, SeqCst);
for _ in 0..super::NUM_SPAWN {
let tx = tx.clone();
let rem = rem.clone();
pool.execute(future::lazy(move || {
if 1 == rem.fetch_sub(1, SeqCst) {
tx.send(()).unwrap();
}
Ok(())
}))
.ok()
.unwrap();
}
let _ = rx.recv().unwrap();
});
}
#[bench]
fn yield_many(b: &mut test::Bencher) {
let pool = CpuPool::new(num_cpus::get());
let tasks = super::TASKS_PER_CPU * num_cpus::get();
let (tx, rx) = mpsc::sync_channel(tasks);
b.iter(move || {
for _ in 0..tasks {
let mut rem = super::NUM_YIELD;
let tx = tx.clone();
pool.execute(future::poll_fn(move || {
rem -= 1;
if rem == 0 {
tx.send(()).unwrap();
Ok(Async::Ready(()))
} else {
// Notify the current task
task::current().notify();
// Not ready
Ok(Async::NotReady)
}
}))
.ok()
.unwrap();
}
for _ in 0..tasks {
let _ = rx.recv().unwrap();
}
});
}
}
@@ -1,72 +0,0 @@
#![cfg(feature = "broken")]
#![feature(test)]
#![warn(rust_2018_idioms)]
extern crate test;
const ITER: usize = 20_000;
mod us {
use futures::future;
use std::sync::mpsc;
use tokio_executor::threadpool::*;
#[bench]
fn chained_spawn(b: &mut test::Bencher) {
let threadpool = ThreadPool::new();
fn spawn(pool_tx: Sender, res_tx: mpsc::Sender<()>, n: usize) {
if n == 0 {
res_tx.send(()).unwrap();
} else {
let pool_tx2 = pool_tx.clone();
pool_tx
.spawn(future::lazy(move || {
spawn(pool_tx2, res_tx, n - 1);
Ok(())
}))
.unwrap();
}
}
b.iter(move || {
let (res_tx, res_rx) = mpsc::channel();
spawn(threadpool.sender().clone(), res_tx, super::ITER);
res_rx.recv().unwrap();
});
}
}
mod cpupool {
use futures::future::{self, Executor};
use futures_cpupool::*;
use num_cpus;
use std::sync::mpsc;
#[bench]
fn chained_spawn(b: &mut test::Bencher) {
let pool = CpuPool::new(num_cpus::get());
fn spawn(pool: CpuPool, res_tx: mpsc::Sender<()>, n: usize) {
if n == 0 {
res_tx.send(()).unwrap();
} else {
let pool2 = pool.clone();
pool.execute(future::lazy(move || {
spawn(pool2, res_tx, n - 1);
Ok(())
}))
.ok()
.unwrap();
}
}
b.iter(move || {
let (res_tx, res_rx) = mpsc::channel();
spawn(pool.clone(), res_tx, super::ITER);
res_rx.recv().unwrap();
});
}
}
+69 -30
View File
@@ -1,4 +1,7 @@
use super::{Executor, SpawnError};
#[cfg(feature = "thread-pool")]
use crate::thread_pool::ThreadPool;
use crate::{Executor, SpawnError};
use std::cell::Cell;
use std::future::Future;
use std::pin::Pin;
@@ -37,17 +40,18 @@ impl DefaultExecutor {
#[inline]
fn with_current<F: FnOnce(&mut dyn Executor) -> R, R>(f: F) -> Option<R> {
EXECUTOR.with(
|current_executor| match current_executor.replace(State::Active) {
State::Ready(executor_ptr) => {
let executor = unsafe { &mut *executor_ptr };
let result = f(executor);
current_executor.set(State::Ready(executor_ptr));
Some(result)
}
State::Empty | State::Active => None,
},
)
EXECUTOR.with(|current_executor| match current_executor.get() {
State::Ready(executor_ptr) => {
let executor = unsafe { &mut *executor_ptr };
Some(f(executor))
}
#[cfg(feature = "thread-pool")]
State::ThreadPool(threadpool_ptr) => {
let mut thread_pool = unsafe { &*threadpool_ptr };
Some(f(&mut thread_pool))
}
State::Empty => None,
})
}
}
@@ -55,10 +59,13 @@ impl DefaultExecutor {
enum State {
// default executor not defined
Empty,
// default executor is defined and ready to be used
// default executor is a thread pool instance.
#[cfg(feature = "thread-pool")]
ThreadPool(*const ThreadPool),
// default executor is set to a custom executor.
Ready(*mut dyn Executor),
// default executor is currently active (used to detect recursive calls)
Active,
}
thread_local! {
@@ -132,7 +139,26 @@ pub fn spawn<T>(future: T)
where
T: Future<Output = ()> + Send + 'static,
{
DefaultExecutor::current().spawn(Box::pin(future)).unwrap()
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 = "thread-pool")]
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 = "thread-pool")]
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
@@ -143,9 +169,24 @@ 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.get();
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.
@@ -159,16 +200,15 @@ where
let _reset = Reset(cell, was);
// 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 _) };
if let State::Ready(executor) = state {
let executor = unsafe { &mut *executor };
cell.set(State::Ready(executor));
if executor.status().is_err() {
panic!("executor not active; is this because `with_default` is called with `DefaultExecutor`?");
}
}
cell.set(state);
f()
})
@@ -183,7 +223,7 @@ unsafe fn hide_lt<'a>(p: *mut (dyn Executor + 'a)) -> *mut (dyn Executor + 'stat
#[cfg(test)]
mod tests {
use super::{with_default, DefaultExecutor, Executor};
use super::{with_default, DefaultExecutor};
#[test]
fn default_executor_is_send_and_sync() {
@@ -193,12 +233,11 @@ mod tests {
}
#[test]
#[should_panic]
fn nested_default_executor_status() {
let _enter = super::super::enter().unwrap();
let mut executor = DefaultExecutor::current();
let result = with_default(&mut executor, || DefaultExecutor::current().status());
assert!(result.err().unwrap().is_shutdown())
let _result = with_default(&mut executor, || ());
}
}
+18 -2
View File
@@ -60,6 +60,17 @@
//! [`DefaultExecutor`]: struct.DefaultExecutor.html
//! [`Park`]: park/index.html
//! [`Future::poll`]: https://doc.rust-lang.org/std/future/trait.Future.html#tymethod.poll
#[cfg(all(test, loom))]
macro_rules! thread_local {
($($tts:tt)+) => { loom::thread_local!{ $($tts)+ } }
}
// At the top due to macros
#[cfg(test)]
#[macro_use]
mod tests;
#[cfg(any(feature = "current-thread", feature = "threadpool"))]
#[macro_use]
mod tracing;
@@ -68,8 +79,13 @@ mod enter;
mod error;
mod executor;
mod global;
mod loom;
pub mod park;
#[cfg(feature = "thread-pool")]
mod task;
mod typed;
#[cfg(feature = "thread-pool")]
mod util;
#[cfg(feature = "blocking")]
pub mod blocking;
@@ -77,8 +93,8 @@ pub mod blocking;
#[cfg(feature = "current-thread")]
pub mod current_thread;
#[cfg(feature = "threadpool")]
pub mod threadpool;
#[cfg(feature = "thread-pool")]
pub mod thread_pool;
pub use crate::enter::{enter, exit, Enter, EnterError};
pub use crate::error::SpawnError;
+25
View File
@@ -0,0 +1,25 @@
//! 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(feature = "thread-pool")]
pub(crate) use self::std::{alloc, cell, rand, sys, thread};
+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)
}
}
@@ -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 = "thread-pool")]
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)
}
}
@@ -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) {}
}
+77
View File
@@ -0,0 +1,77 @@
#[cfg(feature = "thread-pool")]
mod atomic_u32;
mod atomic_usize;
#[cfg(feature = "thread-pool")]
mod causal_cell;
#[cfg(feature = "thread-pool")]
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 = "thread-pool")]
pub(crate) mod cell {
pub(crate) use super::causal_cell::{CausalCell, CausalCheck};
}
#[cfg(feature = "thread-pool")]
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 = "thread-pool")]
pub(crate) use crate::loom::std::atomic_u32::AtomicU32;
pub(crate) use crate::loom::std::atomic_usize::AtomicUsize;
#[cfg(feature = "thread-pool")]
pub(crate) use std::sync::atomic::{fence, spin_loop_hint, AtomicPtr};
}
}
#[cfg(feature = "thread-pool")]
pub(crate) mod sys {
pub(crate) fn num_cpus() -> usize {
usize::max(1, num_cpus::get_physical())
}
}
#[cfg(feature = "thread-pool")]
pub(crate) use std::thread;
+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()
}
}
@@ -1,146 +1,14 @@
//! 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
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::{Arc, Condvar, Mutex};
use crate::park::{Park, Unpark};
use std::marker::PhantomData;
use std::mem;
use std::rc::Rc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::sync::atomic::Ordering;
use std::task::{RawWaker, RawWakerVTable, Waker};
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()
}
}
/// Blocks the current thread using a condition variable.
///
/// Implements the [`Park`] functionality by using a condition variable. An
@@ -353,7 +221,7 @@ impl Unpark for UnparkThread {
}
}
static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop_waker);
impl UnparkThread {
pub(crate) fn into_waker(self) -> Waker {
@@ -377,6 +245,10 @@ unsafe fn clone(raw: *const ()) -> RawWaker {
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();
+153
View File
@@ -0,0 +1,153 @@
use crate::loom::alloc::Track;
use crate::loom::cell::CausalCell;
use crate::task::raw::{self, Vtable};
use crate::task::state::State;
use crate::task::waker::waker_ref;
use crate::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) }
}
}
+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 {}
+546
View File
@@ -0,0 +1,546 @@
use crate::loom::alloc::Track;
use crate::loom::cell::CausalCheck;
use crate::task::core::{Cell, Core, Header, Trailer};
use crate::task::state::Snapshot;
use crate::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))
}
}
+74
View File
@@ -0,0 +1,74 @@
use crate::loom::alloc::Track;
use crate::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();
}
}
}
+70
View File
@@ -0,0 +1,70 @@
use crate::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()
}
}
+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()
}
}
+190
View File
@@ -0,0 +1,190 @@
use crate::loom::alloc::Track;
use crate::task::core::Cell;
use crate::task::harness::Harness;
use crate::task::state::{Snapshot, State};
use crate::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)
}
+85
View File
@@ -0,0 +1,85 @@
use crate::loom::sync::atomic::AtomicPtr;
use crate::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)
}
}
+502
View File
@@ -0,0 +1,502 @@
use crate::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::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::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::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::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()
}
}
+277
View File
@@ -0,0 +1,277 @@
use crate::task;
use crate::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()))
}))
})
}
+5
View File
@@ -0,0 +1,5 @@
#[cfg(loom)]
mod loom;
#[cfg(not(loom))]
mod task;
+644
View File
@@ -0,0 +1,644 @@
use crate::task::{self, Header};
use crate::tests::backoff::*;
use crate::tests::mock_schedule::{mock, Mock};
use crate::tests::track_drop::track_drop;
use tokio::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());
}
+107
View File
@@ -0,0 +1,107 @@
use crate::task::harness::Harness;
use crate::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 ()) {}
+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
}
}
+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();
}
}
}
+51
View File
@@ -0,0 +1,51 @@
use crate::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();
}
}
+66
View File
@@ -0,0 +1,66 @@
#![allow(warnings)]
use crate::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);
}
}
+131
View File
@@ -0,0 +1,131 @@
#![allow(warnings)]
use crate::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>) {}
}
+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),
_ => {}
}
}};
}
+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)
}
}
+253
View File
@@ -0,0 +1,253 @@
use crate::loom::sync::Arc;
use crate::loom::sys::num_cpus;
use crate::loom::thread;
use crate::park::Park;
use crate::thread_pool::park::DefaultPark;
use crate::thread_pool::{shutdown, worker, Spawner, ThreadPool};
use std::{fmt, usize};
/// Builds a thread pool with custom configuration values.
pub struct Builder {
/// Number of threads to spawn
pool_size: usize,
/// Thread name prefix
name_prefix: String,
/// Thread stack size
stack_size: Option<usize>,
/// Around worker callback
around_worker: Option<Arc<Callback>>,
}
type Callback = 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_prefix: "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 prefix of threads spawned by the scheduler
///
/// Thread name prefix is used for generating thread names. For example, if
/// prefix is `my-pool-`, then threads in the pool will get names like
/// `my-pool-1` etc.
///
/// 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_prefix("my-pool-")
/// .build();
/// ```
pub fn name_prefix<S: Into<String>>(&mut self, val: S) -> &mut Self {
self.name_prefix = 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 (pool, workers) = worker::create_set(self.pool_size, |i| BoxedPark::new(build_park(i)));
// Spawn threads for each worker
for (idx, mut worker) in workers.into_iter().enumerate() {
let around_worker = self.around_worker.clone();
let shutdown_tx = shutdown_tx.clone();
let mut th = thread::Builder::new().name(format!("{}{}", self.name_prefix, idx));
if let Some(stack) = self.stack_size {
th = th.stack_size(stack);
}
let res = th.spawn(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 {
cb(idx, &mut || worker.run());
} else {
worker.run();
}
// Worker must be dropped before the `shutdown_tx`
drop(worker);
// Dropping the handle must happen __after__ the callback
drop(shutdown_tx);
});
if let Err(err) = res {
panic!("failed to spawn worker thread: {:?}", err);
}
}
let spawner = Spawner::new(pool);
ThreadPool::from_parts(spawner, shutdown_rx)
}
}
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_prefix", &self.name_prefix)
.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::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)
}
}
+85
View File
@@ -0,0 +1,85 @@
use crate::loom::sync::Arc;
use crate::park::Unpark;
use crate::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,
}
}
}
+229
View File
@@ -0,0 +1,229 @@
//! Coordinates idling workers
use crate::loom::sync::atomic::AtomicUsize;
use crate::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());
}
+42
View File
@@ -0,0 +1,42 @@
use crate::park::Unpark;
use crate::task;
use crate::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()
}
}
+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::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;
+77
View File
@@ -0,0 +1,77 @@
use crate::task::{self, Task};
use crate::thread_pool::{queue, Shared};
use crate::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);
}
}
+182
View File
@@ -0,0 +1,182 @@
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::{Arc, Condvar, Mutex};
use crate::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();
}
}
+101
View File
@@ -0,0 +1,101 @@
use crate::thread_pool::{shutdown, Builder, JoinHandle, Spawner};
use crate::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,
}
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) -> ThreadPool {
ThreadPool {
spawner,
shutdown_rx,
}
}
/// 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::global::with_threadpool(self, || {
let mut enter = crate::enter().expect("attempting to block while on a Tokio executor");
enter.block_on(future)
})
}
/// Shutdown the thread pool.
pub fn shutdown_now(&mut self) {
if self.spawner.workers().close() {
self.shutdown_rx.wait();
}
}
}
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::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();
}
}
@@ -0,0 +1,195 @@
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::Mutex;
use crate::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;
}
@@ -0,0 +1,36 @@
use crate::loom::sync::Arc;
use crate::task::Task;
use crate::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();
}
}
@@ -0,0 +1,298 @@
use crate::loom::cell::{CausalCell, CausalCheck};
use crate::loom::sync::atomic::{self, AtomicU32};
use crate::task::Task;
use crate::thread_pool::queue::global;
use crate::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()
}
}
@@ -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::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() {}
}
}
@@ -0,0 +1,129 @@
use crate::task::Task;
use crate::thread_pool::queue::{local, Cluster, Inject};
// Loom primitive
use crate::loom::sync::Arc;
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()
}
}
+196
View File
@@ -0,0 +1,196 @@
//! Putting a worker to sleep.
//!
//! - Attempt to spin.
use crate::loom::rand::seed;
use crate::park::Unpark;
use crate::task::{self, Task};
use crate::thread_pool::{current, queue, BoxFuture, Idle, JoinHandle, Owned, Shared};
use crate::util::{CachePadded, FastRand};
use crate::{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,
}
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) -> 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),
}
}
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(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(())
}
}
+104
View File
@@ -0,0 +1,104 @@
use crate::park::Unpark;
use crate::task::{self, Schedule, Task};
use crate::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);
}
}
@@ -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::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::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);
}
}
+57
View File
@@ -0,0 +1,57 @@
use crate::loom::sync::Arc;
use crate::park::Unpark;
use crate::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);
}
/// 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()
}
}
@@ -0,0 +1,134 @@
use crate::loom::sync::atomic::Ordering::{Acquire, Relaxed, Release};
use crate::loom::sync::atomic::{AtomicBool, AtomicUsize};
use crate::loom::sync::{Arc, Mutex};
use crate::spawn;
use crate::tests::loom_oneshot as oneshot;
use crate::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;
});
drop(pool);
});
}
fn gated() -> impl Future<Output = &'static str> {
gated2(false)
}
fn gated2(thread: bool) -> impl Future<Output = &'static str> {
use crate::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
}
})
}
@@ -0,0 +1,68 @@
use crate::task::{self, Task};
use crate::tests::mock_schedule::{Noop, NOOP_SCHEDULE};
use crate::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 })
}
@@ -0,0 +1,11 @@
#[cfg(loom)]
mod loom_pool;
#[cfg(loom)]
mod loom_queue;
#[cfg(not(loom))]
mod queue;
#[cfg(not(loom))]
mod worker;
@@ -0,0 +1,281 @@
use crate::task::{self, Task};
use crate::tests::mock_schedule::{Noop, NOOP_SCHEDULE};
use crate::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())
}
@@ -0,0 +1,66 @@
use crate::tests::track_drop::track_drop;
use crate::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::tests::mock_park::MockPark::new();
let (pool, workers) = thread_pool::create_pool($n, |index| mock_park.mk_park(index));
(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 tokio::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();
}
+394
View File
@@ -0,0 +1,394 @@
use crate::loom::sync::Arc;
use crate::park::{Park, Unpark};
use crate::task::Task;
use crate::thread_pool::{current, Owned, Shared};
use std::time::Duration;
// TODO: remove this re-export
pub(super) use crate::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,
}
struct Entry<P: 'static> {
pool: Arc<Set<P>>,
index: usize,
}
pub(crate) fn create_set<F, P>(
pool_size: usize,
mk_park: F,
) -> (Arc<Set<P::Unpark>>, Vec<Worker<P>>)
where
P: 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()));
// 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)| 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: Park + 'static,
{
pub(super) fn new(pool: Arc<Set<P::Unpark>>, index: usize, park: P) -> Self {
Worker {
entry: Entry { pool, index },
park,
}
}
pub(super) fn run(&mut self) {
let mut executor = &*self.entry.pool;
let entry = &self.entry;
let park = &mut self.park;
// Track the current worker
current::set(&entry.pool, entry.index, || {
let _enter = crate::enter().expect("executor already running on thread");
crate::with_default(&mut executor, || {
entry.run(park);
})
})
}
#[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);
}
}
impl<P> Entry<P>
where
P: Unpark,
{
fn run(&self, park: &mut impl Park<Unpark = P>) {
while self.is_running() {
if self.tick(park) {
self.park(park);
}
}
self.shutdown(park);
}
fn is_running(&self) -> bool {
self.owned().is_running.get()
}
/// Returns `true` if the worker needs to park
fn tick(&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(&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(&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(&self) {
// Free any completed tasks
self.drain_tasks_pending_drop();
// Update the pool state cache
self.owned()
.is_running
.set(!self.owned().work_queue.is_closed());
}
fn search_for_work(&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(&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(&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(&self) -> bool {
let ret = self
.set()
.idle()
.transition_worker_to_parked(self.index, self.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(&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(&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(&self) {
if !self.owned().work_queue.is_empty() {
self.set().notify_work();
}
}
fn park(&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(&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(&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(&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(&self) -> u16 {
let tick = self.owned().tick.get();
self.owned().tick.set(tick.wrapping_add(1));
tick
}
fn is_searching(&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(&self) -> &Owned<P> {
// safety: we own the slot
unsafe { &*self.set().owned()[self.index].get() }
}
}
-174
View File
@@ -1,174 +0,0 @@
use super::worker::Worker;
use futures_core::ready;
use std::error::Error;
use std::fmt;
use std::task::Poll;
/// Error raised by `blocking`.
pub struct BlockingError {
_p: (),
}
/// Enter a blocking section of code.
///
/// The `blocking` function annotates a section of code that performs a blocking
/// operation, either by issuing a blocking syscall or by performing a long
/// running CPU-bound computation.
///
/// When the `blocking` function enters, it hands off the responsibility of
/// processing the current work queue to another thread. Then, it calls the
/// supplied closure. The closure is permitted to block indefinitely.
///
/// If the maximum number of concurrent `blocking` calls has been reached, then
/// `NotReady` is returned and the task is notified once existing `blocking`
/// calls complete. The maximum value is specified when creating a thread pool
/// using [`Builder::max_blocking`][build]
///
/// NB: The entire task that called `blocking` is blocked whenever the supplied
/// closure blocks, even if you have used future combinators such as `select` -
/// the other futures in this task will not make progress until the closure
/// returns.
/// If this is not desired, ensure that `blocking` runs in its own task (e.g.
/// using `futures::sync::oneshot::spawn`).
///
/// [build]: struct.Builder.html#method.max_blocking
///
/// # Return
///
/// When the blocking closure is executed, `Ok(Ready(T))` is returned, where
/// `T` is the closure's return value.
///
/// If the thread pool has shutdown, `Err` is returned.
///
/// If the number of concurrent `blocking` calls has reached the maximum,
/// `Ok(NotReady)` is returned and the current task is notified when a call to
/// `blocking` will succeed.
///
/// If `blocking` is called from outside the context of a Tokio thread pool,
/// `Err` is returned.
///
/// # Background
///
/// By default, the Tokio thread pool expects that tasks will only run for short
/// periods at a time before yielding back to the thread pool. This is the basic
/// premise of cooperative multitasking.
///
/// However, it is common to want to perform a blocking operation while
/// processing an asynchronous computation. Examples of blocking operation
/// include:
///
/// * Performing synchronous file operations (reading and writing).
/// * Blocking on acquiring a mutex.
/// * Performing a CPU bound computation, like cryptographic encryption or
/// decryption.
///
/// One option for dealing with blocking operations in an asynchronous context
/// is to use a thread pool dedicated to performing these operations. This not
/// ideal as it requires bidirectional message passing as well as a channel to
/// communicate which adds a level of buffering.
///
/// Instead, `blocking` hands off the responsibility of processing the work queue
/// to another thread. This hand off is light compared to a channel and does not
/// require buffering.
///
/// # Examples
///
/// Block on receiving a message from a `std` channel. This example is a little
/// silly as using the non-blocking channel from the `futures` crate would make
/// more sense. The blocking receive can be replaced with any blocking operation
/// that needs to be performed.
///
/// ```rust
/// use tokio_executor::threadpool::{ThreadPool, blocking};
///
/// use futures_util::future::poll_fn;
/// use std::sync::mpsc;
/// use std::thread;
/// use std::time::Duration;
///
/// pub fn main() {
/// // This is a *blocking* channel
/// let (tx, rx) = mpsc::channel();
///
/// // Spawn a thread to send a message
/// thread::spawn(move || {
/// thread::sleep(Duration::from_millis(500));
/// tx.send("hello").unwrap();
/// });
///
/// let pool = ThreadPool::new();
///
/// pool.spawn(async move {
/// // Because `blocking` returns `Poll`, it is intended to be used
/// // from the context of a `Future` implementation. Since we don't
/// // have a complicated requirement, we can use `poll_fn` in this
/// // case.
/// let _ = poll_fn(move |_| {
/// blocking(|| {
/// let msg = rx.recv().unwrap();
/// println!("message = {}", msg);
/// }).map_err(|_| panic!("the threadpool shut down"))
/// }).await;
/// });
///
/// // Wait for the task we just spawned to complete.
/// pool.shutdown_on_idle().wait();
/// }
/// ```
pub fn blocking<F, T>(f: F) -> Poll<Result<T, BlockingError>>
where
F: FnOnce() -> T,
{
let res = Worker::with_current(|worker| {
let worker = match worker {
Some(worker) => worker,
None => {
return Poll::Ready(Err(BlockingError { _p: () }));
}
};
// Transition the worker state to blocking. This will exit the fn early
// with `NotReady` if the pool does not have enough capacity to enter
// blocking mode.
worker.transition_to_blocking()
});
// If the transition cannot happen, exit early
ready!(res)?;
// Currently in blocking mode, so call the inner closure
//
// "Exit" the current executor in case the blocking function wants
// to call a different executor.
let ret = crate::exit(move || f());
// Try to transition out of blocking mode. This is a fast path that takes
// back ownership of the worker if the worker handoff didn't complete yet.
Worker::with_current(|worker| {
// Worker must be set since it was above.
worker.unwrap().transition_from_blocking();
});
// Return the result
Poll::Ready(Ok(ret))
}
impl fmt::Display for BlockingError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
fmt,
"`blocking` annotation used from outside the context of a thread pool"
)
}
}
impl fmt::Debug for BlockingError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BlockingError")
.field("reason", &format!("{}", self))
.finish()
}
}
impl Error for BlockingError {}
-426
View File
@@ -1,426 +0,0 @@
use super::callback::Callback;
use super::config::{Config, MAX_WORKERS};
use super::park::{BoxPark, BoxedPark, DefaultPark};
use super::pool::{Pool, MAX_BACKUP};
use super::shutdown::ShutdownTrigger;
use super::thread_pool::ThreadPool;
use super::worker::{self, Worker, WorkerId};
use crate::park::Park;
use crossbeam_deque::Injector;
use num_cpus;
use std::any::Any;
use std::cmp::max;
use std::error::Error;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;
/// Builds a thread pool with custom configuration values.
///
/// Methods can be chained in order to set the configuration values. The thread
/// pool is constructed by calling [`build`].
///
/// New instances of `Builder` are obtained via [`Builder::new`].
///
/// See function level documentation for details on the various configuration
/// settings.
///
/// [`build`]: #method.build
/// [`Builder::new`]: #method.new
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// use std::time::Duration;
///
/// let thread_pool = Builder::new()
/// .pool_size(4)
/// .keep_alive(Some(Duration::from_secs(30)))
/// .build();
///
/// thread_pool.spawn(async {
/// println!("called from a worker thread");
/// });
///
/// // Gracefully shutdown the threadpool
/// thread_pool.shutdown().wait();
/// ```
pub struct Builder {
/// Thread pool specific configuration values
config: Config,
/// Number of workers to spawn
pool_size: usize,
/// Maximum number of futures that can be in a blocking section
/// concurrently.
max_blocking: usize,
/// Generates the `Park` instances
new_park: Box<dyn Fn(&WorkerId) -> BoxPark>,
}
impl Builder {
/// Returns a new thread pool builder initialized with default configuration
/// values.
///
/// Configuration methods can be chained on the return value.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
/// use std::time::Duration;
///
/// let thread_pool = Builder::new()
/// .pool_size(4)
/// .keep_alive(Some(Duration::from_secs(30)))
/// .build();
/// ```
pub fn new() -> Builder {
let num_cpus = max(1, num_cpus::get());
let new_park =
Box::new(|_: &WorkerId| Box::new(BoxedPark::new(DefaultPark::new())) as BoxPark);
Builder {
pool_size: num_cpus,
max_blocking: 100,
config: Config {
keep_alive: None,
name_prefix: None,
stack_size: None,
around_worker: None,
after_start: None,
before_stop: None,
panic_handler: None,
},
new_park,
}
}
/// Set the maximum number of worker threads for the thread pool instance.
///
/// This must be a number between 1 and 32,768 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::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .pool_size(4)
/// .build();
/// ```
pub fn pool_size(&mut self, val: usize) -> &mut Self {
assert!(val >= 1, "at least one thread required");
assert!(val <= MAX_WORKERS, "max value is {}", MAX_WORKERS);
self.pool_size = val;
self
}
/// Set the maximum number of concurrent blocking sections.
///
/// When the maximum concurrent `blocking` calls is reached, any further
/// calls to `blocking` will return `NotReady` and the task is notified once
/// previously in-flight calls to `blocking` return.
///
/// This must be a number between 1 and 32,768 though it is advised to keep
/// this value on the smaller side.
///
/// The default value is 100.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .max_blocking(200)
/// .build();
/// ```
pub fn max_blocking(&mut self, val: usize) -> &mut Self {
assert!(val <= MAX_BACKUP, "max value is {}", MAX_BACKUP);
self.max_blocking = val;
self
}
/// Set the thread keep alive duration
///
/// If set, a thread that has completed a `blocking` call will wait for up
/// to the specified duration to become a worker thread again. Once the
/// duration elapses, the thread will shutdown.
///
/// When the value is `None`, the thread will wait to become a worker
/// thread forever.
///
/// The default value is `None`.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
/// use std::time::Duration;
///
/// let thread_pool = Builder::new()
/// .keep_alive(Some(Duration::from_secs(30)))
/// .build();
/// ```
pub fn keep_alive(&mut self, val: Option<Duration>) -> &mut Self {
self.config.keep_alive = val;
self
}
/// Sets a callback to be triggered when a panic during a future bubbles up
/// to Tokio. By default Tokio catches these panics, and they will be
/// ignored. The parameter passed to this callback is the same error value
/// returned from std::panic::catch_unwind(). To abort the process on
/// panics, use std::panic::resume_unwind() in this callback as shown
/// below.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .panic_handler(|err| std::panic::resume_unwind(err))
/// .build();
/// ```
pub fn panic_handler<F>(&mut self, f: F) -> &mut Self
where
F: Fn(Box<dyn Any + Send>) + Send + Sync + 'static,
{
self.config.panic_handler = Some(Arc::new(f));
self
}
/// Set name prefix of threads spawned by the scheduler
///
/// Thread name prefix is used for generating thread names. For example, if
/// prefix is `my-pool-`, then threads in the pool will get names like
/// `my-pool-1` etc.
///
/// If this configuration is not set, then the thread will use the system
/// default naming scheme.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .name_prefix("my-pool-")
/// .build();
/// ```
pub fn name_prefix<S: Into<String>>(&mut self, val: S) -> &mut Self {
self.config.name_prefix = Some(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::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .stack_size(32 * 1024)
/// .build();
/// ```
pub fn stack_size(&mut self, val: usize) -> &mut Self {
self.config.stack_size = Some(val);
self
}
/// Execute function `f` on each worker thread.
///
/// This function is provided a handle to the worker and is expected to call
/// [`Worker::run`], otherwise the worker thread will shutdown without doing
/// any work.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .around_worker(|worker| {
/// println!("worker is starting up");
/// worker.run();
/// println!("worker is shutting down");
/// })
/// .build();
/// ```
///
/// [`Worker::run`]: struct.Worker.html#method.run
pub fn around_worker<F>(&mut self, f: F) -> &mut Self
where
F: Fn(&Worker) + Send + Sync + 'static,
{
self.config.around_worker = Some(Callback::new(f));
self
}
/// Execute function `f` after each thread is started but before it starts
/// doing work.
///
/// This is intended for bookkeeping and monitoring use cases.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .after_start(|| {
/// println!("thread started");
/// })
/// .build();
/// ```
pub fn after_start<F>(&mut self, f: F) -> &mut Self
where
F: Fn() + Send + Sync + 'static,
{
self.config.after_start = Some(Arc::new(f));
self
}
/// Execute function `f` before each thread stops.
///
/// This is intended for bookkeeping and monitoring use cases.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .before_stop(|| {
/// println!("thread stopping");
/// })
/// .build();
/// ```
pub fn before_stop<F>(&mut self, f: F) -> &mut Self
where
F: Fn() + Send + Sync + 'static,
{
self.config.before_stop = Some(Arc::new(f));
self
}
/// Customize the `park` instance used by each worker thread.
///
/// The provided closure `f` is called once per worker and returns a `Park`
/// instance that is used by the worker to put itself to sleep.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
/// use tokio_executor::threadpool::park::DefaultPark;
/// # fn decorate<F>(f: F) -> F { f }
///
/// let thread_pool = Builder::new()
/// .custom_park(|_| {
/// // This is the default park type that the worker would use if we
/// // did not customize it.
/// let park = DefaultPark::new();
///
/// // Decorate the `park` instance, allowing us to customize work
/// // that happens when a worker thread goes to sleep.
/// decorate(park)
/// })
/// .build();
/// ```
pub fn custom_park<F, P>(&mut self, f: F) -> &mut Self
where
F: Fn(&WorkerId) -> P + 'static,
P: Park + Send + 'static,
P::Error: Error,
{
self.new_park = Box::new(move |id| Box::new(BoxedPark::new(f(id))));
self
}
/// Create the configured `ThreadPool`.
///
/// The returned `ThreadPool` instance is ready to spawn tasks.
///
/// # Examples
///
/// ```
/// use tokio_executor::threadpool::Builder;
///
/// let thread_pool = Builder::new()
/// .build();
/// ```
pub fn build(&self) -> ThreadPool {
trace!(message = "build;", num_workers = self.pool_size);
// Create the worker entry list
let workers: Arc<[worker::Entry]> = {
let mut workers = vec![];
for i in 0..self.pool_size {
let id = WorkerId::new(i);
let park = (self.new_park)(&id);
let unpark = park.unpark();
workers.push(worker::Entry::new(park, unpark));
}
workers.into()
};
let queue = Arc::new(Injector::new());
// Create a trigger that will clean up resources on shutdown.
//
// The `Pool` contains a weak reference to it, while `Worker`s and the `ThreadPool` contain
// strong references.
let trigger = Arc::new(ShutdownTrigger::new(workers.clone(), queue.clone()));
// Create the pool
let pool = Arc::new(Pool::new(
workers,
Arc::downgrade(&trigger),
self.max_blocking,
self.config.clone(),
queue,
));
ThreadPool::new2(pool, trigger)
}
}
impl fmt::Debug for Builder {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Builder")
.field("config", &self.config)
.field("pool_size", &self.pool_size)
.field("new_park", &"Box<Fn() -> BoxPark>")
.finish()
}
}
impl Default for Builder {
fn default() -> Self {
Self::new()
}
}
-28
View File
@@ -1,28 +0,0 @@
use super::worker::Worker;
use std::fmt;
use std::sync::Arc;
#[derive(Clone)]
pub(crate) struct Callback {
f: Arc<dyn Fn(&Worker) + Send + Sync>,
}
impl Callback {
pub(crate) fn new<F>(f: F) -> Self
where
F: Fn(&Worker) + Send + Sync + 'static,
{
Callback { f: Arc::new(f) }
}
pub(crate) fn call(&self, worker: &Worker) {
(self.f)(worker)
}
}
impl fmt::Debug for Callback {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "Fn")
}
}
-37
View File
@@ -1,37 +0,0 @@
use super::callback::Callback;
use std::any::Any;
use std::fmt;
use std::sync::Arc;
use std::time::Duration;
/// Thread pool specific configuration values
#[derive(Clone)]
pub(crate) struct Config {
pub(crate) keep_alive: Option<Duration>,
// Used to configure a worker thread
pub(crate) name_prefix: Option<String>,
pub(crate) stack_size: Option<usize>,
pub(crate) around_worker: Option<Callback>,
pub(crate) after_start: Option<Arc<dyn Fn() + Send + Sync>>,
pub(crate) before_stop: Option<Arc<dyn Fn() + Send + Sync>>,
pub(crate) panic_handler: Option<PanicHandler>,
}
// Define type alias to avoid clippy::type_complexity.
type PanicHandler = Arc<dyn Fn(Box<dyn Any + Send>) + Send + Sync>;
/// Max number of workers that can be part of a pool. This is the most that can
/// fit in the scheduler state. Note, that this is the max number of **active**
/// threads. There can be more standby threads.
pub(crate) const MAX_WORKERS: usize = 1 << 15;
impl fmt::Debug for Config {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Config")
.field("keep_alive", &self.keep_alive)
.field("name_prefix", &self.name_prefix)
.field("stack_size", &self.stack_size)
.finish()
}
}
-146
View File
@@ -1,146 +0,0 @@
//! A work-stealing based thread pool for executing futures.
//!
//! The Tokio thread pool supports scheduling futures and processing them on
//! multiple CPU cores. It is optimized for the primary Tokio use case of many
//! independent tasks with limited computation and with most tasks waiting on
//! I/O. Usually, users will not create a `ThreadPool` instance directly, but
//! will use one via a [`runtime`].
//!
//! The `ThreadPool` structure manages two sets of threads:
//!
//! * Worker threads.
//! * Backup threads.
//!
//! Worker threads are used to schedule futures using a work-stealing strategy.
//! Backup threads, on the other hand, are intended only to support the
//! `blocking` API. Threads will transition between the two sets.
//!
//! The advantage of the work-stealing strategy is minimal cross-thread
//! coordination. The thread pool attempts to make as much progress as possible
//! without communicating across threads.
//!
//! ## Worker overview
//!
//! Each worker has two queues: a deque and a mpsc channel. The deque is the
//! primary queue for tasks that are scheduled to run on the worker thread. Tasks
//! can only be pushed onto the deque by the worker, but other workers may
//! "steal" from that deque. The mpsc channel is used to submit futures while
//! external to the pool.
//!
//! As long as the thread pool has not been shutdown, a worker will run in a
//! loop. Each loop, it consumes all tasks on its mpsc channel and pushes it onto
//! the deque. It then pops tasks off of the deque and executes them.
//!
//! If a worker has no work, i.e., both queues are empty. It attempts to steal.
//! To do this, it randomly scans other workers' deques and tries to pop a task.
//! If it finds no work to steal, the thread goes to sleep.
//!
//! When the worker detects that the pool has been shut down, it exits the loop,
//! cleans up its state, and shuts the thread down.
//!
//! ## Thread pool initialization
//!
//! Note, users normally will use the threadpool created by a [`runtime`].
//!
//! By default, no threads are spawned on creation. Instead, when new futures are
//! spawned, the pool first checks if there are enough active worker threads. If
//! not, a new worker thread is spawned.
//!
//! ## Spawning futures
//!
//! The spawning behavior depends on whether a future was spawned from within a
//! worker or thread or if it was spawned from an external handle.
//!
//! When spawning a future while external to the thread pool, the current
//! strategy is to randomly pick a worker to submit the task to. The task is then
//! pushed onto that worker's mpsc channel.
//!
//! When spawning a future while on a worker thread, the task is pushed onto the
//! back of the current worker's deque.
//!
//! ## Blocking annotation strategy
//!
//! The [`blocking`] function is used to annotate a section of code that
//! performs a blocking operation, either by issuing a blocking syscall or
//! performing any long running CPU-bound computation.
//!
//! The strategy for handling blocking closures is to hand off the worker to a
//! new thread. This implies handing off the `deque` and `mpsc`. Once this is
//! done, the new thread continues to process the work queue and the original
//! thread is able to block. Once it finishes processing the blocking future, the
//! thread has no additional work and is inserted into the backup pool. This
//! makes it available to other workers that encounter a [`blocking`] call.
//!
//! [`blocking`]: fn.blocking.html
//! [`runtime`]: https://docs.rs/tokio/0.1/tokio/runtime/
// ## Crate layout
//
// The primary type, `Pool`, holds the majority of a thread pool's state,
// including the state for each worker. Each worker's state is maintained in an
// instance of `worker::Entry`.
//
// `Worker` contains the logic that runs on each worker thread. It holds an
// `Arc` to `Pool` and is able to access its state from `Pool`.
//
// `Task` is a harness around an individual future. It manages polling and
// scheduling that future.
//
// ## Sleeping workers
//
// Sleeping workers are tracked using a [Treiber stack]. This results in the
// thread that most recently went to sleep getting woken up first. When the pool
// is not under load, this helps threads shutdown faster.
//
// Sleeping is done by using `tokio_executor::Park` implementations. This allows
// the user of the thread pool to customize the work that is performed to sleep.
// This is how injecting timers and other functionality into the thread pool is
// done.
//
// ## Notifying workers
//
// When there is work to be done, workers must be notified. However, notifying a
// worker requires cross thread coordination. Ideally, a worker would only be
// notified when it is sleeping, but there is no way to know if a worker is
// sleeping without cross thread communication.
//
// The two cases when a worker might need to be notified are:
//
// 1. A task is externally submitted to a worker via the mpsc channel.
// 2. A worker has a back log of work and needs other workers to steal from it.
//
// In the first case, the worker will always be notified. However, it could be
// possible to avoid the notification if the mpsc channel has two or greater
// number of tasks *after* the task is submitted. In this case, we are able to
// assume that the worker has previously been notified.
//
// The second case is trickier. Currently, whenever a worker spawns a new future
// (pushing it onto its deque) and when it pops a future from its mpsc, it tries
// to notify a sleeping worker to wake up and start stealing. This is a lot of
// notification and it **might** be possible to reduce it.
//
// Also, whenever a worker is woken up via a signal and it does find work, it,
// in turn, will try to wake up a new worker.
//
// [Treiber stack]: https://en.wikipedia.org/wiki/Treiber_Stack
pub mod park;
mod blocking;
mod builder;
mod callback;
mod config;
mod pool;
mod sender;
mod shutdown;
mod task;
mod thread_pool;
mod waker;
mod worker;
pub use self::blocking::{blocking, BlockingError};
pub use self::builder::Builder;
pub use self::sender::Sender;
pub use self::shutdown::Shutdown;
pub use self::thread_pool::ThreadPool;
pub use self::worker::{Worker, WorkerId};
@@ -1,46 +0,0 @@
use crate::park::{Park, Unpark};
use std::error::Error;
use std::time::Duration;
pub(crate) type BoxPark = Box<dyn Park<Unpark = BoxUnpark, Error = ()> + Send>;
pub(crate) type BoxUnpark = Box<dyn Unpark>;
pub(crate) struct BoxedPark<T>(T);
impl<T> BoxedPark<T> {
pub(crate) fn new(inner: T) -> Self {
BoxedPark(inner)
}
}
impl<T: Park + Send> Park for BoxedPark<T>
where
T::Error: Error,
{
type Unpark = BoxUnpark;
type Error = ();
fn unpark(&self) -> Self::Unpark {
Box::new(self.0.unpark())
}
fn park(&mut self) -> Result<(), Self::Error> {
self.0.park().map_err(|_e| {
// if tracing is disabled, the compiler will flag this as unused.
warn!(
message = "calling `park` on worker thread errored -- shutting down thread",
error = %_e
);
})
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
self.0.park_timeout(duration).map_err(|_e| {
warn!(
message = "calling `park` on worker thread errored -- shutting down thread",
error = %_e,
);
})
}
}
@@ -1,98 +0,0 @@
use crate::park::{Park, Unpark};
use crossbeam_utils::sync::{Parker, Unparker};
use std::error::Error;
use std::fmt;
use std::time::Duration;
/// Parks the thread.
#[derive(Debug)]
pub struct DefaultPark {
inner: Parker,
}
/// Unparks threads that were parked by `DefaultPark`.
#[derive(Debug)]
pub struct DefaultUnpark {
inner: Unparker,
}
/// 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: (),
}
// ===== impl DefaultPark =====
impl DefaultPark {
/// Creates a new `DefaultPark` instance.
pub fn new() -> DefaultPark {
DefaultPark {
inner: Parker::new(),
}
}
/// Unpark the thread without having to clone the unpark handle.
///
/// Named `notify` to avoid conflicting with the `unpark` fn.
pub(crate) fn notify(&self) {
self.inner.unparker().unpark();
}
pub(crate) fn park_sync(&self, duration: Option<Duration>) {
match duration {
None => self.inner.park(),
Some(duration) => self.inner.park_timeout(duration),
}
}
}
impl Park for DefaultPark {
type Unpark = DefaultUnpark;
type Error = ParkError;
fn unpark(&self) -> Self::Unpark {
DefaultUnpark {
inner: self.inner.unparker().clone(),
}
}
fn park(&mut self) -> Result<(), Self::Error> {
self.inner.park();
Ok(())
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
self.inner.park_timeout(duration);
Ok(())
}
}
impl Default for DefaultPark {
fn default() -> Self {
Self::new()
}
}
// ===== impl DefaultUnpark =====
impl Unpark for DefaultUnpark {
fn unpark(&self) {
self.inner.unpark();
}
}
// ===== impl ParkError =====
impl fmt::Display for ParkError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "unknown park error")
}
}
impl Error for ParkError {}
@@ -1,8 +0,0 @@
//! Thread parking utilities.
mod boxed;
mod default_park;
pub use self::default_park::{DefaultPark, DefaultUnpark, ParkError};
pub(crate) use self::boxed::{BoxPark, BoxUnpark, BoxedPark};
@@ -1,308 +0,0 @@
use super::super::park::DefaultPark;
use super::super::worker::WorkerId;
use std::cell::UnsafeCell;
use std::fmt;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{self, AcqRel, Acquire, Relaxed};
use std::time::{Duration, Instant};
/// State associated with a thread in the thread pool.
///
/// The pool manages a number of threads. Some of those threads are considered
/// "primary" threads and process the work queue. When a task being run on a
/// primary thread enters a blocking context, the responsibility of processing
/// the work queue must be handed off to another thread. This is done by first
/// checking for idle threads on the backup stack. If one is found, the worker
/// token (`WorkerId`) is handed off to that running thread. If none are found,
/// a new thread is spawned.
///
/// This state manages the exchange. A thread that is idle, not assigned to a
/// work queue, sits around for a specified amount of time. When the worker
/// token is handed off, it is first stored in `handoff`. The backup thread is
/// then signaled. At this point, the backup thread wakes up from sleep and
/// reads `handoff`. At that point, it has been promoted to a primary thread and
/// will begin processing inbound work on the work queue.
///
/// The name `Backup` isn't really great for what the type does, but I have not
/// come up with a better name... Maybe it should just be named `Thread`.
#[derive(Debug)]
pub(crate) struct Backup {
/// Worker ID that is being handed to this thread.
handoff: UnsafeCell<Option<WorkerId>>,
/// Thread state.
///
/// This tracks:
///
/// * Is queued flag
/// * If the pool is shutting down.
/// * If the thread is running
state: AtomicUsize,
/// Next entry in the Treiber stack.
next_sleeper: UnsafeCell<BackupId>,
/// Used to put the thread to sleep
park: DefaultPark,
}
#[derive(Debug, Eq, PartialEq, Copy, Clone)]
pub(crate) struct BackupId(pub(crate) usize);
#[derive(Debug)]
pub(crate) enum Handoff {
Worker(WorkerId),
Idle,
Terminated,
}
/// Tracks thread state.
#[derive(Clone, Copy, Eq, PartialEq)]
struct State(usize);
/// Set when the worker is pushed onto the scheduler's stack of sleeping
/// threads.
///
/// This flag also serves as a "notification" bit. If another thread is
/// attempting to hand off a worker to the backup thread, then the pushed bit
/// will not be set when the thread tries to shutdown.
pub(crate) const PUSHED: usize = 0b001;
/// Set when the thread is running
pub(crate) const RUNNING: usize = 0b010;
/// Set when the thread pool has terminated
pub(crate) const TERMINATED: usize = 0b100;
// ===== impl Backup =====
impl Backup {
pub(crate) fn new() -> Backup {
Backup {
handoff: UnsafeCell::new(None),
state: AtomicUsize::new(State::new().into()),
next_sleeper: UnsafeCell::new(BackupId(0)),
park: DefaultPark::new(),
}
}
/// Called when the thread is starting
pub(crate) fn start(&self, worker_id: &WorkerId) {
debug_assert!({
let state: State = self.state.load(Relaxed).into();
debug_assert!(!state.is_pushed());
debug_assert!(state.is_running());
debug_assert!(!state.is_terminated());
true
});
// The handoff value is equal to `worker_id`
debug_assert_eq!(unsafe { (*self.handoff.get()).as_ref() }, Some(worker_id));
unsafe {
*self.handoff.get() = None;
}
}
pub(crate) fn is_running(&self) -> bool {
let state: State = self.state.load(Relaxed).into();
state.is_running()
}
/// Hands off the worker to a thread.
///
/// Returns `true` if the thread needs to be spawned.
pub(crate) fn worker_handoff(&self, worker_id: WorkerId) -> bool {
unsafe {
// The backup worker should not already have been handoff a worker.
debug_assert!((*self.handoff.get()).is_none());
// Set the handoff
*self.handoff.get() = Some(worker_id);
}
// This *probably* can just be `Release`... memory orderings, how do
// they work?
let prev = State::worker_handoff(&self.state);
debug_assert!(prev.is_pushed());
if prev.is_running() {
// Wakeup the backup thread
self.park.notify();
false
} else {
true
}
}
/// Terminate the worker
pub(crate) fn signal_stop(&self) {
let prev: State = self.state.fetch_xor(TERMINATED | PUSHED, AcqRel).into();
debug_assert!(!prev.is_terminated());
debug_assert!(prev.is_pushed());
if prev.is_running() {
self.park.notify();
}
}
/// Release the worker
pub(crate) fn release(&self) {
let prev: State = self.state.fetch_xor(RUNNING, AcqRel).into();
debug_assert!(prev.is_running());
}
/// Wait for a worker handoff
pub(crate) fn wait_for_handoff(&self, timeout: Option<Duration>) -> Handoff {
let sleep_until = timeout.map(|dur| Instant::now() + dur);
let mut state: State = self.state.load(Acquire).into();
// Run in a loop since there can be spurious wakeups
loop {
if !state.is_pushed() {
if state.is_terminated() {
return Handoff::Terminated;
}
let worker_id = unsafe { (*self.handoff.get()).take().expect("no worker handoff") };
return Handoff::Worker(worker_id);
}
match sleep_until {
None => {
self.park.park_sync(None);
state = self.state.load(Acquire).into();
}
Some(when) => {
let now = Instant::now();
if now < when {
self.park.park_sync(Some(when - now));
state = self.state.load(Acquire).into();
} else {
debug_assert!(state.is_running());
// Transition out of running
let mut next = state;
next.unset_running();
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
debug_assert!(!next.is_running());
return Handoff::Idle;
}
state = actual;
}
}
}
}
}
pub(crate) fn is_pushed(&self) -> bool {
let state: State = self.state.load(Relaxed).into();
state.is_pushed()
}
pub(crate) fn set_pushed(&self, ordering: Ordering) {
let prev: State = self.state.fetch_or(PUSHED, ordering).into();
debug_assert!(!prev.is_pushed());
}
#[inline]
pub(crate) fn next_sleeper(&self) -> BackupId {
unsafe { *self.next_sleeper.get() }
}
#[inline]
pub(crate) fn set_next_sleeper(&self, val: BackupId) {
unsafe {
*self.next_sleeper.get() = val;
}
}
}
// ===== impl State =====
impl State {
/// Returns a new, default, thread `State`
pub(crate) fn new() -> State {
State(0)
}
/// Returns true if the thread entry is pushed in the sleeper stack
pub(crate) fn is_pushed(self) -> bool {
self.0 & PUSHED == PUSHED
}
fn unset_pushed(&mut self) {
self.0 &= !PUSHED;
}
pub(crate) fn is_running(self) -> bool {
self.0 & RUNNING == RUNNING
}
pub(crate) fn set_running(&mut self) {
self.0 |= RUNNING;
}
pub(crate) fn unset_running(&mut self) {
self.0 &= !RUNNING;
}
pub(crate) fn is_terminated(self) -> bool {
self.0 & TERMINATED == TERMINATED
}
fn worker_handoff(state: &AtomicUsize) -> State {
let mut curr: State = state.load(Acquire).into();
loop {
let mut next = curr;
next.set_running();
next.unset_pushed();
let actual = state
.compare_and_swap(curr.into(), next.into(), AcqRel)
.into();
if actual == curr {
return curr;
}
curr = actual;
}
}
}
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("backup::State")
.field("is_pushed", &self.is_pushed())
.field("is_running", &self.is_running())
.field("is_terminated", &self.is_terminated())
.finish()
}
}
@@ -1,188 +0,0 @@
use super::{Backup, BackupId};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{AcqRel, Acquire};
#[derive(Debug)]
pub(crate) struct BackupStack {
state: AtomicUsize,
}
#[derive(Debug, Eq, PartialEq, Clone, Copy)]
struct State(usize);
pub(crate) const MAX_BACKUP: usize = 1 << 15;
/// Extracts the head of the backup stack from the state
const STACK_MASK: usize = ((1 << 16) - 1);
/// Used to mark the stack as empty
pub(crate) const EMPTY: BackupId = BackupId(MAX_BACKUP);
/// Used to mark the stack as terminated
pub(crate) const TERMINATED: BackupId = BackupId(EMPTY.0 + 1);
/// How many bits the Treiber ABA guard is offset by
const ABA_GUARD_SHIFT: usize = 16;
#[cfg(target_pointer_width = "64")]
const ABA_GUARD_MASK: usize = (1 << (64 - ABA_GUARD_SHIFT)) - 1;
#[cfg(target_pointer_width = "32")]
const ABA_GUARD_MASK: usize = (1 << (32 - ABA_GUARD_SHIFT)) - 1;
// ===== impl BackupStack =====
impl BackupStack {
pub(crate) fn new() -> BackupStack {
let state = AtomicUsize::new(State::new().into());
BackupStack { state }
}
/// Push a backup thread onto the stack
///
/// # Return
///
/// Returns `Ok` on success.
///
/// Returns `Err` if the pool has transitioned to the `TERMINATED` state.
/// When terminated, pushing new entries is no longer permitted.
pub(crate) fn push(&self, entries: &[Backup], id: BackupId) -> Result<(), ()> {
let mut state: State = self.state.load(Acquire).into();
entries[id.0].set_pushed(AcqRel);
loop {
let mut next = state;
let head = state.head();
if head == TERMINATED {
// The pool is terminated, cannot push the sleeper.
return Err(());
}
entries[id.0].set_next_sleeper(head);
next.set_head(id);
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if state == actual {
return Ok(());
}
state = actual;
}
}
/// Pop a backup thread off the stack.
///
/// If `terminate` is set and the stack is empty when this function is
/// called, the state of the stack is transitioned to "terminated". At this
/// point, no further entries can be pushed onto the stack.
///
/// # Return
///
/// * Returns the index of the popped worker and the worker's observed
/// state.
///
/// * `Ok(None)` if the stack is empty.
/// * `Err(_)` is returned if the pool has been shutdown.
pub(crate) fn pop(&self, entries: &[Backup], terminate: bool) -> Result<Option<BackupId>, ()> {
// Figure out the empty value
let terminal = if terminate { TERMINATED } else { EMPTY };
let mut state: State = self.state.load(Acquire).into();
loop {
let head = state.head();
if head == EMPTY {
let mut next = state;
next.set_head(terminal);
if next == state {
debug_assert!(terminal == EMPTY);
return Ok(None);
}
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual != state {
state = actual;
continue;
}
return Ok(None);
} else if head == TERMINATED {
return Err(());
}
debug_assert!(head.0 < MAX_BACKUP);
let mut next = state;
let next_head = entries[head.0].next_sleeper();
// TERMINATED can never be set as the "next pointer" on a worker.
debug_assert!(next_head != TERMINATED);
if next_head == EMPTY {
next.set_head(terminal);
} else {
next.set_head(next_head);
}
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
debug_assert!(entries[head.0].is_pushed());
return Ok(Some(head));
}
state = actual;
}
}
}
// ===== impl State =====
impl State {
fn new() -> State {
State(EMPTY.0)
}
fn head(self) -> BackupId {
BackupId(self.0 & STACK_MASK)
}
fn set_head(&mut self, val: BackupId) {
let val = val.0;
// The ABA guard protects against the ABA problem w/ Treiber stacks
let aba_guard = ((self.0 >> ABA_GUARD_SHIFT) + 1) & ABA_GUARD_MASK;
self.0 = (aba_guard << ABA_GUARD_SHIFT) | val;
}
}
impl From<usize> for State {
fn from(src: usize) -> Self {
State(src)
}
}
impl From<State> for usize {
fn from(src: State) -> Self {
src.0
}
}
-481
View File
@@ -1,481 +0,0 @@
mod backup;
mod backup_stack;
mod state;
pub(crate) use self::backup::{Backup, BackupId};
pub(crate) use self::backup_stack::MAX_BACKUP;
pub(crate) use self::state::{Lifecycle, State, MAX_FUTURES};
use self::backup::Handoff;
use self::backup_stack::BackupStack;
use super::config::Config;
use super::shutdown::ShutdownTrigger;
use super::task::{Blocking, Task};
use super::worker::{self, Worker, WorkerId};
use super::BlockingError;
use crossbeam_deque::Injector;
use crossbeam_utils::CachePadded;
use lazy_static::lazy_static;
use std::cell::Cell;
use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hash, Hasher};
use std::num::Wrapping;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{AcqRel, Acquire};
use std::sync::{Arc, Weak};
use std::task::Poll;
use std::thread;
#[derive(Debug)]
pub(crate) struct Pool {
// Tracks the state of the thread pool (running, shutting down, ...).
//
// While workers check this field as a hint to detect shutdown, it is
// **not** used as a primary point of coordination for workers. The sleep
// stack is used as the primary point of coordination for workers.
//
// The value of this atomic is deserialized into a `pool::State` instance.
// See comments for that type.
pub(crate) state: CachePadded<AtomicUsize>,
// Stack tracking sleeping workers.
sleep_stack: CachePadded<worker::Stack>,
// Worker state
//
// A worker is a thread that is processing the work queue and polling
// futures.
//
// The number of workers will *usually* be small.
pub(crate) workers: Arc<[worker::Entry]>,
// The global MPMC queue of tasks.
//
// Spawned tasks are pushed into this queue. Although worker threads have their own dedicated
// task queues, they periodically steal tasks from this global queue, too.
pub(crate) queue: Arc<Injector<Arc<Task>>>,
// Completes the shutdown process when the `ThreadPool` and all `Worker`s get dropped.
//
// When spawning a new `Worker`, this weak reference is upgraded and handed out to the new
// thread.
pub(crate) trigger: Weak<ShutdownTrigger>,
// Backup thread state
//
// In order to efficiently support `blocking`, a pool of backup threads is
// needed. These backup threads are ready to take over a worker if the
// future being processed requires blocking.
backup: Box<[Backup]>,
// Stack of sleeping backup threads
pub(crate) backup_stack: BackupStack,
// State regarding coordinating blocking sections and tracking tasks that
// are pending blocking capacity.
blocking: Blocking,
// Configuration
pub(crate) config: Config,
}
impl Pool {
/// Create a new `Pool`
pub(crate) fn new(
workers: Arc<[worker::Entry]>,
trigger: Weak<ShutdownTrigger>,
max_blocking: usize,
config: Config,
queue: Arc<Injector<Arc<Task>>>,
) -> Pool {
let pool_size = workers.len();
let total_size = max_blocking + pool_size;
// Create the set of backup entries
//
// This is `backup + pool_size` because the core thread pool running the
// workers is spawned from backup as well.
let backup = (0..total_size)
.map(|_| Backup::new())
.collect::<Vec<_>>()
.into_boxed_slice();
let backup_stack = BackupStack::new();
for i in (0..backup.len()).rev() {
backup_stack.push(&backup, BackupId(i)).unwrap();
}
// Initialize the blocking state
let blocking = Blocking::new(max_blocking);
let ret = Pool {
state: CachePadded::new(AtomicUsize::new(State::new().into())),
sleep_stack: CachePadded::new(worker::Stack::new()),
workers,
queue,
trigger,
backup,
backup_stack,
blocking,
config,
};
// Now, we prime the sleeper stack
for i in 0..pool_size {
ret.sleep_stack.push(&ret.workers, i).unwrap();
}
ret
}
/// Start shutting down the pool. This means that no new futures will be
/// accepted.
#[cfg_attr(feature = "tracing", tracing::instrument(level = "trace"))]
pub(crate) fn shutdown(&self, now: bool, purge_queue: bool) {
let mut state: State = self.state.load(Acquire).into();
trace!(?state);
// For now, this must be true
debug_assert!(!purge_queue || now);
// Start by setting the shutdown flag
loop {
let mut next = state;
let num_futures = next.num_futures();
if next.lifecycle() == Lifecycle::ShutdownNow {
// Already transitioned to shutting down state
if !purge_queue || num_futures == 0 {
// Nothing more to do
return;
}
// The queue must be purged
debug_assert!(purge_queue);
next.clear_num_futures();
} else {
next.set_lifecycle(if now || num_futures == 0 {
// If already idle, always transition to shutdown now.
Lifecycle::ShutdownNow
} else {
Lifecycle::ShutdownOnIdle
});
if purge_queue {
next.clear_num_futures();
}
}
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if state == actual {
state = next;
break;
}
state = actual;
}
trace!("transitioned to shutdown");
// Only transition to terminate if there are no futures currently on the
// pool
if state.num_futures() != 0 {
return;
}
self.terminate_sleeping_workers();
}
/// Called by `Worker` as it tries to enter a sleeping state. Before it
/// sleeps, it must push itself onto the sleep stack. This enables other
/// threads to see it when signaling work.
pub(crate) fn push_sleeper(&self, idx: usize) -> Result<(), ()> {
self.sleep_stack.push(&self.workers, idx)
}
pub(crate) fn terminate_sleeping_workers(&self) {
use super::worker::Lifecycle::Signaled;
trace!("shutting down workers");
// Wakeup all sleeping workers. They will wake up, see the state
// transition, and terminate.
while let Some((idx, worker_state)) = self.sleep_stack.pop(&self.workers, Signaled, true) {
self.workers[idx].signal_stop(worker_state);
}
// Now terminate any backup threads
//
// The call to `pop` must be successful because shutting down the pool
// is coordinated and at this point, this is the only thread that will
// attempt to transition the backup stack to "terminated".
while let Ok(Some(backup_id)) = self.backup_stack.pop(&self.backup, true) {
self.backup[backup_id.0].signal_stop();
}
}
pub(crate) fn poll_blocking_capacity(
&self,
task: &Arc<Task>,
) -> Poll<Result<(), BlockingError>> {
self.blocking.poll_blocking_capacity(task)
}
/// Submit a task to the scheduler.
///
/// Called from either inside or outside of the scheduler. If currently on
/// the scheduler, then a fast path is taken.
pub(crate) fn submit(&self, task: Arc<Task>, pool: &Arc<Pool>) {
debug_assert_eq!(*self, **pool);
Worker::with_current(|worker| {
if let Some(worker) = worker {
// If the worker is in blocking mode, then even though the
// thread-local variable is set, the current thread does not
// have ownership of that worker entry. This is because the
// worker entry has already been handed off to another thread.
//
// The second check handles the case where the current thread is
// part of a different threadpool than the one being submitted
// to.
if !worker.is_blocking() && *self == *worker.pool {
let idx = worker.id.0;
trace!(message = "submit internal;", idx);
worker.pool.workers[idx].submit_internal(task);
worker.pool.signal_work(pool);
return;
}
}
self.submit_external(task, pool);
});
}
/// Submit a task to the scheduler from off worker
///
/// Called from outside of the scheduler, this function is how new tasks
/// enter the system.
pub(crate) fn submit_external(&self, task: Arc<Task>, pool: &Arc<Pool>) {
debug_assert_eq!(*self, **pool);
trace!("submit external");
self.queue.push(task);
self.signal_work(pool);
}
pub(crate) fn release_backup(&self, backup_id: BackupId) -> Result<(), ()> {
// First update the state, this cannot fail because the caller must have
// exclusive access to the backup token.
self.backup[backup_id.0].release();
// Push the backup entry back on the stack
self.backup_stack.push(&self.backup, backup_id)
}
pub(crate) fn notify_blocking_task(&self, pool: &Arc<Pool>) {
debug_assert_eq!(*self, **pool);
self.blocking.notify_task(&pool);
}
/// Provision a thread to run a worker
pub(crate) fn spawn_thread(&self, id: WorkerId, pool: &Arc<Pool>) {
debug_assert_eq!(*self, **pool);
let backup_id = match self.backup_stack.pop(&self.backup, false) {
Ok(Some(backup_id)) => backup_id,
Ok(None) => panic!("no thread available"),
Err(_) => {
debug!("failed to spawn worker thread due to the thread pool shutting down");
return;
}
};
let need_spawn = self.backup[backup_id.0].worker_handoff(id.clone());
if !need_spawn {
return;
}
let trigger = match self.trigger.upgrade() {
None => {
// The pool is shutting down.
return;
}
Some(t) => t,
};
let mut th = thread::Builder::new();
if let Some(ref prefix) = pool.config.name_prefix {
th = th.name(format!("{}{}", prefix, backup_id.0));
}
if let Some(stack) = pool.config.stack_size {
th = th.stack_size(stack);
}
let pool = pool.clone();
let res = th.spawn(move || {
if let Some(ref f) = pool.config.after_start {
f();
}
let mut worker_id = id;
pool.backup[backup_id.0].start(&worker_id);
loop {
// The backup token should be in the running state.
debug_assert!(pool.backup[backup_id.0].is_running());
// TODO: Avoid always cloning
let worker = Worker::new(worker_id, backup_id, pool.clone(), trigger.clone());
// Run the worker. If the worker transitioned to a "blocking"
// state, then `is_blocking` will be true.
if !worker.do_run() {
// The worker shutdown, so exit the thread.
break;
}
debug_assert!(!pool.backup[backup_id.0].is_pushed());
// Push the thread back onto the backup stack. This makes it
// available for future handoffs.
//
// This **must** happen before notifying the task.
let res = pool.backup_stack.push(&pool.backup, backup_id);
if res.is_err() {
// The pool is being shutdown.
break;
}
// The task switched the current thread to blocking mode.
// Now that the blocking task completed, any tasks
pool.notify_blocking_task(&pool);
debug_assert!(pool.backup[backup_id.0].is_running());
// Wait for a handoff
let handoff = pool.backup[backup_id.0].wait_for_handoff(pool.config.keep_alive);
match handoff {
Handoff::Worker(id) => {
debug_assert!(pool.backup[backup_id.0].is_running());
worker_id = id;
}
Handoff::Idle | Handoff::Terminated => {
break;
}
}
}
if let Some(ref f) = pool.config.before_stop {
f();
}
});
if let Err(err) = res {
error!(message = "failed to spawn worker thread;", ?err);
panic!("failed to spawn worker thread: {:?}", err);
}
}
/// If there are any other workers currently relaxing, signal them that work
/// is available so that they can try to find more work to process.
#[allow(clippy::cognitive_complexity)] // https://github.com/rust-lang/rust-clippy/issues/3900
pub(crate) fn signal_work(&self, pool: &Arc<Pool>) {
debug_assert_eq!(*self, **pool);
use super::worker::Lifecycle::Signaled;
if let Some((idx, worker_state)) = self.sleep_stack.pop(&self.workers, Signaled, false) {
let span = trace_span!("signal_work", idx);
let _enter = span.enter();
let entry = &self.workers[idx];
debug_assert!(
worker_state.lifecycle() != Signaled,
"actual={:?}",
worker_state.lifecycle(),
);
trace!("notify");
if !entry.notify(worker_state) {
trace!("spawn;");
self.spawn_thread(WorkerId(idx), pool);
}
}
}
/// Generates a random number
///
/// Uses a thread-local random number generator based on XorShift.
pub(crate) fn rand_usize(&self) -> usize {
thread_local! {
static RNG: Cell<Wrapping<u32>> = Cell::new(Wrapping(prng_seed()));
}
RNG.with(|rng| {
// This is the 32-bit variant of Xorshift.
// https://en.wikipedia.org/wiki/Xorshift
let mut x = rng.get();
x ^= x << 13;
x ^= x >> 17;
x ^= x << 5;
rng.set(x);
x.0 as usize
})
}
}
impl PartialEq for Pool {
fn eq(&self, other: &Pool) -> bool {
self as *const _ == other as *const _
}
}
unsafe impl Send for Pool {}
unsafe impl Sync for Pool {}
// Return a thread-specific, 32-bit, non-zero seed value suitable for a 32-bit
// PRNG. This uses one libstd RandomState for a default hasher and hashes on
// the current thread ID to obtain an unpredictable, collision resistant seed.
fn prng_seed() -> u32 {
// This obtains a small number of random bytes from the host system (for
// example, on unix via getrandom(2)) in order to seed an unpredictable and
// HashDoS resistant 64-bit hash function (currently: `SipHasher13` with
// 128-bit state). We only need one of these, to make the seeds for all
// process threads different via hashed IDs, collision resistant, and
// unpredictable.
lazy_static! {
static ref RND_STATE: RandomState = RandomState::new();
}
// Hash the current thread ID to produce a u32 value
let mut hasher = RND_STATE.build_hasher();
thread::current().id().hash(&mut hasher);
let hash: u64 = hasher.finish();
let seed = (hash as u32) ^ ((hash >> 32) as u32);
// Ensure non-zero seed (Xorshift yields only zero's for that seed)
if seed == 0 {
0x9b4e_6d25 // misc bits, could be any non-zero
} else {
seed
}
}
-132
View File
@@ -1,132 +0,0 @@
use std::{fmt, usize};
/// ThreadPool state.
///
/// The two least significant bits are the shutdown flags. (0 for active, 1 for
/// shutdown on idle, 2 for shutting down). The remaining bits represent the
/// number of futures that still need to complete.
#[derive(Eq, PartialEq, Clone, Copy)]
pub(crate) struct State(usize);
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Clone, Copy)]
#[repr(usize)]
pub(crate) enum Lifecycle {
/// The thread pool is currently running
Running = 0,
/// The thread pool should shutdown once it reaches an idle state.
ShutdownOnIdle = 1,
/// The thread pool should start the process of shutting down.
ShutdownNow = 2,
}
/// Mask used to extract the number of futures from the state
const LIFECYCLE_MASK: usize = 0b11;
const NUM_FUTURES_MASK: usize = !LIFECYCLE_MASK;
const NUM_FUTURES_OFFSET: usize = 2;
/// Max number of futures the pool can handle.
pub(crate) const MAX_FUTURES: usize = usize::MAX >> NUM_FUTURES_OFFSET;
// ===== impl State =====
impl State {
#[inline]
pub(crate) fn new() -> State {
State(0)
}
/// Returns the number of futures still pending completion.
pub(crate) fn num_futures(self) -> usize {
self.0 >> NUM_FUTURES_OFFSET
}
/// Increment the number of futures pending completion.
///
/// Returns false on failure.
pub(crate) fn inc_num_futures(&mut self) {
debug_assert!(self.num_futures() < MAX_FUTURES);
debug_assert!(self.lifecycle() < Lifecycle::ShutdownNow);
self.0 += 1 << NUM_FUTURES_OFFSET;
}
/// Decrement the number of futures pending completion.
pub(crate) fn dec_num_futures(&mut self) {
let num_futures = self.num_futures();
if num_futures == 0 {
// Already zero
return;
}
self.0 -= 1 << NUM_FUTURES_OFFSET;
if self.lifecycle() == Lifecycle::ShutdownOnIdle && num_futures == 1 {
self.set_lifecycle(Lifecycle::ShutdownNow);
}
}
/// Set the number of futures pending completion to zero
pub(crate) fn clear_num_futures(&mut self) {
self.0 &= LIFECYCLE_MASK;
}
pub(crate) fn lifecycle(self) -> Lifecycle {
(self.0 & LIFECYCLE_MASK).into()
}
pub(crate) fn set_lifecycle(&mut self, val: Lifecycle) {
self.0 = (self.0 & NUM_FUTURES_MASK) | (val as usize);
}
pub(crate) fn is_terminated(self) -> bool {
self.lifecycle() == Lifecycle::ShutdownNow && self.num_futures() == 0
}
}
impl From<usize> for State {
fn from(src: usize) -> Self {
State(src)
}
}
impl From<State> for usize {
fn from(src: State) -> Self {
src.0
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("pool::State")
.field("lifecycle", &self.lifecycle())
.field("num_futures", &self.num_futures())
.finish()
}
}
// ===== impl Lifecycle =====
impl From<usize> for Lifecycle {
fn from(src: usize) -> Lifecycle {
use self::Lifecycle::*;
debug_assert!(
src == Running as usize
|| src == ShutdownOnIdle as usize
|| src == ShutdownNow as usize
);
unsafe { ::std::mem::transmute(src) }
}
}
impl From<Lifecycle> for usize {
fn from(src: Lifecycle) -> usize {
let v = src as usize;
debug_assert!(v & LIFECYCLE_MASK == v);
v
}
}
-193
View File
@@ -1,193 +0,0 @@
use super::pool::{self, Lifecycle, Pool, MAX_FUTURES};
use super::task::Task;
use crate::{Executor, SpawnError, TypedExecutor};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering::{AcqRel, Acquire};
use std::sync::Arc;
/// Submit futures to the associated thread pool for execution.
///
/// A `Sender` instance is a handle to a single thread pool, allowing the owner
/// of the handle to spawn futures onto the thread pool. New futures are spawned
/// using [`Sender::spawn`].
///
/// The `Sender` handle is *only* used for spawning new futures. It does not
/// impact the lifecycle of the thread pool in any way.
///
/// `Sender` instances are obtained by calling [`ThreadPool::sender`]. The
/// `Sender` struct implements the `Executor` trait.
///
/// [`Sender::spawn`]: #method.spawn
/// [`ThreadPool::sender`]: struct.ThreadPool.html#method.sender
#[derive(Debug)]
pub struct Sender {
pub(crate) pool: Arc<Pool>,
}
impl Sender {
/// Spawn a future onto the thread pool
///
/// This function takes ownership of the future and spawns it onto the
/// thread pool, assigning it to a worker thread. The exact strategy used to
/// assign a future to a worker depends on if the caller is already on a
/// worker thread or external to the thread pool.
///
/// If the caller is currently on the thread pool, the spawned future will
/// be assigned to the same worker that the caller is on. If the caller is
/// external to the thread pool, the future will be assigned to a random
/// worker.
///
/// If `spawn` returns `Ok`, this does not mean that the future will be
/// executed. The thread pool can be forcibly shutdown between the time
/// `spawn` is called and the future has a chance to execute.
///
/// If `spawn` returns `Err`, then the future failed to be spawned. There
/// are two possible causes:
///
/// * The thread pool is at capacity and is unable to spawn a new future.
/// This is a temporary failure. At some point in the future, the thread
/// pool might be able to spawn new futures.
/// * The thread pool is shutdown. This is a permanent failure indicating
/// that the handle will never be able to spawn new futures.
///
/// The status of the thread pool can be queried before calling `spawn`
/// using the `status` function (part of the `Executor` trait).
///
/// # Examples
///
/// ```rust
/// use tokio_executor::threadpool::ThreadPool;
///
/// // Create a thread pool with default configuration values
/// let thread_pool = ThreadPool::new();
///
/// thread_pool.sender().spawn(async {
/// println!("called from a worker thread");
/// }).unwrap();
///
/// // Gracefully shutdown the threadpool
/// thread_pool.shutdown().wait();
/// ```
pub fn spawn<F>(&self, future: F) -> Result<(), SpawnError>
where
F: Future<Output = ()> + Send + 'static,
{
let mut s = self;
Executor::spawn(&mut s, Box::pin(future))
}
/// Logic to prepare for spawning
fn prepare_for_spawn(&self) -> Result<(), SpawnError> {
let mut state: pool::State = self.pool.state.load(Acquire).into();
// Increment the number of futures spawned on the pool as well as
// validate that the pool is still running/
loop {
let mut next = state;
if next.num_futures() == MAX_FUTURES {
// No capacity
return Err(SpawnError::at_capacity());
}
if next.lifecycle() == Lifecycle::ShutdownNow {
// Cannot execute the future, executor is shutdown.
return Err(SpawnError::shutdown());
}
next.inc_num_futures();
let actual = self
.pool
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
trace!(message = "execute;", count = next.num_futures());
break;
}
state = actual;
}
Ok(())
}
}
impl Executor for Sender {
fn status(&self) -> Result<(), SpawnError> {
let s = self;
Executor::status(&s)
}
fn spawn(
&mut self,
future: Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), SpawnError> {
let mut s = &*self;
Executor::spawn(&mut s, future)
}
}
impl Executor for &Sender {
fn status(&self) -> Result<(), SpawnError> {
let state: pool::State = self.pool.state.load(Acquire).into();
if state.num_futures() == MAX_FUTURES {
// No capacity
return Err(SpawnError::at_capacity());
}
if state.lifecycle() == Lifecycle::ShutdownNow {
// Cannot execute the future, executor is shutdown.
return Err(SpawnError::shutdown());
}
Ok(())
}
fn spawn(
&mut self,
future: Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), SpawnError> {
self.prepare_for_spawn()?;
// At this point, the pool has accepted the future, so schedule it for
// execution.
// Create a new task for the future
let task = Arc::new(Task::new(future));
// Call `submit_external()` in order to place the task into the global
// queue. This way all workers have equal chance of running this task,
// which means IO handles will be assigned to reactors more evenly.
self.pool.submit_external(task, &self.pool);
Ok(())
}
}
impl<T> TypedExecutor<T> for Sender
where
T: Future<Output = ()> + Send + 'static,
{
fn status(&self) -> Result<(), SpawnError> {
Executor::status(self)
}
fn spawn(&mut self, future: T) -> Result<(), SpawnError> {
Executor::spawn(self, Box::pin(future))
}
}
impl Clone for Sender {
#[inline]
fn clone(&self) -> Sender {
let pool = self.pool.clone();
Sender { pool }
}
}
-110
View File
@@ -1,110 +0,0 @@
use super::task::Task;
use super::worker;
use tokio_sync::AtomicWaker;
use crossbeam_deque::Injector;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
/// Future that resolves when the thread pool is shutdown.
///
/// A `ThreadPool` is shutdown once all the worker have drained their queues and
/// shutdown their threads.
///
/// `Shutdown` is returned by [`shutdown`], [`shutdown_on_idle`], and
/// [`shutdown_now`].
///
/// [`shutdown`]: struct.ThreadPool.html#method.shutdown
/// [`shutdown_on_idle`]: struct.ThreadPool.html#method.shutdown_on_idle
/// [`shutdown_now`]: struct.ThreadPool.html#method.shutdown_now
#[derive(Debug)]
pub struct Shutdown {
inner: Arc<Mutex<Inner>>,
}
/// Shared state between `Shutdown` and `ShutdownTrigger`.
///
/// This is used for notifying the `Shutdown` future when `ShutdownTrigger` gets dropped.
#[derive(Debug)]
struct Inner {
/// The task to notify when the threadpool completes the shutdown process.
task: AtomicWaker,
/// `true` if the threadpool has been shut down.
completed: bool,
}
impl Shutdown {
pub(crate) fn new(trigger: &ShutdownTrigger) -> Shutdown {
Shutdown {
inner: trigger.inner.clone(),
}
}
/// Wait for the shutdown to complete
pub fn wait(self) {
let mut enter = crate::enter().unwrap();
enter.block_on(self);
}
}
impl Future for Shutdown {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let inner = self.inner.lock().unwrap();
if !inner.completed {
inner.task.register_by_ref(cx.waker());
Poll::Pending
} else {
Poll::Ready(())
}
}
}
/// When dropped, cleans up threadpool's resources and completes the shutdown process.
#[derive(Debug)]
pub(crate) struct ShutdownTrigger {
inner: Arc<Mutex<Inner>>,
workers: Arc<[worker::Entry]>,
queue: Arc<Injector<Arc<Task>>>,
}
unsafe impl Send for ShutdownTrigger {}
unsafe impl Sync for ShutdownTrigger {}
impl ShutdownTrigger {
pub(crate) fn new(
workers: Arc<[worker::Entry]>,
queue: Arc<Injector<Arc<Task>>>,
) -> ShutdownTrigger {
ShutdownTrigger {
inner: Arc::new(Mutex::new(Inner {
task: AtomicWaker::new(),
completed: false,
})),
workers,
queue,
}
}
}
impl Drop for ShutdownTrigger {
fn drop(&mut self) {
// Drain the global task queue.
while !self.queue.steal().is_empty() {}
// Drop the remaining incomplete tasks and parkers assosicated with workers.
for worker in self.workers.iter() {
worker.shutdown();
}
// Notify the task interested in shutdown.
let mut inner = self.inner.lock().unwrap();
inner.completed = true;
inner.task.wake();
}
}
@@ -1,499 +0,0 @@
use super::super::pool::Pool;
use super::super::task::{BlockingState, Task};
use crate::threadpool::BlockingError;
use std::cell::UnsafeCell;
use std::fmt;
use std::ptr;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release};
use std::sync::Arc;
use std::task::Poll;
use std::thread;
/// Manages the state around entering a blocking section and tasks that are
/// queued pending the ability to block.
///
/// This is a hybrid counter and intrusive mpsc channel (like `Queue`).
#[derive(Debug)]
pub(crate) struct Blocking {
/// Queue head.
///
/// This is either the current remaining capacity for blocking sections
/// **or** if the max has been reached, the head of a pending blocking
/// capacity channel of tasks.
///
/// When this points to a task, it represents a strong reference, i.e.
/// `Arc<Task>`.
state: AtomicUsize,
/// Tail pointer. This is `Arc<Task>` unless it points to `stub`.
tail: UnsafeCell<*mut Task>,
/// Stub pointer, used as part of the intrusive mpsc channel algorithm
/// described by 1024cores.
stub: Box<Task>,
/// The channel algorithm is MPSC. This means that, in order to pop tasks,
/// coordination is required.
///
/// Since it doesn't matter *which* task pops & notifies the queued task, we
/// can avoid a full mutex and make the "lock" lock free.
///
/// Instead, threads race to set the "entered" bit. When the transition is
/// successfully made, the thread has permission to pop tasks off of the
/// queue. If a thread loses the race, instead of waiting to pop a task, it
/// signals to the winning thread that it should pop an additional task.
lock: AtomicUsize,
}
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub(crate) enum CanBlock {
/// Blocking capacity has been allocated to this task.
///
/// The capacity allocation is initially checked before a task is polled. If
/// capacity has been allocated, it is consumed and tracked as `Allocated`.
Allocated,
/// Allocation capacity must be either available to the task when it is
/// polled or not available. This means that a task can only ask for
/// capacity once. This state is used to track a task that has not yet asked
/// for blocking capacity. When a task needs blocking capacity, if it is in
/// this state, it can immediately try to get an allocation.
CanRequest,
/// The task has requested blocking capacity, but none is available.
NoCapacity,
}
/// Decorates the `usize` value of `Blocking::state`, providing fns to
/// manipulate the state instead of requiring bit ops.
#[derive(Copy, Clone, Eq, PartialEq)]
struct State(usize);
/// Flag differentiating between remaining capacity and task pointers.
///
/// If we assume pointers are properly aligned, then the least significant bit
/// will always be zero. So, we use that bit to track if the value represents a
/// number.
const NUM_FLAG: usize = 1;
/// When representing "numbers", the state has to be shifted this much (to get
/// rid of the flag bit).
const NUM_SHIFT: usize = 1;
// ====== impl Blocking =====
//
impl Blocking {
/// Create a new `Blocking`.
pub(crate) fn new(capacity: usize) -> Blocking {
assert!(capacity > 0, "blocking capacity must be greater than zero");
let stub = Box::new(Task::stub());
let ptr = &*stub as *const _ as *mut _;
// Allocations are aligned
debug_assert!(ptr as usize & NUM_FLAG == 0);
// The initial state value. This starts at the max capacity.
let init = State::new(capacity);
Blocking {
state: AtomicUsize::new(init.into()),
tail: UnsafeCell::new(ptr),
stub,
lock: AtomicUsize::new(0),
}
}
/// Atomically either acquire blocking capacity or queue the task to be
/// notified once capacity becomes available.
///
/// The caller must ensure that `task` has not previously been queued to be
/// notified when capacity becomes available.
pub(crate) fn poll_blocking_capacity(
&self,
task: &Arc<Task>,
) -> Poll<Result<(), BlockingError>> {
// This requires atomically claiming blocking capacity and if none is
// available, queuing &task.
// The task cannot be queued at this point. The caller must ensure this.
debug_assert!(!BlockingState::from(task.blocking.load(Acquire)).is_queued());
// Don't bump the ref count unless necessary.
let mut strong: Option<*const Task> = None;
// Load the state
let mut curr: State = self.state.load(Acquire).into();
loop {
let mut next = curr;
if !next.claim_capacity(&self.stub) {
debug_assert!(curr.ptr().is_some());
// Unable to claim capacity, so we must queue `task` onto the
// channel.
//
// This guard also serves to ensure that queuing work that is
// only needed to run once only gets run once.
if strong.is_none() {
// First, transition the task to a "queued" state. This
// prevents double queuing.
//
// This is also the only thread that can set the queued flag
// at this point. And, the goal is for this to only be
// visible when the task node is polled from the channel.
// The memory ordering is established by MPSC queue
// operation.
//
// Note that, if the task doesn't get queued (because the
// CAS fails and capacity is now available) then this flag
// must be unset. Again, there is no race because until the
// task is queued, no other thread can see it.
let prev = BlockingState::toggle_queued(&task.blocking, Relaxed);
debug_assert!(!prev.is_queued());
// Bump the ref count
strong = Some(Arc::into_raw(task.clone()));
// Set the next pointer. This does not require an atomic
// operation as this node is not currently accessible to
// other threads via the queue.
task.next_blocking.store(ptr::null_mut(), Relaxed);
}
let ptr = strong.unwrap();
// Update the head to point to the new node. We need to see the
// previous node in order to update the next pointer as well as
// release `task` to any other threads calling `push`.
next.set_ptr(ptr);
}
debug_assert_ne!(curr.0, 0);
debug_assert_ne!(next.0, 0);
let actual = self
.state
.compare_and_swap(curr.into(), next.into(), AcqRel)
.into();
if curr == actual {
break;
}
curr = actual;
}
match curr.ptr() {
Some(prev) => {
let ptr = strong.unwrap();
// Finish pushing
unsafe {
(*prev).next_blocking.store(ptr as *mut _, Release);
}
// The node was queued to be notified once capacity is made
// available.
Poll::Pending
}
None => {
debug_assert!(curr.remaining_capacity() > 0);
// If `strong` is set, gotta undo a bunch of work
if let Some(ptr) = strong {
let _ = unsafe { Arc::from_raw(ptr) };
// Unset the queued flag.
let prev = BlockingState::toggle_queued(&task.blocking, Relaxed);
debug_assert!(prev.is_queued());
}
// Capacity has been obtained
Poll::Ready(Ok(()))
}
}
}
unsafe fn push_stub(&self) {
let task: *mut Task = &*self.stub as *const _ as *mut _;
// Set the next pointer. This does not require an atomic operation as
// this node is not accessible. The write will be flushed with the next
// operation
(*task).next_blocking.store(ptr::null_mut(), Relaxed);
// Update the head to point to the new node. We need to see the previous
// node in order to update the next pointer as well as release `task`
// to any other threads calling `push`.
let prev = self.state.swap(task as usize, AcqRel);
// The stub is only pushed when there are pending tasks. Because of
// this, the state must *always* be in pointer mode.
debug_assert!(State::from(prev).is_ptr());
let prev = prev as *const Task;
// We don't want the *existing* pointer to be a stub.
debug_assert_ne!(prev, task);
// Release `task` to the consume end.
(*prev).next_blocking.store(task, Release);
}
pub(crate) fn notify_task(&self, pool: &Arc<Pool>) {
let prev = self.lock.fetch_add(1, AcqRel);
if prev != 0 {
// Another thread has the lock and will be responsible for notifying
// pending tasks.
return;
}
let mut dec = 1;
loop {
let mut remaining_pops = dec;
while remaining_pops > 0 {
remaining_pops -= 1;
let task = match self.pop(remaining_pops) {
Some(t) => t,
None => break,
};
Task::notify_blocking(task, pool);
}
// Decrement the number of handled notifications
let actual = self.lock.fetch_sub(dec, AcqRel);
if actual == dec {
break;
}
// This can only be greater than expected as we are the only thread
// that is decrementing.
debug_assert!(actual > dec);
dec = actual - dec;
}
}
/// Pop a task
///
/// `rem` represents the remaining number of times the caller will pop. If
/// there are no more tasks to pop, `rem` is used to set the remaining
/// capacity.
fn pop(&self, rem: usize) -> Option<Arc<Task>> {
'outer: loop {
unsafe {
let mut tail = *self.tail.get();
let mut next = (*tail).next_blocking.load(Acquire);
let stub = &*self.stub as *const _ as *mut _;
if tail == stub {
if next.is_null() {
// This loop is not part of the standard intrusive mpsc
// channel algorithm. This is where we atomically pop
// the last task and add `rem` to the remaining capacity.
//
// This modification to the pop algorithm works because,
// at this point, we have not done any work (only done
// reading). We have a *pretty* good idea that there is
// no concurrent pusher.
//
// The capacity is then atomically added by doing an
// AcqRel CAS on `state`. The `state` cell is the
// linchpin of the algorithm.
//
// By successfully CASing `head` w/ AcqRel, we ensure
// that, if any thread was racing and entered a push, we
// see that and abort pop, retrying as it is
// "inconsistent".
let mut curr: State = self.state.load(Acquire).into();
loop {
if curr.has_task(&self.stub) {
// Inconsistent state, yield the thread and try
// again.
thread::yield_now();
continue 'outer;
}
let mut after = curr;
// +1 here because `rem` represents the number of
// pops that will come after the current one.
after.add_capacity(rem + 1, &self.stub);
let actual: State = self
.state
.compare_and_swap(curr.into(), after.into(), AcqRel)
.into();
if actual == curr {
// Successfully returned the remaining capacity
return None;
}
curr = actual;
}
}
*self.tail.get() = next;
tail = next;
next = (*next).next_blocking.load(Acquire);
}
if !next.is_null() {
*self.tail.get() = next;
// No ref_count inc is necessary here as this poll is paired
// with a `push` which "forgets" the handle.
return Some(Arc::from_raw(tail));
}
let state = self.state.load(Acquire);
// This must always be a pointer
debug_assert!(State::from(state).is_ptr());
if state != tail as usize {
// Try again
thread::yield_now();
continue 'outer;
}
self.push_stub();
next = (*tail).next_blocking.load(Acquire);
if !next.is_null() {
*self.tail.get() = next;
return Some(Arc::from_raw(tail));
}
thread::yield_now();
// Try again
}
}
}
}
// ====== impl State =====
impl State {
/// Return a new `State` representing the remaining capacity at the maximum
/// value.
fn new(capacity: usize) -> State {
State((capacity << NUM_SHIFT) | NUM_FLAG)
}
fn remaining_capacity(self) -> usize {
if !self.has_remaining_capacity() {
return 0;
}
self.0 >> 1
}
fn has_remaining_capacity(self) -> bool {
self.0 & NUM_FLAG == NUM_FLAG
}
fn has_task(self, stub: &Task) -> bool {
!(self.has_remaining_capacity() || self.is_stub(stub))
}
fn is_stub(self, stub: &Task) -> bool {
self.0 == stub as *const _ as usize
}
/// Try to claim blocking capacity.
///
/// # Return
///
/// Returns `true` if the capacity was claimed, `false` otherwise. If
/// `false` is returned, it can be assumed that `State` represents the head
/// pointer in the mpsc channel.
fn claim_capacity(&mut self, stub: &Task) -> bool {
if !self.has_remaining_capacity() {
return false;
}
debug_assert!(self.0 != 1);
self.0 -= 1 << NUM_SHIFT;
if self.0 == NUM_FLAG {
// Set the state to the stub pointer.
self.0 = stub as *const _ as usize;
}
true
}
/// Add blocking capacity.
fn add_capacity(&mut self, capacity: usize, stub: &Task) -> bool {
debug_assert!(capacity > 0);
if self.is_stub(stub) {
self.0 = (capacity << NUM_SHIFT) | NUM_FLAG;
true
} else if self.has_remaining_capacity() {
self.0 += capacity << NUM_SHIFT;
true
} else {
false
}
}
fn is_ptr(self) -> bool {
self.0 & NUM_FLAG == 0
}
fn ptr(self) -> Option<*const Task> {
if self.is_ptr() {
Some(self.0 as *const Task)
} else {
None
}
}
fn set_ptr(&mut self, ptr: *const Task) {
let ptr = ptr as usize;
debug_assert!(ptr & NUM_FLAG == 0);
self.0 = ptr
}
}
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 {
let mut fmt = fmt.debug_struct("State");
if self.is_ptr() {
fmt.field("ptr", &self.0);
} else {
fmt.field("remaining", &self.remaining_capacity());
}
fmt.finish()
}
}
@@ -1,89 +0,0 @@
use super::CanBlock;
use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};
/// State tracking task level state to support `blocking`.
///
/// This tracks two separate flags.
///
/// a) If the task is queued in the pending blocking channel. This prevents
/// double queuing (which would break the linked list).
///
/// b) If the task has been allocated capacity to block.
#[derive(Eq, PartialEq)]
pub(crate) struct BlockingState(usize);
const QUEUED: usize = 0b01;
const ALLOCATED: usize = 0b10;
impl BlockingState {
/// Create a new, default, `BlockingState`.
pub(crate) fn new() -> BlockingState {
BlockingState(0)
}
/// Returns `true` if the state represents the associated task being queued
/// in the pending blocking capacity channel
pub(crate) fn is_queued(&self) -> bool {
self.0 & QUEUED == QUEUED
}
/// Toggle the queued flag
///
/// Returns the state before the flag has been toggled.
pub(crate) fn toggle_queued(state: &AtomicUsize, ordering: Ordering) -> BlockingState {
state.fetch_xor(QUEUED, ordering).into()
}
/// Returns `true` if the state represents the associated task having been
/// allocated capacity to block.
pub(crate) fn is_allocated(&self) -> bool {
self.0 & ALLOCATED == ALLOCATED
}
/// Atomically consume the capacity allocation and return if the allocation
/// was present.
///
/// If this returns `true`, then the task has the ability to block for the
/// duration of the `poll`.
pub(crate) fn consume_allocation(state: &AtomicUsize, ordering: Ordering) -> CanBlock {
let state: Self = state.fetch_and(!ALLOCATED, ordering).into();
if state.is_allocated() {
CanBlock::Allocated
} else if state.is_queued() {
CanBlock::NoCapacity
} else {
CanBlock::CanRequest
}
}
pub(crate) fn notify_blocking(state: &AtomicUsize, ordering: Ordering) {
let prev: Self = state.fetch_xor(ALLOCATED | QUEUED, ordering).into();
debug_assert!(prev.is_queued());
debug_assert!(!prev.is_allocated());
}
}
impl From<usize> for BlockingState {
fn from(src: usize) -> BlockingState {
BlockingState(src)
}
}
impl From<BlockingState> for usize {
fn from(src: BlockingState) -> usize {
src.0
}
}
impl fmt::Debug for BlockingState {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("BlockingState")
.field("is_queued", &self.is_queued())
.field("is_allocated", &self.is_allocated())
.finish()
}
}
-323
View File
@@ -1,323 +0,0 @@
mod blocking;
mod blocking_state;
mod state;
pub(crate) use self::blocking::{Blocking, CanBlock};
use self::blocking_state::BlockingState;
use self::state::State;
use super::pool::Pool;
use super::waker::Waker;
use futures_util::task;
use std::cell::{Cell, UnsafeCell};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
use std::sync::atomic::{AtomicPtr, AtomicUsize};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::{fmt, panic, ptr};
/// Harness around a future.
///
/// This also behaves as a node in the inbound work queue and the blocking
/// queue.
pub(crate) struct Task {
/// Task lifecycle state
state: AtomicUsize,
/// Task blocking related state
blocking: AtomicUsize,
/// Next pointer in the queue of tasks pending blocking capacity.
next_blocking: AtomicPtr<Task>,
/// ID of the worker that polled this task first.
///
/// This field can be a `Cell` because it's only accessed by the worker thread that is
/// executing the task.
///
/// The worker ID is represented by a `u32` rather than `usize` in order to save some space
/// on 64-bit platforms.
pub(crate) reg_worker: Cell<Option<u32>>,
/// The key associated with this task in the `Slab` it was registered in.
///
/// This field can be a `Cell` because it's only accessed by the worker thread that has
/// registered the task.
pub(crate) reg_index: Cell<usize>,
/// Store the future at the head of the struct
///
/// The future is dropped immediately when it transitions to Complete
future: UnsafeCell<Option<BoxFuture>>,
}
#[derive(Debug)]
pub(crate) enum Run {
Idle,
Schedule,
Complete,
}
type BoxFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
// ===== impl Task =====
impl Task {
/// Create a new `Task` as a harness for `future`.
pub(crate) fn new(future: BoxFuture) -> Task {
Task {
state: AtomicUsize::new(State::new().into()),
blocking: AtomicUsize::new(BlockingState::new().into()),
next_blocking: AtomicPtr::new(ptr::null_mut()),
reg_worker: Cell::new(None),
reg_index: Cell::new(0),
future: UnsafeCell::new(Some(future)),
}
}
/// Create a fake `Task` to be used as part of the intrusive mpsc channel
/// algorithm.
fn stub() -> Task {
let future = Box::pin(Empty) as BoxFuture;
Task {
state: AtomicUsize::new(State::stub().into()),
blocking: AtomicUsize::new(BlockingState::new().into()),
next_blocking: AtomicPtr::new(ptr::null_mut()),
reg_worker: Cell::new(None),
reg_index: Cell::new(0),
future: UnsafeCell::new(Some(future)),
}
}
/// Execute the task returning `Run::Schedule` if the task needs to be
/// scheduled again.
///
// tracing macro expansion adds enough branches to make clippy angry here.
#[allow(clippy::cognitive_complexity)] // https://github.com/rust-lang/rust-clippy/issues/3900
pub(crate) fn run(me: &Arc<Task>, pool: &Arc<Pool>) -> Run {
use self::State::*;
#[cfg(feature = "tracing")]
use std::sync::atomic::Ordering::Relaxed;
// Transition task to running state. At this point, the task must be
// scheduled.
let actual: State = me
.state
.compare_and_swap(Scheduled.into(), Running.into(), AcqRel)
.into();
match actual {
Scheduled => {}
_ => panic!("unexpected task state; {:?}", actual),
}
let span = trace_span!("Task::run");
let _enter = span.enter();
trace!(state = ?State::from(me.state.load(Relaxed)));
// The transition to `Running` done above ensures that a lock on the
// future has been obtained.
let fut = unsafe { &mut (*me.future.get()) };
// This block deals with the future panicking while being polled.
//
// If the future panics, then the drop handler must be called such that
// `thread::panicking() -> true`. To do this, the future is dropped from
// within the catch_unwind block.
let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
struct Guard<'a>(&'a mut Option<BoxFuture>, bool);
impl Drop for Guard<'_> {
fn drop(&mut self) {
// This drops the future
if self.1 {
let _ = self.0.take();
}
}
}
let mut g = Guard(fut, true);
let waker = task::waker(Arc::new(Waker {
task: me.clone(),
pool: pool.clone(),
}));
let mut cx = Context::from_waker(&waker);
let ret = g.0.as_mut().unwrap().as_mut().poll(&mut cx);
g.1 = false;
ret
}));
match res {
Ok(Poll::Ready(_)) | Err(_) => {
trace!("task complete");
// The future has completed. Drop it immediately to free
// resources and run drop handlers.
//
// The `Task` harness will stay around longer if it is contained
// by any of the various queues.
me.drop_future();
// Transition to the completed state
me.state.store(State::Complete.into(), Release);
if let Err(panic_err) = res {
if let Some(ref f) = pool.config.panic_handler {
f(panic_err);
}
}
Run::Complete
}
Ok(Poll::Pending) => {
trace!("not ready");
// Attempt to transition from Running -> Idle, if successful,
// then the task does not need to be scheduled again. If the CAS
// fails, then the task has been unparked concurrent to running,
// in which case it transitions immediately back to scheduled
// and we return `true`.
let prev: State = me
.state
.compare_and_swap(Running.into(), Idle.into(), AcqRel)
.into();
match prev {
Running => Run::Idle,
Notified => {
me.state.store(Scheduled.into(), Release);
Run::Schedule
}
_ => unreachable!(),
}
}
}
}
/// Aborts this task.
///
/// This is called when the threadpool shuts down and the task has already beed polled but not
/// completed.
pub(crate) fn abort(&self) {
use self::State::*;
let mut state = self.state.load(Acquire).into();
loop {
match state {
Idle | Scheduled => {}
Running | Notified | Complete | Aborted => {
// It is assumed that no worker threads are running so the task must be either
// in the idle or scheduled state.
panic!("unexpected state while aborting task: {:?}", state);
}
}
let actual = self
.state
.compare_and_swap(state.into(), Aborted.into(), AcqRel)
.into();
if actual == state {
// The future has been aborted. Drop it immediately to free resources and run drop
// handlers.
self.drop_future();
break;
}
state = actual;
}
}
/// Notify the task it has been allocated blocking capacity
pub(crate) fn notify_blocking(me: Arc<Task>, pool: &Arc<Pool>) {
BlockingState::notify_blocking(&me.blocking, AcqRel);
Task::schedule(&me, pool);
}
pub(crate) fn schedule(me: &Arc<Self>, pool: &Arc<Pool>) {
if me.schedule2() {
let task = me.clone();
pool.submit(task, &pool);
}
}
/// Transition the task state to scheduled.
///
/// Returns `true` if the caller is permitted to schedule the task.
fn schedule2(&self) -> bool {
use self::State::*;
loop {
// Scheduling can only be done from the `Idle` state.
let actual = self
.state
.compare_and_swap(Idle.into(), Scheduled.into(), AcqRel)
.into();
match actual {
Idle => return true,
Running => {
// The task is already running on another thread. Transition
// the state to `Notified`. If this CAS fails, then restart
// the logic again from `Idle`.
let actual = self
.state
.compare_and_swap(Running.into(), Notified.into(), AcqRel)
.into();
match actual {
Idle => continue,
_ => return false,
}
}
Complete | Aborted | Notified | Scheduled => return false,
}
}
}
/// Consumes any allocated capacity to block.
///
/// Returns `true` if capacity was allocated, `false` otherwise.
pub(crate) fn consume_blocking_allocation(&self) -> CanBlock {
// This flag is the primary point of coordination. The queued flag
// happens "around" setting the blocking capacity.
BlockingState::consume_allocation(&self.blocking, AcqRel)
}
/// Drop the future
///
/// This must only be called by the thread that successfully transitioned
/// the future state to `Running`.
fn drop_future(&self) {
let _ = unsafe { (*self.future.get()).take() };
}
}
impl fmt::Debug for Task {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Task")
.field("state", &self.state)
.field("future", &"BoxFuture")
.finish()
}
}
struct Empty;
impl Future for Empty {
type Output = ();
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> {
// Never used
unreachable!();
}
}
@@ -1,57 +0,0 @@
#[repr(usize)]
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub(crate) enum State {
/// Task is currently idle
Idle = 0,
/// Task is currently running
Running = 1,
/// Task is currently running, but has been notified that it must run again.
Notified = 2,
/// Task has been scheduled
Scheduled = 3,
/// Task is complete
Complete = 4,
/// Task was aborted because the thread pool has been shut down
Aborted = 5,
}
// ===== impl State =====
impl State {
/// Returns the initial task state.
///
/// Tasks start in the scheduled state as they are immediately scheduled on
/// creation.
pub(crate) fn new() -> State {
State::Scheduled
}
pub(crate) fn stub() -> State {
State::Idle
}
}
impl From<usize> for State {
fn from(src: usize) -> Self {
use self::State::*;
debug_assert!(
src >= Idle as usize && src <= Aborted as usize,
"actual={}",
src
);
unsafe { ::std::mem::transmute(src) }
}
}
impl From<State> for usize {
fn from(src: State) -> Self {
src as usize
}
}
@@ -1,226 +0,0 @@
use super::builder::Builder;
use super::pool::Pool;
use super::sender::Sender;
use super::shutdown::{Shutdown, ShutdownTrigger};
use std::future::Future;
use std::sync::Arc;
/// Work-stealing based thread pool for executing futures.
///
/// If a `ThreadPool` instance is dropped without explicitly being shutdown,
/// `shutdown_now` is called implicitly, forcing all tasks that have not yet
/// completed to be dropped.
///
/// Create `ThreadPool` instances using `Builder`.
#[derive(Debug)]
pub struct ThreadPool {
inner: Option<Inner>,
}
#[derive(Debug)]
struct Inner {
sender: Sender,
trigger: Arc<ShutdownTrigger>,
}
impl ThreadPool {
/// Create a new `ThreadPool` with default values.
///
/// Use [`Builder`] for creating a configured thread pool.
///
/// [`Builder`]: struct.Builder.html
pub fn new() -> ThreadPool {
Builder::new().build()
}
pub(crate) fn new2(pool: Arc<Pool>, trigger: Arc<ShutdownTrigger>) -> ThreadPool {
ThreadPool {
inner: Some(Inner {
sender: Sender { pool },
trigger,
}),
}
}
/// Spawn a future onto the thread pool.
///
/// This function takes ownership of the future and randomly assigns it to a
/// worker thread. The thread will then start executing the future.
///
/// # Examples
///
/// ```rust
/// use tokio_executor::threadpool::ThreadPool;
///
/// // Create a thread pool with default configuration values
/// let thread_pool = ThreadPool::new();
///
/// thread_pool.spawn(async {
/// println!("called from a worker thread");
/// });
///
/// // Gracefully shutdown the threadpool
/// thread_pool.shutdown().wait();
/// ```
///
/// # Panics
///
/// This function panics if the spawn fails. Use [`Sender::spawn`] for a
/// version that returns a `Result` instead of panicking.
pub fn spawn<F>(&self, future: F)
where
F: Future<Output = ()> + Send + 'static,
{
self.sender().spawn(future).unwrap();
}
/*
* TODO: Bring back
/// Spawn a future on to the thread pool, return a future representing
/// the produced value.
///
/// The SpawnHandle returned is a future that is a proxy for future itself.
/// When future completes on this thread pool then the SpawnHandle will itself
/// be resolved.
///
/// # Examples
///
/// ```rust
/// # use tokio_executor::threadpool::ThreadPool;
/// use futures::future::{Future, lazy};
///
/// // Create a thread pool with default configuration values
/// let thread_pool = ThreadPool::new();
///
/// let handle = thread_pool.spawn_handle(lazy(|| Ok::<_, ()>(42)));
///
/// let value = handle.wait().unwrap();
/// assert_eq!(value, 42);
///
/// // Gracefully shutdown the threadpool
/// thread_pool.shutdown().wait().unwrap();
/// ```
///
/// # Panics
///
/// This function panics if the spawn fails.
pub fn spawn_handle<F>(&self, future: F) -> SpawnHandle<F::Item, F::Error>
where
F: Future + Send + 'static,
F::Item: Send + 'static,
F::Error: Send + 'static,
{
SpawnHandle(oneshot::spawn(future, self.sender()))
}
*/
/// Return a reference to the sender handle
///
/// The handle is used to spawn futures onto the thread pool. It also
/// implements the `Executor` trait.
pub fn sender(&self) -> &Sender {
&self.inner.as_ref().unwrap().sender
}
/// Return a mutable reference to the sender handle
pub fn sender_mut(&mut self) -> &mut Sender {
&mut self.inner.as_mut().unwrap().sender
}
/// Shutdown the pool once it becomes idle.
///
/// Idle is defined as the completion of all futures that have been spawned
/// onto the thread pool. There may still be outstanding handles when the
/// thread pool reaches an idle state.
///
/// Once the idle state is reached, calling `spawn` on any outstanding
/// handle will result in an error. All worker threads are signaled and will
/// shutdown. The returned future completes once all worker threads have
/// completed the shutdown process.
pub fn shutdown_on_idle(mut self) -> Shutdown {
let inner = self.inner.take().unwrap();
inner.sender.pool.shutdown(false, false);
Shutdown::new(&inner.trigger)
}
/// Shutdown the pool
///
/// This prevents the thread pool from accepting new tasks but will allow
/// any existing tasks to complete.
///
/// Calling `spawn` on any outstanding handle will result in an error. All
/// worker threads are signaled and will shutdown. The returned future
/// completes once all worker threads have completed the shutdown process.
pub fn shutdown(mut self) -> Shutdown {
let inner = self.inner.take().unwrap();
inner.sender.pool.shutdown(true, false);
Shutdown::new(&inner.trigger)
}
/// Shutdown the pool immediately
///
/// This will prevent the thread pool from accepting new tasks **and**
/// abort any tasks that are currently running on the thread pool.
///
/// Calling `spawn` on any outstanding handle will result in an error. All
/// worker threads are signaled and will shutdown. The returned future
/// completes once all worker threads have completed the shutdown process.
pub fn shutdown_now(mut self) -> Shutdown {
let inner = self.inner.take().unwrap();
inner.sender.pool.shutdown(true, true);
Shutdown::new(&inner.trigger)
}
}
impl Drop for ThreadPool {
fn drop(&mut self) {
if let Some(inner) = self.inner.take() {
// Begin the shutdown process.
inner.sender.pool.shutdown(true, true);
let shutdown = Shutdown::new(&inner.trigger);
// Drop `inner` in order to drop its shutdown trigger.
drop(inner);
// Wait until all worker threads terminate and the threadpool's resources clean up.
let mut enter = match crate::enter() {
Ok(e) => e,
Err(_) => return,
};
enter.block_on(shutdown);
}
}
}
impl Default for ThreadPool {
fn default() -> Self {
Self::new()
}
}
/*
* TODO: Bring back
/// Handle returned from ThreadPool::spawn_handle.
///
/// This handle is a future representing the completion of a different future
/// spawned on to the thread pool. Created through the ThreadPool::spawn_handle
/// function this handle will resolve when the future provided resolves on the
/// thread pool.
#[derive(Debug)]
pub struct SpawnHandle<T, E>(oneshot::SpawnHandle<T, E>);
impl<T, E> Future for SpawnHandle<T, E> {
type Item = T;
type Error = E;
fn poll(&mut self) -> Poll<T, E> {
self.0.poll()
}
}
*/
-24
View File
@@ -1,24 +0,0 @@
use super::pool::Pool;
use super::task::Task;
use futures_util::task::ArcWake;
use std::sync::Arc;
/// Implements the future `Waker` API.
///
/// This is how external events are able to signal the task, informing it to try
/// to poll the future again.
#[derive(Debug)]
pub(crate) struct Waker {
pub(crate) pool: Arc<Pool>,
pub(crate) task: Arc<Task>,
}
unsafe impl Send for Waker {}
unsafe impl Sync for Waker {}
impl ArcWake for Waker {
fn wake_by_ref(me: &Arc<Self>) {
Task::schedule(&me.task, &me.pool);
}
}
@@ -1,329 +0,0 @@
use super::super::park::{BoxPark, BoxUnpark};
use super::super::task::Task;
use super::state::{State, PUSHED_MASK};
use crossbeam_deque::{Steal, Stealer, Worker};
use crossbeam_queue::SegQueue;
use crossbeam_utils::CachePadded;
use slab::Slab;
use std::cell::UnsafeCell;
use std::fmt;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
// TODO: None of the fields should be public
//
// It would also be helpful to split up the state across what fields /
// operations are thread-safe vs. which ones require ownership of the worker.
pub(crate) struct WorkerEntry {
// Worker state. This is mutated when notifying the worker.
//
// The `usize` value is deserialized to a `worker::State` instance. See
// comments on that type.
pub(crate) state: CachePadded<AtomicUsize>,
// Next entry in the parked Trieber stack
next_sleeper: UnsafeCell<usize>,
// Worker half of deque
pub(crate) worker: Worker<Arc<Task>>,
// Stealer half of deque
stealer: Stealer<Arc<Task>>,
// Thread parker
park: UnsafeCell<Option<BoxPark>>,
// Thread unparker
unpark: UnsafeCell<Option<BoxUnpark>>,
// Tasks that have been first polled by this worker, but not completed yet.
running_tasks: UnsafeCell<Slab<Arc<Task>>>,
// Tasks that have been first polled by this worker, but completed by another worker.
remotely_completed_tasks: SegQueue<Arc<Task>>,
// Set to `true` when `remotely_completed_tasks` has tasks that need to be removed from
// `running_tasks`.
needs_drain: AtomicBool,
}
impl WorkerEntry {
pub(crate) fn new(park: BoxPark, unpark: BoxUnpark) -> Self {
let w = Worker::new_fifo();
let s = w.stealer();
WorkerEntry {
state: CachePadded::new(AtomicUsize::new(State::default().into())),
next_sleeper: UnsafeCell::new(0),
worker: w,
stealer: s,
park: UnsafeCell::new(Some(park)),
unpark: UnsafeCell::new(Some(unpark)),
running_tasks: UnsafeCell::new(Slab::new()),
remotely_completed_tasks: SegQueue::new(),
needs_drain: AtomicBool::new(false),
}
}
/// Atomically unset the pushed flag.
///
/// # Return
///
/// The state *before* the push flag is unset.
///
/// # Ordering
///
/// The specified ordering is established on the entry's state variable.
pub(crate) fn fetch_unset_pushed(&self, ordering: Ordering) -> State {
self.state.fetch_and(!PUSHED_MASK, ordering).into()
}
/// Submit a task to this worker while currently on the same thread that is
/// running the worker.
#[inline]
pub(crate) fn submit_internal(&self, task: Arc<Task>) {
self.push_internal(task);
}
/// Notifies the worker and returns `false` if it needs to be spawned.
///
/// # Ordering
///
/// The `state` must have been obtained with an `Acquire` ordering.
#[inline]
pub(crate) fn notify(&self, mut state: State) -> bool {
use super::Lifecycle::*;
loop {
let mut next = state;
next.notify();
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if state == actual {
break;
}
state = actual;
}
match state.lifecycle() {
Sleeping => {
// The worker is currently sleeping, the condition variable must
// be signaled
self.unpark();
true
}
Shutdown => false,
Running | Notified | Signaled => {
// In these states, the worker is active and will eventually see
// the task that was just submitted.
true
}
}
}
/// Signals to the worker that it should stop
///
/// `state` is the last observed state for the worker. This allows skipping
/// the initial load from the state atomic.
///
/// # Return
///
/// Returns `Ok` when the worker was successfully signaled.
///
/// Returns `Err` if the worker has already terminated.
pub(crate) fn signal_stop(&self, mut state: State) {
use super::Lifecycle::*;
// Transition the worker state to signaled
loop {
let mut next = state;
match state.lifecycle() {
Shutdown => {
return;
}
Running | Sleeping => {}
Notified | Signaled => {
// These two states imply that the worker is active, thus it
// will eventually see the shutdown signal, so we don't need
// to do anything.
//
// The worker is forced to see the shutdown signal
// eventually as:
//
// a) No more work will arrive
// b) The shutdown signal is stored as the head of the
// sleep, stack which will prevent the worker from going to
// sleep again.
return;
}
}
next.set_lifecycle(Signaled);
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
break;
}
state = actual;
}
// Wakeup the worker
self.unpark();
}
/// Pop a task
///
/// This **must** only be called by the thread that owns the worker entry.
/// This function is not `Sync`.
#[inline]
pub(crate) fn pop_task(&self) -> Option<Arc<Task>> {
self.worker.pop()
}
/// Steal tasks
///
/// This is called by *other* workers to steal a task for processing. This
/// function is `Sync`.
///
/// At the same time, this method steals some additional tasks and moves
/// them into `dest` in order to balance the work distribution among
/// workers.
pub(crate) fn steal_tasks(&self, dest: &Self) -> Steal<Arc<Task>> {
self.stealer.steal_batch_and_pop(&dest.worker)
}
/// Drain (and drop) all tasks that are queued for work.
///
/// This is called when the pool is shutting down.
pub(crate) fn drain_tasks(&self) {
while self.worker.pop().is_some() {}
}
/// Parks the worker thread.
pub(crate) fn park(&self) {
if let Some(park) = unsafe { (*self.park.get()).as_mut() } {
park.park().unwrap();
}
}
/// Parks the worker thread for at most `duration`.
pub(crate) fn park_timeout(&self, duration: Duration) {
if let Some(park) = unsafe { (*self.park.get()).as_mut() } {
park.park_timeout(duration).unwrap();
}
}
/// Unparks the worker thread.
#[inline]
pub(crate) fn unpark(&self) {
if let Some(park) = unsafe { (*self.unpark.get()).as_ref() } {
park.unpark();
}
}
/// Registers a task in this worker.
///
/// Called when the task is being polled for the first time.
#[inline]
pub(crate) fn register_task(&self, task: &Arc<Task>) {
let running_tasks = unsafe { &mut *self.running_tasks.get() };
let key = running_tasks.insert(task.clone());
task.reg_index.set(key);
}
/// Unregisters a task from this worker.
///
/// Called when the task is completed and was previously registered in this worker.
#[inline]
pub(crate) fn unregister_task(&self, task: Arc<Task>) {
let running_tasks = unsafe { &mut *self.running_tasks.get() };
running_tasks.remove(task.reg_index.get());
self.drain_remotely_completed_tasks();
}
/// Unregisters a task from this worker.
///
/// Called when the task is completed by another worker and was previously registered in this
/// worker.
#[inline]
pub(crate) fn remotely_complete_task(&self, task: Arc<Task>) {
self.remotely_completed_tasks.push(task);
self.needs_drain.store(true, Release);
}
/// Drops the remaining incomplete tasks and the parker associated with this worker.
///
/// This function is called by the shutdown trigger.
pub(crate) fn shutdown(&self) {
self.drain_remotely_completed_tasks();
// Abort all incomplete tasks.
let running_tasks = unsafe { &mut *self.running_tasks.get() };
for (_, task) in running_tasks.iter() {
task.abort();
}
running_tasks.clear();
unsafe {
*self.park.get() = None;
*self.unpark.get() = None;
}
}
/// Drains the `remotely_completed_tasks` queue and removes tasks from `running_tasks`.
#[inline]
fn drain_remotely_completed_tasks(&self) {
if self.needs_drain.compare_and_swap(true, false, Acquire) {
let running_tasks = unsafe { &mut *self.running_tasks.get() };
while let Ok(task) = self.remotely_completed_tasks.pop() {
running_tasks.remove(task.reg_index.get());
}
}
}
#[inline]
pub(crate) fn push_internal(&self, task: Arc<Task>) {
self.worker.push(task);
}
#[inline]
pub(crate) fn next_sleeper(&self) -> usize {
unsafe { *self.next_sleeper.get() }
}
#[inline]
pub(crate) fn set_next_sleeper(&self, val: usize) {
unsafe {
*self.next_sleeper.get() = val;
}
}
}
impl fmt::Debug for WorkerEntry {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("WorkerEntry")
.field("state", &self.state.load(Relaxed))
.field("next_sleeper", &"UnsafeCell<usize>")
.field("worker", &self.worker)
.field("stealer", &self.stealer)
.field("park", &"UnsafeCell<BoxPark>")
.field("unpark", &"BoxUnpark")
.finish()
}
}
-792
View File
@@ -1,792 +0,0 @@
mod entry;
mod stack;
mod state;
pub(crate) use self::entry::WorkerEntry as Entry;
pub(crate) use self::stack::Stack;
pub(crate) use self::state::{Lifecycle, State};
use super::pool::{self, BackupId, Pool};
use super::sender::Sender;
use super::shutdown::ShutdownTrigger;
use super::task::{self, CanBlock, Task};
use super::BlockingError;
use std::cell::Cell;
use std::marker::PhantomData;
use std::ptr;
use std::rc::Rc;
use std::sync::atomic::Ordering::{AcqRel, Acquire};
use std::sync::Arc;
use std::task::Poll;
use std::thread;
use std::time::Duration;
/// Thread worker
///
/// This is passed to the [`around_worker`] callback set on [`Builder`]. This
/// callback is only expected to call [`run`] on it.
///
/// [`Builder`]: struct.Builder.html
/// [`around_worker`]: struct.Builder.html#method.around_worker
/// [`run`]: struct.Worker.html#method.run
#[derive(Debug)]
pub struct Worker {
// Shared scheduler data
pub(crate) pool: Arc<Pool>,
// WorkerEntry index
pub(crate) id: WorkerId,
// Backup thread ID assigned to processing this worker.
backup_id: BackupId,
// Set to the task that is currently being polled by the worker. This is
// needed so that `blocking` blocks are able to interact with this task.
//
// This has to be a raw pointer to make it compile, but great care is taken
// when this is set.
current_task: CurrentTask,
// Set when the thread is in blocking mode.
is_blocking: Cell<bool>,
// Set when the worker should finalize on drop
should_finalize: Cell<bool>,
// Completes the shutdown process when the `ThreadPool` and all `Worker`s get dropped.
trigger: Arc<ShutdownTrigger>,
// Keep the value on the current thread.
_p: PhantomData<Rc<()>>,
}
/// Tracks the state related to the currently running task.
#[derive(Debug)]
struct CurrentTask {
/// This has to be a raw pointer to make it compile, but great care is taken
/// when this is set.
task: Cell<Option<*const Arc<Task>>>,
/// Tracks the blocking capacity allocation state.
can_block: Cell<CanBlock>,
}
/// Identifies a thread pool worker.
///
/// This identifier is unique scoped by the thread pool. It is possible that
/// different thread pool instances share worker identifier values.
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct WorkerId(pub(crate) usize);
// Pointer to the current worker info
thread_local!(static CURRENT_WORKER: Cell<*const Worker> = Cell::new(ptr::null()));
impl Worker {
pub(crate) fn new(
id: WorkerId,
backup_id: BackupId,
pool: Arc<Pool>,
trigger: Arc<ShutdownTrigger>,
) -> Worker {
Worker {
pool,
id,
backup_id,
current_task: CurrentTask::new(),
is_blocking: Cell::new(false),
should_finalize: Cell::new(false),
trigger,
_p: PhantomData,
}
}
pub(crate) fn is_blocking(&self) -> bool {
self.is_blocking.get()
}
/// Run the worker
///
/// Returns `true` if the thread should keep running as a `backup` thread.
pub(crate) fn do_run(&self) -> bool {
// Create another worker... It's ok, this is just a new type around
// `Pool` that is expected to stay on the current thread.
CURRENT_WORKER.with(|c| {
c.set(self as *const _);
let pool = self.pool.clone();
let mut sender = Sender { pool };
// Enter an execution context
let _enter = crate::enter().unwrap();
crate::with_default(&mut sender, || {
if let Some(ref callback) = self.pool.config.around_worker {
callback.call(self);
} else {
self.run();
}
});
});
// Can't be in blocking mode and finalization mode
debug_assert!(!self.is_blocking.get() || !self.should_finalize.get());
self.is_blocking.get()
}
pub(crate) fn with_current<F: FnOnce(Option<&Worker>) -> R, R>(f: F) -> R {
CURRENT_WORKER.with(move |c| {
let ptr = c.get();
if ptr.is_null() {
f(None)
} else {
f(Some(unsafe { &*ptr }))
}
})
}
/// Transition the current worker to a blocking worker
pub(crate) fn transition_to_blocking(&self) -> Poll<Result<(), BlockingError>> {
use self::CanBlock::*;
// If we get this far, then `current_task` has been set.
let task_ref = self.current_task.get_ref();
// First step is to acquire blocking capacity for the task.
match self.current_task.can_block() {
// Capacity to block has already been allocated to this task.
Allocated => {}
// The task has already requested capacity to block, but there is
// none yet available.
NoCapacity => return Poll::Pending,
// The task has yet to ask for capacity
CanRequest => {
// Atomically attempt to acquire blocking capacity, and if none
// is available, register the task to be notified once capacity
// becomes available.
match self.pool.poll_blocking_capacity(task_ref)? {
Poll::Ready(()) => {
self.current_task.set_can_block(Allocated);
}
Poll::Pending => {
self.current_task.set_can_block(NoCapacity);
return Poll::Pending;
}
}
}
}
// The task has been allocated blocking capacity. At this point, this is
// when the current thread transitions from a worker to a backup thread.
// To do so requires handing over the worker to another backup thread.
if self.is_blocking.get() {
// The thread is already in blocking mode, so there is nothing else
// to do. Return `Ready` and allow the caller to block the thread.
return Poll::Ready(Ok(()));
}
trace!("transition to blocking state");
// Transitioning to blocking requires handing over the worker state to
// another thread so that the work queue can continue to be processed.
self.pool.spawn_thread(self.id.clone(), &self.pool);
// Track that the thread has now fully entered the blocking state.
self.is_blocking.set(true);
Poll::Ready(Ok(()))
}
/// Transition from blocking
pub(crate) fn transition_from_blocking(&self) {
// TODO: Attempt to take ownership of the worker again.
}
/// Returns a reference to the worker's identifier.
///
/// This identifier is unique scoped by the thread pool. It is possible that
/// different thread pool instances share worker identifier values.
pub fn id(&self) -> &WorkerId {
&self.id
}
/// Run the worker
///
/// This function blocks until the worker is shutting down.
pub fn run(&self) {
const MAX_SPINS: usize = 3;
const LIGHT_SLEEP_INTERVAL: usize = 32;
let mut first = true;
let mut spin_cnt = 0;
let mut tick = 0;
while self.check_run_state(first) {
first = false;
// Run the next available task
if self.try_run_task(&self.pool) {
if self.is_blocking.get() {
// Exit out of the run state
return;
}
// Poll the reactor and the global queue every now and then to
// ensure no task gets left behind.
if tick % LIGHT_SLEEP_INTERVAL == 0 {
self.sleep_light();
}
tick = tick.wrapping_add(1);
spin_cnt = 0;
// As long as there is work, keep looping.
continue;
}
spin_cnt += 1;
// Yield the thread several times before it actually goes to sleep.
if spin_cnt <= MAX_SPINS {
thread::yield_now();
continue;
}
tick = 0;
spin_cnt = 0;
// Starting to get sleeeeepy
if !self.sleep() {
return;
}
// If there still isn't any work to do, shutdown the worker?
}
// The pool is terminating. However, transitioning the pool state to
// terminated is the very first step of the finalization process. Other
// threads may not see this state and try to spawn a new thread. To
// ensure consistency, before the current thread shuts down, it must
// return the backup token to the stack.
//
// The returned result is ignored because `Err` represents the pool
// shutting down. We are currently aware of this fact.
let _ = self.pool.release_backup(self.backup_id);
self.should_finalize.set(true);
}
/// Try to run a task
///
/// Returns `true` if work was found.
#[inline]
fn try_run_task(&self, pool: &Arc<Pool>) -> bool {
if self.try_run_owned_task(pool) {
return true;
}
self.try_steal_task(pool)
}
/// Checks the worker's current state, updating it as needed.
///
/// Returns `true` if the worker should run.
#[inline]
fn check_run_state(&self, first: bool) -> bool {
use self::Lifecycle::*;
debug_assert!(!self.is_blocking.get());
let mut state: State = self.entry().state.load(Acquire).into();
loop {
let pool_state: pool::State = self.pool.state.load(Acquire).into();
if pool_state.is_terminated() {
return false;
}
let mut next = state;
match state.lifecycle() {
Running => break,
Notified | Signaled => {
// transition back to running
next.set_lifecycle(Running);
}
Shutdown | Sleeping => {
// The worker should never be in these states when calling
// this function.
panic!("unexpected worker state; lifecycle={:?}", state.lifecycle());
}
}
let actual = self
.entry()
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
break;
}
state = actual;
}
// `first` is set to true the first time this function is called after
// the thread has started.
//
// This check is to handle the scenario where a worker gets signaled
// while it is already happily running. The `is_signaled` state is
// intended to wake up a worker that has been previously sleeping in
// effect increasing the number of active workers. If this is the first
// time `check_run_state` is called, then being in a signalled state is
// normal and the thread was started to handle it. However, if this is
// **not** the first time the fn was called, then the number of active
// workers has not been increased by the signal, so `signal_work` has to
// be called again to try to wake up another worker.
//
// For example, if the thread pool is configured to allow 4 workers.
// Worker 1 is processing tasks from its `deque`. Worker 2 receives its
// first task. Worker 2 will pick a random worker to signal. It does
// this by popping off the sleep stack, but there is no guarantee that
// workers on the sleep stack are actually sleeping. It is possible that
// Worker 1 gets signaled.
//
// Without this check, in the above case, no additional workers will get
// started, which results in the thread pool permanently being at 2
// workers even though it should reach 4.
if !first && state.is_signaled() {
trace!("Worker::check_run_state; delegate signal");
// This worker is not ready to be signaled, so delegate the signal
// to another worker.
self.pool.signal_work(&self.pool);
}
true
}
/// Runs the next task on this worker's queue.
///
/// Returns `true` if work was found.
fn try_run_owned_task(&self, pool: &Arc<Pool>) -> bool {
// Poll the internal queue for a task to run
match self.entry().pop_task() {
Some(task) => {
self.run_task(task, pool);
true
}
None => false,
}
}
/// Tries to steal a task from another worker.
///
/// Returns `true` if work was found
fn try_steal_task(&self, pool: &Arc<Pool>) -> bool {
use crossbeam_deque::Steal;
debug_assert!(!self.is_blocking.get());
let len = self.pool.workers.len();
let mut idx = self.pool.rand_usize() % len;
let mut found_work = false;
let start = idx;
loop {
if idx < len {
match self.pool.workers[idx].steal_tasks(self.entry()) {
Steal::Success(task) => {
trace!("stole task from another worker");
self.run_task(task, pool);
trace!(
message = "try_steal_task -- signal_work;",
self = self.id.0,
from = idx,
);
// Signal other workers that work is available
//
// TODO: Should this be called here or before
// `run_task`?
self.pool.signal_work(&self.pool);
return true;
}
Steal::Empty => {}
Steal::Retry => found_work = true,
}
idx += 1;
} else {
idx = 0;
}
if idx == start {
break;
}
}
found_work
}
fn run_task(&self, task: Arc<Task>, pool: &Arc<Pool>) {
use super::task::Run::*;
// If this is the first time this task is being polled, register it so that we can keep
// track of tasks that are in progress.
if task.reg_worker.get().is_none() {
task.reg_worker.set(Some(self.id.0 as u32));
self.entry().register_task(&task);
}
let run = self.run_task2(&task, pool);
// TODO: Try to claim back the worker state in case the backup thread
// did not start up fast enough. This is a performance optimization.
match run {
Idle => {}
Schedule => {
if self.is_blocking.get() {
// The future has been notified while it was running.
// However, the future also entered a blocking section,
// which released the worker state from this thread.
//
// This means that scheduling the future must be done from
// a point of view external to the worker set.
//
// We have to call `submit_external` instead of `submit`
// here because `self` is still set as the current worker.
self.pool.submit_external(task, &self.pool);
} else {
self.entry().push_internal(task);
}
}
Complete => {
let mut state: pool::State = self.pool.state.load(Acquire).into();
loop {
let mut next = state;
next.dec_num_futures();
let actual = self
.pool
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
trace!(message = "task complete;", state = ?next);
if state.num_futures() == 1 {
// If the thread pool has been flagged as shutdown,
// start terminating workers. This involves waking
// up any sleeping worker so that they can notice
// the shutdown state.
if next.is_terminated() {
self.pool.terminate_sleeping_workers();
}
}
// Find which worker polled this task first.
let worker = task.reg_worker.get().unwrap() as usize;
// Unregister the task from the worker it was registered in.
if !self.is_blocking.get() && worker == self.id.0 {
self.entry().unregister_task(task);
} else {
self.pool.workers[worker].remotely_complete_task(task);
}
// The worker's run loop will detect the shutdown state
// next iteration.
return;
}
state = actual;
}
}
}
}
/// Actually run the task. This is where `Worker::current_task` is set.
///
/// Great care is needed to ensure that `current_task` is unset in this
/// function.
fn run_task2(&self, task: &Arc<Task>, pool: &Arc<Pool>) -> task::Run {
struct Guard<'a> {
worker: &'a Worker,
}
impl Drop for Guard<'_> {
fn drop(&mut self) {
// A task is allocated at run when it was explicitly notified
// that the task has capacity to block. When this happens, that
// capacity is automatically allocated to the notified task.
// This capacity is "use it or lose it", so if the thread is not
// transitioned to blocking in this call, then another task has
// to be notified.
//
// If the task has consumed its blocking allocation but hasn't
// used it, it must be given to some other task instead.
if !self.worker.is_blocking.get() {
let can_block = self.worker.current_task.can_block();
if can_block == CanBlock::Allocated {
self.worker.pool.notify_blocking_task(&self.worker.pool);
}
}
self.worker.current_task.clear();
}
}
// Set `current_task`
self.current_task.set(task, CanBlock::CanRequest);
// Create the guard, this ensures that `current_task` is unset when the
// function returns, even if the return is caused by a panic.
let _g = Guard { worker: self };
Task::run(task, pool)
}
/// Put the worker to sleep
///
/// Returns `true` if woken up due to new work arriving.
// tracing macro expansion adds enough branches to make clippy angry here.
#[allow(clippy::cognitive_complexity)] // https://github.com/rust-lang/rust-clippy/issues/3900
fn sleep(&self) -> bool {
use self::Lifecycle::*;
// Putting a worker to sleep is a multipart operation. This is, in part,
// due to the fact that a worker can be notified without it being popped
// from the sleep stack. Extra care is needed to deal with this.
let span = trace_span!("Worker::sleep", idx = self.id.0, id = ?self.id);
let _e = span.enter();
let mut state: State = self.entry().state.load(Acquire).into();
// The first part of the sleep process is to transition the worker state
// to "pushed". Now, it may be that the worker is already pushed on the
// sleeper stack, in which case, we don't push again.
loop {
let mut next = state;
match state.lifecycle() {
Running => {
// Try setting the pushed state
next.set_pushed();
// Transition the worker state to sleeping
next.set_lifecycle(Sleeping);
}
Notified | Signaled => {
// No need to sleep, transition back to running and move on.
next.set_lifecycle(Running);
}
Shutdown | Sleeping => {
// The worker cannot transition to sleep when already in a
// sleeping state.
panic!("unexpected worker state; actual={:?}", state.lifecycle());
}
}
let actual = self
.entry()
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
if state.is_notified() {
// The previous state was notified, so we don't need to
// sleep.
return true;
}
if !state.is_pushed() {
debug_assert!(next.is_pushed());
trace!("push to stack");
// We obtained permission to push the worker into the
// sleeper queue.
if self.pool.push_sleeper(self.id.0).is_err() {
trace!("push to stack failed");
// The push failed due to the pool being terminated.
//
// This is true because the "work" being woken up for is
// shutting down.
return true;
}
}
break;
}
state = actual;
}
trace!("starting to sleep");
// Do a quick check to see if there are any notifications in the
// reactor or new tasks in the global queue. Since this call will
// clear the wakeup token, we need to check the state again and
// only after that go to sleep.
self.sleep_light();
// The state has been transitioned to sleeping, we can now wait by
// calling the parker. This is done in a loop as condvars can wakeup
// spuriously.
loop {
// Reload the state
state = self.entry().state.load(Acquire).into();
// If the worker has been notified, transition back to running.
match state.lifecycle() {
Sleeping => {
// Still sleeping. Park again.
}
Notified | Signaled => {
// Transition back to running
loop {
let mut next = state;
next.set_lifecycle(Running);
let actual = self
.entry()
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
return true;
}
state = actual;
}
}
Shutdown | Running => {
// To get here, the block above transitioned the state to
// `Sleeping`. No other thread can concurrently
// transition to `Shutdown` or `Running`.
unreachable!();
}
}
self.entry().park();
trace!("wakeup");
}
}
/// This doesn't actually put the thread to sleep. It calls
/// `park.park_timeout` with a duration of 0. This allows the park
/// implementation to perform any work that might be done on an interval.
///
/// Returns `true` if this worker has tasks in its queue.
fn sleep_light(&self) {
self.entry().park_timeout(Duration::from_millis(0));
use crossbeam_deque::Steal;
loop {
match self.pool.queue.steal_batch(&self.entry().worker) {
Steal::Success(()) => {
self.pool.signal_work(&self.pool);
break;
}
Steal::Empty => break,
Steal::Retry => {}
}
}
}
fn entry(&self) -> &Entry {
debug_assert!(!self.is_blocking.get());
&self.pool.workers[self.id.0]
}
}
impl Drop for Worker {
fn drop(&mut self) {
trace!(message = "shutting down thread", idx = self.id.0);
if self.should_finalize.get() {
// Drain the work queue
self.entry().drain_tasks();
}
}
}
// ===== impl CurrentTask =====
impl CurrentTask {
/// Returns a default `CurrentTask` representing no task.
fn new() -> CurrentTask {
CurrentTask {
task: Cell::new(None),
can_block: Cell::new(CanBlock::CanRequest),
}
}
/// Returns a reference to the task.
fn get_ref(&self) -> &Arc<Task> {
unsafe { &*self.task.get().unwrap() }
}
fn can_block(&self) -> CanBlock {
use self::CanBlock::*;
match self.can_block.get() {
Allocated => Allocated,
CanRequest | NoCapacity => {
let can_block = self.get_ref().consume_blocking_allocation();
self.can_block.set(can_block);
can_block
}
}
}
fn set_can_block(&self, can_block: CanBlock) {
self.can_block.set(can_block);
}
fn set(&self, task: &Arc<Task>, can_block: CanBlock) {
self.task.set(Some(task as *const _));
self.can_block.set(can_block);
}
/// Reset the `CurrentTask` to null state.
fn clear(&self) {
self.task.set(None);
self.can_block.set(CanBlock::CanRequest);
}
}
// ===== impl WorkerId =====
impl WorkerId {
/// Returns a `WorkerId` representing the worker entry at index `idx`.
pub(crate) fn new(idx: usize) -> WorkerId {
WorkerId(idx)
}
/// Returns this identifier represented as an integer.
///
/// Worker identifiers in a single thread pool are guaranteed to correspond to integers in the
/// range `0..pool_size`.
pub fn to_usize(&self) -> usize {
self.0
}
}
@@ -1,257 +0,0 @@
use super::super::config::MAX_WORKERS;
use super::super::worker;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed};
use std::{fmt, usize};
/// Lock-free stack of sleeping workers.
///
/// This is implemented as a Treiber stack and references to nodes are
/// `usize` values, indexing the entry in the `[worker::Entry]` array stored by
/// `Pool`. Each `Entry` instance maintains a `pushed` bit in its state. This
/// bit tracks if the entry is already pushed onto the stack or not. A single
/// entry can only be stored on the stack a single time.
///
/// By using indexes instead of pointers, that allows a much greater amount of
/// data to be used for the ABA guard (see correctness section of wikipedia
/// page).
///
/// Treiber stack: https://en.wikipedia.org/wiki/Treiber_Stack
#[derive(Debug)]
pub(crate) struct Stack {
state: AtomicUsize,
}
/// State related to the stack of sleeping workers.
///
/// - Parked head 16 bits
/// - Sequence remaining
///
/// The parked head value has a couple of special values:
///
/// - EMPTY: No sleepers
/// - TERMINATED: Don't spawn more threads
#[derive(Eq, PartialEq, Clone, Copy)]
pub struct State(usize);
/// Extracts the head of the worker stack from the scheduler state
///
/// The 16 relates to the value of MAX_WORKERS
const STACK_MASK: usize = ((1 << 16) - 1);
/// Used to mark the stack as empty
pub(crate) const EMPTY: usize = MAX_WORKERS;
/// Used to mark the stack as terminated
pub(crate) const TERMINATED: usize = EMPTY + 1;
/// How many bits the Treiber ABA guard is offset by
const ABA_GUARD_SHIFT: usize = 16;
#[cfg(target_pointer_width = "64")]
const ABA_GUARD_MASK: usize = (1 << (64 - ABA_GUARD_SHIFT)) - 1;
#[cfg(target_pointer_width = "32")]
const ABA_GUARD_MASK: usize = (1 << (32 - ABA_GUARD_SHIFT)) - 1;
// ===== impl Stack =====
impl Stack {
/// Create a new `Stack` representing the empty state.
pub(crate) fn new() -> Stack {
let state = AtomicUsize::new(State::new().into());
Stack { state }
}
/// Push a worker onto the stack
///
/// # Return
///
/// Returns `Ok` on success.
///
/// Returns `Err` if the pool has transitioned to the `TERMINATED` state.
/// When terminated, pushing new entries is no longer permitted.
pub(crate) fn push(&self, entries: &[worker::Entry], idx: usize) -> Result<(), ()> {
let mut state: State = self.state.load(Acquire).into();
debug_assert!(worker::State::from(entries[idx].state.load(Relaxed)).is_pushed());
loop {
let mut next = state;
let head = state.head();
if head == TERMINATED {
// The pool is terminated, cannot push the sleeper.
return Err(());
}
entries[idx].set_next_sleeper(head);
next.set_head(idx);
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if state == actual {
return Ok(());
}
state = actual;
}
}
/// Pop a worker off the stack.
///
/// If `terminate` is set and the stack is empty when this function is
/// called, the state of the stack is transitioned to "terminated". At this
/// point, no further workers can be pushed onto the stack.
///
/// # Return
///
/// Returns the index of the popped worker and the worker's observed state.
///
/// `None` if the stack is empty.
pub(crate) fn pop(
&self,
entries: &[worker::Entry],
max_lifecycle: worker::Lifecycle,
terminate: bool,
) -> Option<(usize, worker::State)> {
// Figure out the empty value
let terminal = if terminate { TERMINATED } else { EMPTY };
// If terminating, the max lifecycle *must* be `Signaled`, which is the
// highest lifecycle. By passing the greatest possible lifecycle value,
// no entries are skipped by this function.
//
// TODO: It would be better to terminate in a separate function that
// atomically takes all values and transitions to a terminated state.
debug_assert!(!terminate || max_lifecycle == worker::Lifecycle::Signaled);
let mut state: State = self.state.load(Acquire).into();
loop {
let head = state.head();
if head == EMPTY {
let mut next = state;
next.set_head(terminal);
if next == state {
debug_assert!(terminal == EMPTY);
return None;
}
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual != state {
state = actual;
continue;
}
return None;
} else if head == TERMINATED {
return None;
}
debug_assert!(head < MAX_WORKERS);
let mut next = state;
let next_head = entries[head].next_sleeper();
// TERMINATED can never be set as the "next pointer" on a worker.
debug_assert!(next_head != TERMINATED);
if next_head == EMPTY {
next.set_head(terminal);
} else {
next.set_head(next_head);
}
let actual = self
.state
.compare_and_swap(state.into(), next.into(), AcqRel)
.into();
if actual == state {
// Release ordering is needed to ensure that unsetting the
// `pushed` flag happens after popping the sleeper from the
// stack.
//
// Acquire ordering is required to acquire any memory associated
// with transitioning the worker's lifecycle.
let state = entries[head].fetch_unset_pushed(AcqRel);
if state.lifecycle() >= max_lifecycle {
// If the worker has already been notified, then it is
// warming up to do more work. In this case, try to pop
// another thread that might be in a relaxed state.
continue;
}
return Some((head, state));
}
state = actual;
}
}
}
// ===== impl State =====
impl State {
#[inline]
fn new() -> State {
State(EMPTY)
}
#[inline]
fn head(self) -> usize {
self.0 & STACK_MASK
}
#[inline]
fn set_head(&mut self, val: usize) {
// The ABA guard protects against the ABA problem w/ Treiber stacks
let aba_guard = ((self.0 >> ABA_GUARD_SHIFT) + 1) & ABA_GUARD_MASK;
self.0 = (aba_guard << ABA_GUARD_SHIFT) | val;
}
}
impl From<usize> for State {
fn from(src: usize) -> Self {
State(src)
}
}
impl From<State> for usize {
fn from(src: State) -> Self {
src.0
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
let head = self.head();
let mut fmt = fmt.debug_struct("stack::State");
if head < MAX_WORKERS {
fmt.field("head", &head);
} else if head == EMPTY {
fmt.field("head", &"EMPTY");
} else if head == TERMINATED {
fmt.field("head", &"TERMINATED");
}
fmt.finish()
}
}
@@ -1,153 +0,0 @@
use std::fmt;
/// Tracks worker state
#[derive(Clone, Copy, Eq, PartialEq)]
pub(crate) struct State(usize);
/// Set when the worker is pushed onto the scheduler's stack of sleeping
/// threads.
pub(crate) const PUSHED_MASK: usize = 0b001;
/// Manages the worker lifecycle part of the state
const LIFECYCLE_MASK: usize = 0b1110;
const LIFECYCLE_SHIFT: usize = 1;
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd, Clone, Copy)]
#[repr(usize)]
pub(crate) enum Lifecycle {
/// The worker does not currently have an associated thread.
Shutdown = 0 << LIFECYCLE_SHIFT,
/// The worker is doing work
Running = 1 << LIFECYCLE_SHIFT,
/// The worker is currently asleep in the condvar
Sleeping = 2 << LIFECYCLE_SHIFT,
/// The worker has been notified it should process more work.
Notified = 3 << LIFECYCLE_SHIFT,
/// A stronger form of notification. In this case, the worker is expected to
/// wakeup and try to acquire more work... if it enters this state while
/// already busy with other work, it is expected to signal another worker.
Signaled = 4 << LIFECYCLE_SHIFT,
}
impl State {
/// Returns true if the worker entry is pushed in the sleeper stack
pub(crate) fn is_pushed(self) -> bool {
self.0 & PUSHED_MASK == PUSHED_MASK
}
pub(crate) fn set_pushed(&mut self) {
self.0 |= PUSHED_MASK
}
pub(crate) fn is_notified(self) -> bool {
use self::Lifecycle::*;
match self.lifecycle() {
Notified | Signaled => true,
_ => false,
}
}
pub(crate) fn lifecycle(self) -> Lifecycle {
Lifecycle::from(self.0 & LIFECYCLE_MASK)
}
pub(crate) fn set_lifecycle(&mut self, val: Lifecycle) {
self.0 = (self.0 & !LIFECYCLE_MASK) | (val as usize)
}
pub(crate) fn is_signaled(self) -> bool {
self.lifecycle() == Lifecycle::Signaled
}
pub(crate) fn notify(&mut self) {
use self::Lifecycle::Signaled;
if self.lifecycle() != Signaled {
self.set_lifecycle(Signaled)
}
}
}
impl Default for State {
fn default() -> State {
// All workers will start pushed in the sleeping stack
State(PUSHED_MASK)
}
}
impl From<usize> for State {
fn from(src: usize) -> Self {
State(src)
}
}
impl From<State> for usize {
fn from(src: State) -> Self {
src.0
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("worker::State")
.field("lifecycle", &self.lifecycle())
.field("is_pushed", &self.is_pushed())
.finish()
}
}
// ===== impl Lifecycle =====
impl From<usize> for Lifecycle {
fn from(src: usize) -> Lifecycle {
use self::Lifecycle::*;
debug_assert!(
src == Shutdown as usize
|| src == Running as usize
|| src == Sleeping as usize
|| src == Notified as usize
|| src == Signaled as usize
);
unsafe { ::std::mem::transmute(src) }
}
}
impl From<Lifecycle> for usize {
fn from(src: Lifecycle) -> usize {
let v = src as usize;
debug_assert!(v & LIFECYCLE_MASK == v);
v
}
}
#[cfg(test)]
mod test {
use super::Lifecycle::*;
use super::*;
#[test]
fn lifecycle_encode() {
let lifecycles = &[Shutdown, Running, Sleeping, Notified, Signaled];
for &lifecycle in lifecycles {
let mut v: usize = lifecycle.into();
v &= LIFECYCLE_MASK;
assert_eq!(lifecycle, Lifecycle::from(v));
}
}
#[test]
fn lifecycle_ord() {
assert!(Running >= Shutdown);
assert!(Signaled >= Notified);
assert!(Signaled >= Sleeping);
}
}
+2
View File
@@ -1,3 +1,5 @@
#![allow(unused_macros)]
//! This module provides a small facade that wraps the `tracing` APIs we use, so
//! that when the `tracing` dependency is disabled, `tracing`'s macros expand to
//! no-ops.
+5
View File
@@ -0,0 +1,5 @@
mod pad;
mod rand;
pub(crate) use self::pad::CachePadded;
pub(crate) use self::rand::FastRand;
+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)
}
}
+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)
}
}
+4
View File
@@ -1,4 +1,5 @@
#![warn(rust_2018_idioms)]
#![cfg(not(miri))]
use tokio::sync::oneshot;
use tokio_executor::current_thread::{self, block_on_all, CurrentThread, TaskExecutor};
@@ -141,6 +142,7 @@ mod from_block_on_future {
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;
@@ -241,6 +243,7 @@ mod run_in_future {
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());
@@ -259,6 +262,7 @@ fn tick_on_infini_future() {
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
@@ -1,11 +1,9 @@
#![warn(rust_2018_idioms)]
use tokio_executor::park::{Park, Unpark};
use tokio_executor::threadpool;
use tokio_executor::threadpool::park::{DefaultPark, DefaultUnpark};
use tokio_executor::threadpool::*;
use tokio_test::assert_pending;
use tokio_executor::thread_pool::*;
use futures_util::future::poll_fn;
use std::cell::Cell;
use std::future::Future;
use std::pin::Pin;
@@ -18,74 +16,7 @@ use std::time::Duration;
thread_local!(static FOO: Cell<u32> = Cell::new(0));
#[test]
fn natural_shutdown_simple_futures() {
for _ in 0..1_000 {
let num_inc = Arc::new(AtomicUsize::new(0));
let num_dec = Arc::new(AtomicUsize::new(0));
FOO.with(|f| {
f.set(1);
let pool = {
let num_inc = num_inc.clone();
let num_dec = num_dec.clone();
Builder::new()
.around_worker(move |w| {
num_inc.fetch_add(1, Relaxed);
w.run();
num_dec.fetch_add(1, Relaxed);
})
.build()
};
let tx = pool.sender().clone();
let a = {
let (t, rx) = mpsc::channel();
tx.spawn(async move {
// Makes sure this runs on a worker thread
FOO.with(|f| assert_eq!(f.get(), 0));
t.send("one").unwrap();
})
.unwrap();
rx
};
let b = {
let (t, rx) = mpsc::channel();
tx.spawn(async move {
// Makes sure this runs on a worker thread
FOO.with(|f| assert_eq!(f.get(), 0));
t.send("two").unwrap();
})
.unwrap();
rx
};
drop(tx);
assert_eq!("one", a.recv().unwrap());
assert_eq!("two", b.recv().unwrap());
// Wait for the pool to shutdown
pool.shutdown().wait();
// Assert that at least one thread started
let num_inc = num_inc.load(Relaxed);
assert!(num_inc > 0);
// Assert that all threads shutdown
let num_dec = num_dec.load(Relaxed);
assert_eq!(num_inc, num_dec);
});
}
}
#[test]
fn force_shutdown_drops_futures() {
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));
@@ -110,19 +41,20 @@ fn force_shutdown_drops_futures() {
let a = num_inc.clone();
let b = num_dec.clone();
let pool = Builder::new()
.around_worker(move |w| {
let mut pool = Builder::new()
.around_worker(move |_, work| {
a.fetch_add(1, Relaxed);
w.run();
work();
b.fetch_add(1, Relaxed);
})
.build();
let tx = pool.sender().clone();
tx.spawn(Never(num_drop.clone())).unwrap();
// let tx = pool.sender().clone();
pool.spawn(Never(num_drop.clone()));
// Wait for the pool to shutdown
pool.shutdown_now().wait();
pool.shutdown_now();
// Assert that only a single thread was spawned.
let a = num_inc.load(Relaxed);
@@ -140,6 +72,8 @@ fn force_shutdown_drops_futures() {
#[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));
@@ -165,24 +99,22 @@ fn drop_threadpool_drops_futures() {
let b = num_dec.clone();
let pool = Builder::new()
.max_blocking(2)
.pool_size(20)
.around_worker(move |w| {
.num_threads(NUM_THREADS)
.around_worker(move |_, work| {
a.fetch_add(1, Relaxed);
w.run();
work();
b.fetch_add(1, Relaxed);
})
.build();
let tx = pool.sender().clone();
tx.spawn(Never(num_drop.clone())).unwrap();
pool.spawn(Never(num_drop.clone()));
// Wait for the pool to shutdown
drop(pool);
// Assert that only a single thread was spawned.
// Assert that all the threads spawned
let a = num_inc.load(Relaxed);
assert!(a >= 1);
assert_eq!(a, NUM_THREADS);
// Assert that all threads shutdown
let b = num_dec.load(Relaxed);
@@ -196,26 +128,32 @@ fn drop_threadpool_drops_futures() {
#[test]
fn many_oneshot_futures() {
// used for notifying the main thread
const NUM: usize = 10_000;
for _ in 0..50 {
let pool = ThreadPool::new();
let tx = pool.sender().clone();
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();
tx.spawn(async move {
cnt.fetch_add(1, Relaxed);
})
.unwrap();
let tx = tx.clone();
pool.spawn(async move {
let num = cnt.fetch_add(1, Relaxed) + 1;
if num == NUM {
tx.send(()).unwrap();
}
});
}
// Wait for the pool to shutdown
pool.shutdown().wait();
rx.recv().unwrap();
let num = cnt.load(Relaxed);
assert_eq!(num, NUM);
// Wait for the pool to shutdown
pool.shutdown_now();
}
}
@@ -228,9 +166,7 @@ fn many_multishot_futures() {
const TRACKS: usize = 50;
for _ in 0..50 {
let pool = ThreadPool::new();
let pool_tx = pool.sender().clone();
let pool = new_pool();
let mut start_txs = Vec::with_capacity(TRACKS);
let mut final_rxs = Vec::with_capacity(TRACKS);
@@ -241,13 +177,11 @@ fn many_multishot_futures() {
let (mut next_tx, next_rx) = mpsc::channel(10);
// Forward all the messages
pool_tx
.spawn(async move {
while let Some(v) = chain_rx.recv().await {
next_tx.send(v).await.unwrap();
}
})
.unwrap();
pool.spawn(async move {
while let Some(v) = chain_rx.recv().await {
next_tx.send(v).await.unwrap();
}
});
chain_rx = next_rx;
}
@@ -257,21 +191,19 @@ fn many_multishot_futures() {
let mut cycle_tx = start_tx.clone();
let mut rem = CYCLES;
pool_tx
.spawn(async move {
for _ in 0..CYCLES {
let msg = chain_rx.recv().await.unwrap();
pool.spawn(async move {
for _ in 0..CYCLES {
let msg = chain_rx.recv().await.unwrap();
rem -= 1;
rem -= 1;
if rem == 0 {
final_tx.send(msg).await.unwrap();
} else {
cycle_tx.send(msg).await.unwrap();
}
if rem == 0 {
final_tx.send(msg).await.unwrap();
} else {
cycle_tx.send(msg).await.unwrap();
}
})
.unwrap();
}
});
start_txs.push(start_tx);
final_rxs.push(final_rx);
@@ -290,80 +222,36 @@ fn many_multishot_futures() {
}
});
}
// Shutdown the pool
pool.shutdown().wait();
}
}
#[test]
fn global_executor_is_configured() {
let pool = ThreadPool::new();
let tx = pool.sender().clone();
let pool = new_pool();
let (signal_tx, signal_rx) = mpsc::channel();
tx.spawn(async move {
pool.spawn(async move {
tokio_executor::spawn(async move {
signal_tx.send(()).unwrap();
});
})
.unwrap();
});
signal_rx.recv().unwrap();
pool.shutdown().wait();
}
#[test]
fn new_threadpool_is_idle() {
let pool = ThreadPool::new();
pool.shutdown_on_idle().wait();
}
#[test]
fn busy_threadpool_is_not_idle() {
use tokio_sync::oneshot;
// let pool = ThreadPool::new();
let pool = Builder::new().pool_size(4).max_blocking(2).build();
let tx = pool.sender().clone();
let (term_tx, term_rx) = oneshot::channel();
tx.spawn(async move {
term_rx.await.unwrap();
})
.unwrap();
let mut idle = pool.shutdown_on_idle();
struct IdleFut<'a>(&'a mut Shutdown);
impl Future for IdleFut<'_> {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
assert_pending!(Pin::new(&mut self.as_mut().0).poll(cx));
Poll::Ready(())
}
}
let idle_fut = IdleFut(&mut idle);
tokio_executor::enter().unwrap().block_on(idle_fut);
term_tx.send(()).unwrap();
let idle_fut = IdleFut(&mut idle);
tokio_executor::enter().unwrap().block_on(idle_fut);
let mut pool = new_pool();
pool.shutdown_now();
}
#[test]
fn panic_in_task() {
let pool = ThreadPool::new();
let tx = pool.sender().clone();
let pool = new_pool();
let (tx, rx) = mpsc::channel();
struct Boom;
struct Boom(mpsc::Sender<()>);
impl Future for Boom {
type Output = ();
@@ -376,37 +264,20 @@ fn panic_in_task() {
impl Drop for Boom {
fn drop(&mut self) {
assert!(::std::thread::panicking());
self.0.send(()).unwrap();
}
}
tx.spawn(Boom).unwrap();
pool.shutdown_on_idle().wait();
}
#[test]
fn count_panics() {
let counter = Arc::new(AtomicUsize::new(0));
let counter_ = counter.clone();
let pool = threadpool::Builder::new()
.panic_handler(move |_err| {
// We caught a panic.
counter_.fetch_add(1, Relaxed);
})
.build();
// Spawn a future that will panic.
pool.spawn(async { panic!() });
pool.shutdown_on_idle().wait();
let counter = counter.load(Relaxed);
assert_eq!(counter, 1);
pool.spawn(Boom(tx));
rx.recv().unwrap();
}
#[test]
fn multi_threadpool() {
use tokio_sync::oneshot;
let pool1 = ThreadPool::new();
let pool2 = ThreadPool::new();
let pool1 = new_pool();
let pool2 = new_pool();
let (tx, rx) = oneshot::channel();
let (done_tx, done_rx) = mpsc::channel();
@@ -425,10 +296,11 @@ fn multi_threadpool() {
#[test]
fn eagerly_drops_futures() {
use std::sync::mpsc;
use std::sync::{mpsc, Mutex};
struct MyPark {
inner: DefaultPark,
rx: mpsc::Receiver<()>,
tx: Mutex<mpsc::Sender<()>>,
#[allow(dead_code)]
park_tx: mpsc::SyncSender<()>,
unpark_tx: mpsc::SyncSender<()>,
@@ -436,33 +308,35 @@ fn eagerly_drops_futures() {
impl Park for MyPark {
type Unpark = MyUnpark;
type Error = <DefaultPark as Park>::Error;
type Error = ();
fn unpark(&self) -> Self::Unpark {
MyUnpark {
inner: self.inner.unpark(),
tx: Mutex::new(self.tx.lock().unwrap().clone()),
unpark_tx: self.unpark_tx.clone(),
}
}
fn park(&mut self) -> Result<(), Self::Error> {
self.inner.park()
let _ = self.rx.recv();
Ok(())
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
self.inner.park_timeout(duration)
let _ = self.rx.recv_timeout(duration);
Ok(())
}
}
struct MyUnpark {
inner: DefaultUnpark,
tx: Mutex<mpsc::Sender<()>>,
#[allow(dead_code)]
unpark_tx: mpsc::SyncSender<()>,
}
impl Unpark for MyUnpark {
fn unpark(&self) {
self.inner.unpark()
let _ = self.tx.lock().unwrap().send(());
}
}
@@ -471,13 +345,15 @@ fn eagerly_drops_futures() {
let (park_tx, park_rx) = mpsc::sync_channel(0);
let (unpark_tx, unpark_rx) = mpsc::sync_channel(0);
let pool = threadpool::Builder::new()
.custom_park(move |_| MyPark {
inner: DefaultPark::new(),
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(),
})
.build();
}
});
struct MyTask {
task_tx: Option<mpsc::Sender<Waker>>,
@@ -523,3 +399,80 @@ fn eagerly_drops_futures() {
// 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()
}
-412
View File
@@ -1,412 +0,0 @@
#![warn(rust_2018_idioms)]
use tokio_executor::threadpool::*;
use tokio_test::*;
use futures_core::ready;
use futures_util::future::poll_fn;
use rand::*;
use std::sync::atomic::Ordering::*;
use std::sync::atomic::*;
use std::sync::*;
use std::task::{Poll, Waker};
use std::thread;
use std::time::Duration;
#[test]
fn basic() {
let pool = Builder::new().pool_size(1).max_blocking(1).build();
let (tx1, rx1) = mpsc::channel();
let (tx2, rx2) = mpsc::channel();
pool.spawn(async move {
let res = blocking(|| {
let v = rx1.recv().unwrap();
tx2.send(v).unwrap();
});
assert_ready!(res).unwrap();
});
pool.spawn(async move {
tx1.send(()).unwrap();
});
rx2.recv().unwrap();
}
#[test]
fn other_executors_can_run_inside_blocking() {
let pool = Builder::new().pool_size(1).max_blocking(1).build();
let (tx, rx) = mpsc::channel();
pool.spawn(async move {
let res = blocking(|| {
let _e = tokio_executor::enter().expect("nested blocking enter");
tx.send(()).unwrap();
});
assert_ready!(res).unwrap();
});
rx.recv().unwrap();
}
#[test]
fn notify_task_on_capacity() {
const BLOCKING: usize = 10;
let pool = Builder::new().pool_size(1).max_blocking(1).build();
let rem = Arc::new(AtomicUsize::new(BLOCKING));
let (tx, rx) = mpsc::channel();
for _ in 0..BLOCKING {
let rem = rem.clone();
let tx = tx.clone();
pool.spawn(async move {
poll_fn(move |_| {
blocking(|| {
thread::sleep(Duration::from_millis(100));
let prev = rem.fetch_sub(1, SeqCst);
if prev == 1 {
tx.send(()).unwrap();
}
})
.map_err(|e| panic!("blocking err {:?}", e))
})
.await
.unwrap()
});
}
rx.recv().unwrap();
assert_eq!(0, rem.load(SeqCst));
}
#[test]
fn capacity_is_use_it_or_lose_it() {
use tokio_sync::oneshot;
// TODO: Run w/ bigger pool size
let pool = Builder::new().pool_size(1).max_blocking(1).build();
let (tx1, rx1) = mpsc::channel();
let (tx2, rx2) = oneshot::channel();
let (tx3, rx3) = mpsc::channel();
let (tx4, rx4) = mpsc::channel();
// First, fill the blocking capacity
pool.spawn(async move {
poll_fn(move |_| {
blocking(|| {
rx1.recv().unwrap();
})
})
.await
.unwrap()
});
pool.spawn(async move {
let task: Waker = rx2.await.unwrap();
poll_fn(move |_| {
blocking(|| {
// Notify the other task
task.wake_by_ref();
// Block until woken
rx3.recv().unwrap();
})
})
.await
.unwrap();
});
// Spawn a future that will try to block, get notified, then not actually
// use the blocking
let mut i = 0;
let mut tx2 = Some(tx2);
pool.spawn(async move {
poll_fn(move |cx| {
match i {
0 => {
i = 1;
let res = blocking(|| unreachable!()).map_err(|_| panic!());
assert_pending!(res);
// Unblock the first blocker
tx1.send(()).unwrap();
return Poll::Pending;
}
1 => {
i = 2;
// Skip blocking, and notify the second task that it should
// start blocking
let me = cx.waker().clone();
tx2.take().unwrap().send(me).unwrap();
return Poll::Pending;
}
2 => {
let res = blocking(|| unreachable!()).map_err(|_| panic!());
assert_pending!(res);
// Unblock the first blocker
tx3.send(()).unwrap();
tx4.send(()).unwrap();
Poll::Ready(())
}
_ => unreachable!(),
}
})
.await
});
rx4.recv().unwrap();
}
#[test]
fn blocking_thread_does_not_take_over_shutdown_worker_thread() {
let pool = Builder::new().pool_size(2).max_blocking(1).build();
let (enter_tx, enter_rx) = mpsc::channel();
let (exit_tx, exit_rx) = mpsc::channel();
let (try_tx, try_rx) = mpsc::channel();
let exited = Arc::new(AtomicBool::new(false));
{
let exited = exited.clone();
pool.spawn(async move {
poll_fn(move |_| {
blocking(|| {
enter_tx.send(()).unwrap();
exit_rx.recv().unwrap();
exited.store(true, SeqCst);
})
})
.await
.unwrap()
});
}
// Wait for the task to block
let _ = enter_rx.recv().unwrap();
// Spawn another task that attempts to block
pool.spawn(async move {
poll_fn(move |_| {
let res = blocking(|| {});
assert_eq!(res.is_ready(), exited.load(SeqCst));
try_tx.send(res.is_ready()).unwrap();
res.map(|_| ())
})
.await
});
// Wait for the second task to try to block (and not be ready).
let res = try_rx.recv().unwrap();
assert!(!res);
// Unblock the first task
exit_tx.send(()).unwrap();
// Wait for the second task to successfully block.
let res = try_rx.recv().unwrap();
assert!(res);
drop(pool);
}
#[test]
fn blocking_one_time_gets_capacity_for_multiple_blocks() {
const ITER: usize = 1;
const BLOCKING: usize = 2;
for _ in 0..ITER {
let pool = Builder::new().pool_size(4).max_blocking(1).build();
let rem = Arc::new(AtomicUsize::new(BLOCKING));
let (tx, rx) = mpsc::channel();
for _ in 0..BLOCKING {
let rem = rem.clone();
let tx = tx.clone();
pool.spawn(async move {
poll_fn(move |_| {
// First block
let res = blocking(|| {
thread::sleep(Duration::from_millis(100));
});
ready!(res).unwrap();
let res = blocking(|| {
thread::sleep(Duration::from_millis(100));
let prev = rem.fetch_sub(1, SeqCst);
if prev == 1 {
tx.send(()).unwrap();
}
});
assert!(res.is_ready());
Poll::Ready(())
})
.await
});
}
rx.recv().unwrap();
assert_eq!(0, rem.load(SeqCst));
}
}
#[test]
fn shutdown() {
const ITER: usize = 1_000;
const BLOCKING: usize = 10;
for _ in 0..ITER {
let num_inc = Arc::new(AtomicUsize::new(0));
let num_dec = Arc::new(AtomicUsize::new(0));
let (tx, rx) = mpsc::channel();
let pool = {
let num_inc = num_inc.clone();
let num_dec = num_dec.clone();
Builder::new()
.pool_size(1)
.max_blocking(BLOCKING)
.after_start(move || {
num_inc.fetch_add(1, SeqCst);
})
.before_stop(move || {
num_dec.fetch_add(1, SeqCst);
})
.build()
};
let barrier = Arc::new(Barrier::new(BLOCKING));
for _ in 0..BLOCKING {
let barrier = barrier.clone();
let tx = tx.clone();
pool.spawn(async move {
let res = blocking(|| {
barrier.wait();
Ok::<_, ()>(())
});
tx.send(()).unwrap();
assert!(res.is_ready());
});
}
for _ in 0..BLOCKING {
rx.recv().unwrap();
}
// Shutdown
drop(pool);
assert_eq!(11, num_inc.load(SeqCst));
assert_eq!(11, num_dec.load(SeqCst));
}
}
#[derive(Debug, Copy, Clone)]
enum Sleep {
Skip,
Yield,
Rand,
Fixed(Duration),
}
#[test]
fn hammer() {
use self::Sleep::*;
const ITER: usize = 5;
let combos = [
(2, 4, 1_000, Skip),
(2, 4, 1_000, Yield),
(2, 4, 100, Rand),
(2, 4, 100, Fixed(Duration::from_millis(3))),
(2, 4, 100, Fixed(Duration::from_millis(12))),
];
for &(size, max_blocking, n, sleep) in &combos {
for _ in 0..ITER {
let pool = Builder::new()
.pool_size(size)
.max_blocking(max_blocking)
.build();
let cnt_task = Arc::new(AtomicUsize::new(0));
let cnt_block = Arc::new(AtomicUsize::new(0));
for _ in 0..n {
let cnt_task = cnt_task.clone();
let cnt_block = cnt_block.clone();
pool.spawn(async move {
cnt_task.fetch_add(1, SeqCst);
poll_fn(move |_| {
blocking(|| {
match sleep {
Skip => {}
Yield => {
thread::yield_now();
}
Rand => {
let ms = thread_rng().gen_range(3, 12);
thread::sleep(Duration::from_millis(ms));
}
Fixed(dur) => {
thread::sleep(dur);
}
}
cnt_block.fetch_add(1, SeqCst);
})
.map_err(|_| panic!())
})
.await
.unwrap()
});
}
// Wait for the work to complete
pool.shutdown_on_idle().wait();
assert_eq!(n, cnt_task.load(SeqCst));
assert_eq!(n, cnt_block.load(SeqCst));
}
}
}
-126
View File
@@ -1,126 +0,0 @@
#![warn(rust_2018_idioms)]
use tokio_executor::threadpool::*;
use tokio_sync::{mpsc, oneshot};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::*;
use std::sync::Arc;
use std::task::{Context, Poll};
#[test]
fn hammer() {
const N: usize = 1000;
const ITER: usize = 20;
struct Counted<T> {
cnt: Arc<AtomicUsize>,
inner: T,
}
impl<T: Future> Future for Counted<T> {
type Output = T::Output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T::Output> {
unsafe {
let inner = &mut self.get_unchecked_mut().inner;
Pin::new_unchecked(inner).poll(cx)
}
}
}
impl<T> Drop for Counted<T> {
fn drop(&mut self) {
self.cnt.fetch_add(1, Relaxed);
}
}
for _ in 0..ITER {
let pool = Builder::new()
// .pool_size(30)
.build();
let cnt = Arc::new(AtomicUsize::new(0));
let (mut listen_tx, mut listen_rx) =
mpsc::unbounded_channel::<oneshot::Sender<oneshot::Sender<()>>>();
pool.spawn({
let c1 = cnt.clone();
let c2 = cnt.clone();
let pool = pool.sender().clone();
let task = async move {
while let Some(tx) = listen_rx.recv().await {
let task = async {
let (tx2, rx2) = oneshot::channel();
tx.send(tx2).unwrap();
rx2.await.unwrap()
};
pool.spawn(Counted {
inner: task,
cnt: c1.clone(),
})
.unwrap();
}
};
/*
let task = listen_rx
.map_err(|e| panic!("accept error = {:?}", e))
.for_each(move |tx| {
let task = future::lazy(|| {
let (tx2, rx2) = oneshot::channel();
tx.send(tx2).unwrap();
rx2
})
.map_err(|e| panic!("e={:?}", e))
.and_then(|_| Ok(()));
pool.spawn(Counted {
inner: task,
cnt: c1.clone(),
})
.unwrap();
Ok(())
});
*/
Counted {
inner: task,
cnt: c2,
}
});
for _ in 0..N {
let cnt = cnt.clone();
let (tx, rx) = oneshot::channel();
listen_tx.try_send(tx).unwrap();
pool.spawn(async {
let task = async {
let tx = rx.await.unwrap();
tx.send(()).unwrap();
};
/*
let task = rx.map_err(|e| panic!("rx err={:?}", e)).and_then(|tx| {
tx.send(()).unwrap();
Ok(())
});
*/
Counted { inner: task, cnt }.await
});
}
drop(listen_tx);
pool.shutdown_on_idle().wait();
assert_eq!(N * 2 + 1, cnt.load(Relaxed));
}
}
-36
View File
@@ -1,36 +0,0 @@
//! Echo everything received on STDIN to STDOUT and STDERR.
use futures_util::{FutureExt, SinkExt, StreamExt, TryFutureExt};
use tokio::codec::{FramedRead, FramedWrite, LinesCodec, LinesCodecError};
use tokio::future::ready;
use tokio_executor::threadpool::Builder;
use tokio_fs::{stderr, stdin, stdout};
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let pool = Builder::new().pool_size(1).build();
pool.spawn(
async {
let mut input = FramedRead::new(stdin(), LinesCodec::new());
let mut output = FramedWrite::new(stdout(), LinesCodec::new());
let mut error = FramedWrite::new(stderr(), LinesCodec::new());
while let Some(line) = input.next().await {
let line = line?;
// https://github.com/rust-lang/rust/pull/64856
let s = format!("OUT: {}", line);
output.send(s).await?;
let s = format!("ERR: {}", line);
error.send(s).await?;
}
Ok(())
}
.map_err(|e: LinesCodecError| panic!(e))
.then(|_| ready(())),
);
pool.shutdown_on_idle().await;
Ok(())
}
+1 -1
View File
@@ -53,7 +53,7 @@ rt-full = [
"sync",
"timer",
"tokio-executor/current-thread",
"tokio-executor/threadpool",
"tokio-executor/thread-pool",
"tracing-core",
]
signal = ["tokio-net/signal"]
+2 -5
View File
@@ -70,9 +70,6 @@
//! }
//! ```
//!
//! In this function, the `run` function blocks until the runtime becomes idle.
//! See [`shutdown_on_idle`][idle] for more shutdown details.
//!
//! From within the context of the runtime, additional tasks are spawned using
//! the [`tokio::spawn`] function. Futures spawned using this function will be
//! executed on the same thread pool used by the [`Runtime`].
@@ -129,7 +126,6 @@
//! [`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
//! [idle]: struct.Runtime.html#method.shutdown_on_idle
//! [`tokio::spawn`]: ../executor/fn.spawn.html
//! [`Timer`]: https://docs.rs/tokio-timer/0.2/tokio_timer/timer/struct.Timer.html
//! [`tokio::main`]: ../../tokio_macros/attr.main.html
@@ -141,8 +137,9 @@ mod threadpool;
#[cfg(feature = "rt-full")]
pub use self::threadpool::{
Builder,
JoinHandle,
Runtime,
TaskExecutor,
Spawner,
};
// Internal export, don't use.
@@ -1,61 +0,0 @@
//! Temporary reactor + timer that runs on a background thread. This it to make
//! `block_on` work.
use tokio_executor::current_thread::CurrentThread;
use tokio_net::driver::{self, Reactor};
use tokio_sync::oneshot;
use tokio_timer::clock::Clock;
use tokio_timer::timer::{self, Timer};
use std::{io, thread};
#[derive(Debug)]
pub(crate) struct Background {
reactor_handle: driver::Handle,
timer_handle: timer::Handle,
shutdown_tx: Option<oneshot::Sender<()>>,
thread: Option<thread::JoinHandle<()>>,
}
pub(crate) fn spawn(clock: &Clock) -> io::Result<Background> {
let clock = clock.clone();
let reactor = Reactor::new()?;
let reactor_handle = reactor.handle();
let timer = Timer::new_with_now(reactor, clock);
let timer_handle = timer.handle();
let (shutdown_tx, shutdown_rx) = oneshot::channel();
let shutdown_tx = Some(shutdown_tx);
let thread = thread::spawn(move || {
let mut rt = CurrentThread::new_with_park(timer);
let _ = rt.block_on(shutdown_rx);
});
let thread = Some(thread);
Ok(Background {
reactor_handle,
timer_handle,
shutdown_tx,
thread,
})
}
impl Background {
pub(super) fn reactor(&self) -> &driver::Handle {
&self.reactor_handle
}
pub(super) fn timer(&self) -> &timer::Handle {
&self.timer_handle
}
}
impl Drop for Background {
fn drop(&mut self) {
let _ = self.shutdown_tx.take().unwrap().send(());
let _ = self.thread.take().unwrap().join();
}
}
+62 -133
View File
@@ -1,16 +1,13 @@
use super::{background, Inner, Runtime};
use super::{Inner, Runtime};
use tokio_executor::threadpool;
use tokio_executor::thread_pool;
use tokio_net::driver::{self, Reactor};
use tokio_timer::clock::{self, Clock};
use tokio_timer::timer::{self, Timer};
use num_cpus;
use tracing_core as trace;
use std::io;
use std::sync::Mutex;
use std::time::Duration;
use std::any::Any;
use std::{fmt, io};
use std::sync::{Arc, Mutex};
/// Builds Tokio Runtime with custom configuration values.
///
@@ -28,18 +25,14 @@ use std::any::Any;
/// # Examples
///
/// ```
/// use std::time::Duration;
///
/// use tokio::runtime::Builder;
/// use tokio_timer::clock::Clock;
///
/// fn main() {
/// // build Runtime
/// let runtime = Builder::new()
/// .blocking_threads(4)
/// .clock(Clock::system())
/// .core_threads(4)
/// .keep_alive(Some(Duration::from_secs(60)))
/// .num_threads(4)
/// .name_prefix("my-custom-name-")
/// .stack_size(3 * 1024 * 1024)
/// .build()
@@ -48,13 +41,18 @@ use std::any::Any;
/// // use runtime ...
/// }
/// ```
#[derive(Debug)]
pub struct Builder {
/// Thread pool specific builder
threadpool_builder: threadpool::Builder,
thread_pool_builder: thread_pool::Builder,
/// The number of worker threads
core_threads: usize,
num_threads: usize,
/// To run after each worker thread starts
after_start: Option<Arc<dyn Fn() + Send + Sync>>,
/// To run before each worker thread stops
before_stop: Option<Arc<dyn Fn() + Send + Sync>>,
/// The clock to use
clock: Clock,
@@ -66,15 +64,18 @@ impl Builder {
///
/// Configuration methods can be chained on the return value.
pub fn new() -> Builder {
let core_threads = num_cpus::get().max(1);
let num_threads = num_cpus::get().max(1);
let mut threadpool_builder = threadpool::Builder::new();
threadpool_builder.name_prefix("tokio-runtime-worker-");
threadpool_builder.pool_size(core_threads);
let mut thread_pool_builder = thread_pool::Builder::new();
thread_pool_builder
.name_prefix("tokio-runtime-worker-")
.num_threads(num_threads);
Builder {
threadpool_builder,
core_threads,
thread_pool_builder,
num_threads,
after_start: None,
before_stop: None,
clock: Clock::new(),
}
}
@@ -85,35 +86,6 @@ impl Builder {
self
}
/// Sets a callback to handle panics in futures.
///
/// The callback is triggered when a panic during a future bubbles up to
/// Tokio. By default Tokio catches these panics, and they will be ignored.
/// The parameter passed to this callback is the same error value returned
/// from `std::panic::catch_unwind()`. To abort the process on panics, use
/// `std::panic::resume_unwind()` in this callback as shown below.
///
/// # Examples
///
/// ```
/// # use tokio::runtime;
///
/// # pub fn main() {
/// let rt = runtime::Builder::new()
/// .panic_handler(|err| std::panic::resume_unwind(err))
/// .build()
/// .unwrap();
/// # }
/// ```
pub fn panic_handler<F>(&mut self, f: F) -> &mut Self
where
F: Fn(Box<dyn Any + Send>) + Send + Sync + 'static,
{
self.threadpool_builder.panic_handler(f);
self
}
/// Set the maximum number of worker threads for the `Runtime`'s thread pool.
///
/// This must be a number between 1 and 32,768 though it is advised to keep
@@ -128,71 +100,14 @@ impl Builder {
///
/// # pub fn main() {
/// let rt = runtime::Builder::new()
/// .core_threads(4)
/// .num_threads(4)
/// .build()
/// .unwrap();
/// # }
/// ```
pub fn core_threads(&mut self, val: usize) -> &mut Self {
self.core_threads = val;
self.threadpool_builder.pool_size(val);
self
}
/// Set the maximum number of concurrent blocking sections in the `Runtime`'s
/// thread pool.
///
/// When the maximum concurrent `blocking` calls is reached, any further
/// calls to `blocking` will return `NotReady` and the task is notified once
/// previously in-flight calls to `blocking` return.
///
/// This must be a number between 1 and 32,768 though it is advised to keep
/// this value on the smaller side.
///
/// The default value is 100.
///
/// # Examples
///
/// ```
/// # use tokio::runtime;
///
/// # pub fn main() {
/// let rt = runtime::Builder::new()
/// .blocking_threads(200)
/// .build();
/// # }
/// ```
pub fn blocking_threads(&mut self, val: usize) -> &mut Self {
self.threadpool_builder.max_blocking(val);
self
}
/// Set the worker thread keep alive duration for threads in the `Runtime`'s
/// thread pool.
///
/// If set, a worker thread will wait for up to the specified duration for
/// work, at which point the thread will shutdown. When work becomes
/// available, a new thread will eventually be spawned to replace the one
/// that shut down.
///
/// When the value is `None`, the thread will wait for work forever.
///
/// The default value is `None`.
///
/// # Examples
///
/// ```
/// # use tokio::runtime;
/// use std::time::Duration;
///
/// # pub fn main() {
/// let rt = runtime::Builder::new()
/// .keep_alive(Some(Duration::from_secs(30)))
/// .build();
/// # }
/// ```
pub fn keep_alive(&mut self, val: Option<Duration>) -> &mut Self {
self.threadpool_builder.keep_alive(val);
pub fn num_threads(&mut self, val: usize) -> &mut Self {
self.num_threads = val;
self.thread_pool_builder.num_threads(val);
self
}
@@ -216,7 +131,7 @@ impl Builder {
/// # }
/// ```
pub fn name_prefix<S: Into<String>>(&mut self, val: S) -> &mut Self {
self.threadpool_builder.name_prefix(val);
self.thread_pool_builder.name_prefix(val);
self
}
@@ -240,7 +155,7 @@ impl Builder {
/// # }
/// ```
pub fn stack_size(&mut self, val: usize) -> &mut Self {
self.threadpool_builder.stack_size(val);
self.thread_pool_builder.stack_size(val);
self
}
@@ -265,7 +180,7 @@ impl Builder {
pub fn after_start<F>(&mut self, f: F) -> &mut Self
where F: Fn() + Send + Sync + 'static
{
self.threadpool_builder.after_start(f);
self.after_start = Some(Arc::new(f));
self
}
@@ -289,7 +204,7 @@ impl Builder {
pub fn before_stop<F>(&mut self, f: F) -> &mut Self
where F: Fn() + Send + Sync + 'static
{
self.threadpool_builder.before_stop(f);
self.before_stop = Some(Arc::new(f));
self
}
@@ -308,14 +223,11 @@ impl Builder {
/// # }
/// ```
pub fn build(&mut self) -> io::Result<Runtime> {
// TODO(stjepang): Once we remove the `threadpool_builder` method, remove this line too.
self.threadpool_builder.pool_size(self.core_threads);
let mut reactor_handles = Vec::new();
let mut timer_handles = Vec::new();
let mut timers = Vec::new();
for _ in 0..self.core_threads {
for _ in 0..self.num_threads {
// Create a new reactor.
let reactor = Reactor::new()?;
reactor_handles.push(reactor.handle());
@@ -336,36 +248,44 @@ impl Builder {
let dispatch = trace::dispatcher::get_default(trace::Dispatch::clone);
let trace = dispatch.clone();
let background = background::spawn(&clock)?;
let around_reactor_handles = reactor_handles.clone();
let around_timer_handles = timer_handles.clone();
let after_start = self.after_start.clone();
let before_stop = self.before_stop.clone();
let pool = self
.threadpool_builder
.around_worker(move |w| {
let index = w.id().to_usize();
let _reactor = driver::set_default(&reactor_handles[index]);
.thread_pool_builder
.around_worker(move |index, next| {
let _reactor = driver::set_default(&around_reactor_handles[index]);
clock::with_default(&clock, || {
let _timer = timer::set_default(&timer_handles[index]);
let _timer = timer::set_default(&around_timer_handles[index]);
trace::dispatcher::with_default(&dispatch, || {
w.run();
if let Some(after_start) = after_start.as_ref() {
after_start();
}
next();
if let Some(before_stop) = before_stop.as_ref() {
before_stop();
}
})
})
})
.custom_park(move |worker_id| {
let index = worker_id.to_usize();
.build_with_park(move |index| {
timers[index]
.lock()
.unwrap()
.take()
.unwrap()
})
.build();
});
Ok(Runtime {
inner: Some(Inner {
pool,
background,
reactor_handles,
timer_handles,
trace,
}),
})
@@ -377,3 +297,12 @@ impl Default for Builder {
Self::new()
}
}
impl fmt::Debug for Builder {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("Builder")
.field("thread_pool_builder", &self.thread_pool_builder)
.field("after_start", &self.after_start.as_ref().map(|_| "..."))
.finish()
}
}
+28 -78
View File
@@ -1,15 +1,15 @@
mod background;
mod builder;
mod task_executor;
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
pub use self::builder::Builder;
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
pub use self::task_executor::TaskExecutor;
use background::Background;
use tokio_executor::enter;
use tokio_executor::threadpool::ThreadPool;
mod spawner;
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
pub use self::spawner::Spawner;
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
pub use tokio_executor::thread_pool::JoinHandle;
use tokio_executor::thread_pool::ThreadPool;
use tokio_net::driver;
use tokio_timer::timer;
@@ -41,18 +41,14 @@ struct Inner {
/// Task execution pool.
pool: ThreadPool,
/// Reactor handles
reactor_handles: Vec<tokio_net::driver::Handle>,
/// Timer handles
timer_handles: Vec<timer::Handle>,
/// Tracing dispatcher
trace: trace::Dispatch,
/// Maintains a reactor and timer that are always running on a background
/// thread. This is to support `runtime.block_on` w/o requiring the future
/// to be `Send`.
///
/// A dedicated background thread is required as the threadpool threads
/// might not be running. However, this is a temporary work around.
///
/// TODO: Delete this
background: Background,
}
// ===== impl Runtime =====
@@ -80,9 +76,6 @@ impl Runtime {
/// .unwrap();
///
/// // Use the runtime...
///
/// // Shutdown the runtime
/// rt.shutdown_now();
/// ```
///
/// [mod]: index.html
@@ -90,27 +83,6 @@ impl Runtime {
Builder::new().build()
}
/// Return a handle to the runtime's executor.
///
/// The returned handle can be used to spawn tasks that run on this runtime.
///
/// # Examples
///
/// ```
/// use tokio::runtime::Runtime;
///
/// let rt = Runtime::new()
/// .unwrap();
///
/// let executor_handle = rt.executor();
///
/// // use `executor_handle`
/// ```
pub fn executor(&self) -> TaskExecutor {
let inner = self.inner().pool.sender().clone();
TaskExecutor { inner }
}
/// Spawn a future onto the Tokio runtime.
///
/// This spawns the given future onto the runtime's executor, usually a
@@ -134,8 +106,6 @@ impl Runtime {
/// rt.spawn(async {
/// println!("now running on a worker thread");
/// });
///
/// rt.shutdown_on_idle();
/// }
/// ```
///
@@ -166,32 +136,19 @@ impl Runtime {
where
F: Future,
{
let mut entered = enter().expect("nested block_on");
let bg = &self.inner().background;
let trace = &self.inner().trace;
tokio_executor::with_default(&mut self.inner().pool.sender(), || {
let _reactor = driver::set_default(bg.reactor());
let _timer = timer::set_default(bg.timer());
trace::dispatcher::with_default(trace, || {
entered.block_on(future)
})
let _reactor = driver::set_default(&self.inner().reactor_handles[0]);
let _timer = timer::set_default(&self.inner().timer_handles[0]);
trace::dispatcher::with_default(trace, || {
self.inner().pool.block_on(future)
})
}
/// Signals the runtime to shutdown once it becomes idle.
/// Return a handle to the runtime's spawner.
///
/// Blocks the current thread until the shutdown operation has completed.
/// This function can be used to perform a graceful shutdown of the runtime.
///
/// The runtime enters an idle state once **all** of the following occur.
///
/// * The thread pool has no tasks to execute, i.e., all tasks that were
/// spawned have completed.
/// * The reactor is not managing any I/O resources.
///
/// See [module level][mod] documentation for more details.
/// The returned handle can be used to spawn tasks that run on this runtime.
///
/// # Examples
///
@@ -201,18 +158,13 @@ impl Runtime {
/// let rt = Runtime::new()
/// .unwrap();
///
/// // Use the runtime...
/// let spawner = rt.spawner();
///
/// // Shutdown the runtime
/// rt.shutdown_on_idle();
/// spawner.spawn(async { println!("hello"); });
/// ```
///
/// [mod]: index.html
pub fn shutdown_on_idle(mut self) {
let mut e = tokio_executor::enter().unwrap();
let inner = self.inner.take().unwrap();
e.block_on(inner.pool.shutdown_on_idle());
pub fn spawner(&self) -> Spawner {
let inner = self.inner().pool.spawner().clone();
Spawner::new(inner)
}
/// Signals the runtime to shutdown immediately.
@@ -248,11 +200,9 @@ impl Runtime {
/// ```
///
/// [mod]: index.html
#[allow(warnings)]
pub fn shutdown_now(mut self) {
let mut e = tokio_executor::enter().unwrap();
let inner = self.inner.take().unwrap();
e.block_on(inner.pool.shutdown_now());
self.inner.unwrap().pool.shutdown_now();
}
fn inner(&self) -> &Inner {
@@ -1,21 +1,25 @@
use tokio_executor::SpawnError;
use tokio_executor::threadpool::Sender;
use crate::runtime::JoinHandle;
use tokio_executor::thread_pool;
use std::future::Future;
use std::pin::Pin;
/// Executes futures on the runtime
/// Spawns futures on the runtime
///
/// All futures spawned using this executor will be submitted to the associated
/// Runtime's executor. This executor is usually a thread pool.
///
/// For more details, see the [module level](index.html) documentation.
#[derive(Debug, Clone)]
pub struct TaskExecutor {
pub(super) inner: Sender,
pub struct Spawner {
inner: thread_pool::Spawner,
}
impl TaskExecutor {
impl Spawner {
pub(super) fn new(inner: thread_pool::Spawner) -> Spawner {
Spawner { inner }
}
/// Spawn a future onto the Tokio runtime.
///
/// This spawns the given future onto the runtime's executor, usually a
@@ -34,10 +38,10 @@ impl TaskExecutor {
/// # fn dox() {
/// // Create the runtime
/// let rt = Runtime::new().unwrap();
/// let executor = rt.executor();
/// let spawner = rt.spawner();
///
/// // Spawn a future onto the runtime
/// executor.spawn(async {
/// spawner.spawn(async {
/// println!("now running on a worker thread");
/// });
/// # }
@@ -47,27 +51,10 @@ impl TaskExecutor {
///
/// This function panics if the spawn fails. Failure occurs if the executor
/// is currently at capacity and is unable to spawn a new future.
pub fn spawn<F>(&self, future: F)
where F: Future<Output = ()> + Send + 'static,
pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
where
F: Future<Output = ()> + Send + 'static,
{
self.inner.spawn(future).unwrap();
}
}
impl tokio_executor::Executor for TaskExecutor {
fn spawn(
&mut self,
future: Pin<Box<dyn Future<Output = ()> + Send>>,
) -> Result<(), SpawnError> {
self.inner.spawn(future)
}
}
impl<T> tokio_executor::TypedExecutor<T> for TaskExecutor
where
T: Future<Output = ()> + Send + 'static,
{
fn spawn(&mut self, future: T) -> Result<(), crate::executor::SpawnError> {
crate::executor::Executor::spawn(self, Box::pin(future))
}
}
+16 -32
View File
@@ -47,14 +47,13 @@ fn spawn_shutdown() {
rt.spawn(client_server(tx.clone()));
// Use executor trait
let f = Box::pin(client_server(tx));
tokio_executor::Executor::spawn(&mut rt.executor(), f).unwrap();
// Use spawner
rt.spawner().spawn(client_server(tx));
rt.shutdown_on_idle();
assert_ok!(rx.recv());
assert_ok!(rx.recv());
assert_ok!(rx.try_recv());
assert_ok!(rx.try_recv());
rt.shutdown_now();
assert_err!(rx.try_recv());
}
@@ -68,8 +67,6 @@ fn block_on_timer() {
});
assert_eq!(v, 42);
rt.shutdown_on_idle();
}
#[test]
@@ -111,30 +108,35 @@ fn block_waits() {
});
assert_ok!(b_rx.try_recv());
rt.shutdown_on_idle();
}
#[test]
fn spawn_many() {
const ITER: usize = 200;
let cnt = Arc::new(Mutex::new(0));
let rt = Runtime::new().unwrap();
let cnt = Arc::new(Mutex::new(0));
let (tx, rx) = mpsc::channel();
let tx = Arc::new(Mutex::new(tx));
let c = cnt.clone();
rt.block_on(async move {
for _ in 0..ITER {
let c = c.clone();
let tx = tx.clone();
tokio::spawn(async move {
let mut x = c.lock().unwrap();
*x = 1 + *x;
if *x == ITER {
tx.lock().unwrap().send(()).unwrap();
}
});
}
});
rt.shutdown_on_idle();
rx.recv().unwrap();
assert_eq!(ITER, *cnt.lock().unwrap());
}
@@ -146,30 +148,12 @@ fn nested_enter() {
rt.block_on(async {
assert_err!(tokio_executor::enter());
// Since this is testing panics in other threads, printing about panics
// is noisy and can give the impression that the test is ignoring panics.
//
// It *is* ignoring them, but on purpose.
let prev_hook = panic::take_hook();
panic::set_hook(Box::new(|info| {
let s = info.to_string();
if s.starts_with("panicked at 'nested ")
|| s.starts_with("panicked at 'Multiple executors at once")
{
// expected, noop
} else {
println!("{}", s);
}
}));
let res = panic::catch_unwind(move || {
let rt = Runtime::new().unwrap();
rt.block_on(async {});
});
assert_err!(res);
panic::set_hook(prev_hook);
});
}
@@ -198,7 +182,7 @@ fn after_start_and_before_stop_is_called() {
rt.block_on(client_server(tx));
rt.shutdown_on_idle();
drop(rt);
assert_ok!(rx.try_recv());
-2
View File
@@ -23,8 +23,6 @@ fn timer_with_threaded_runtime() {
tx.send(()).unwrap();
});
rt.shutdown_on_idle();
rx.recv().unwrap();
}