threadpool: move threadpool into tokio-executor (#1452)

The threadpool is behind a feature flag.

Refs: #1264
This commit is contained in:
Carl Lerche
2019-08-15 13:09:02 -07:00
committed by GitHub
parent 37131b2114
commit 3b27dc31d2
49 changed files with 166 additions and 349 deletions
+176
View File
@@ -0,0 +1,176 @@
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
/// #![feature(async_await)]
///
/// 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.
/// 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 {}
+429
View File
@@ -0,0 +1,429 @@
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 log::trace;
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
///
/// ```
/// #![feature(async_await)]
///
/// 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!("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
@@ -0,0 +1,28 @@
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
@@ -0,0 +1,37 @@
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
@@ -0,0 +1,146 @@
//! 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};
@@ -0,0 +1,46 @@
use crate::park::{Park, Unpark};
use log::warn;
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| {
warn!(
"calling `park` on worker thread errored -- shutting down thread: {}",
e
);
})
}
fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> {
self.0.park_timeout(duration).map_err(|e| {
warn!(
"calling `park` on worker thread errored -- shutting down thread: {}",
e
);
})
}
}
@@ -0,0 +1,98 @@
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 {}
@@ -0,0 +1,8 @@
//! Thread parking utilities.
mod boxed;
mod default_park;
pub use self::default_park::{DefaultPark, DefaultUnpark, ParkError};
pub(crate) use self::boxed::{BoxPark, BoxUnpark, BoxedPark};
@@ -0,0 +1,308 @@
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()
}
}
@@ -0,0 +1,188 @@
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
}
}
+478
View File
@@ -0,0 +1,478 @@
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 log::{debug, error, trace};
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.
pub(crate) fn shutdown(&self, now: bool, purge_queue: bool) {
let mut state: State = self.state.load(Acquire).into();
trace!("shutdown; state={:?}", 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!(" -> submit internal; idx={}", 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(e) = res {
error!("failed to spawn worker thread; err={:?}", e);
panic!("failed to spawn worker thread: {:?}", e);
}
}
/// 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.
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 entry = &self.workers[idx];
debug_assert!(
worker_state.lifecycle() != Signaled,
"actual={:?}",
worker_state.lifecycle(),
);
trace!("signal_work -- notify; idx={}", idx);
if !entry.notify(worker_state) {
trace!("signal_work -- spawn; idx={}", idx);
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
@@ -0,0 +1,132 @@
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
}
}
+196
View File
@@ -0,0 +1,196 @@
use super::pool::{self, Lifecycle, Pool, MAX_FUTURES};
use super::task::Task;
use crate::{Executor, SpawnError, TypedExecutor};
use log::trace;
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
/// #![feature(async_await)]
///
/// 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!("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
@@ -0,0 +1,110 @@
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();
}
}
@@ -0,0 +1,499 @@
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()
}
}
@@ -0,0 +1,89 @@
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()
}
}
+317
View File
@@ -0,0 +1,317 @@
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 log::trace;
use std::cell::{Cell, UnsafeCell};
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, 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.
pub(crate) fn run(me: &Arc<Task>, pool: &Arc<Pool>) -> Run {
use self::State::*;
// 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),
}
trace!("Task::run; 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!();
}
}
@@ -0,0 +1,57 @@
#[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
}
}
@@ -0,0 +1,228 @@
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
/// #![feature(async_await)]
///
/// 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
@@ -0,0 +1,24 @@
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);
}
}
@@ -0,0 +1,329 @@
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()
}
}
+791
View File
@@ -0,0 +1,791 @@
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 log::trace;
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!(
"try_steal_task -- signal_work; self={}; from={}",
self.id.0,
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!("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.
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.
trace!("Worker::sleep; worker={:?}", self.id);
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!(" sleeping -- push to stack; idx={}", self.id.0);
// We obtained permission to push the worker into the
// sleeper queue.
if self.pool.push_sleeper(self.id.0).is_err() {
trace!(" sleeping -- push to stack failed; idx={}", self.id.0);
// 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; idx={}", self.id.0);
// 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; idx={}", self.id.0);
}
}
/// 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!("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
}
}
@@ -0,0 +1,257 @@
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()
}
}
@@ -0,0 +1,153 @@
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);
}
}