diff --git a/Cargo.toml b/Cargo.toml index f16ca80b2..0a5c5d7ef 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,9 @@ keywords = ["io", "async", "non-blocking", "futures"] members = [ "./", + "tokio-executor", "tokio-io", + "tokio-threadpool", ] [badges] @@ -31,12 +33,14 @@ travis-ci = { repository = "tokio-rs/tokio" } appveyor = { repository = "carllerche/tokio" } [dependencies] +tokio-io = "0.1" +tokio-executor = { version = "0.1", path = "tokio-executor" } +tokio-threadpool = { version = "0.1", path = "tokio-threadpool" } bytes = "0.4" log = "0.4" mio = "0.6.13" slab = "0.4" iovec = "0.1" -tokio-io = "0.1" futures = "0.1.16" [dev-dependencies] diff --git a/examples/chat.rs b/examples/chat.rs index 1b155427d..334a9b87b 100644 --- a/examples/chat.rs +++ b/examples/chat.rs @@ -439,6 +439,8 @@ pub fn main() { println!("accept error = {:?}", err); }); + println!("server running on localhost:6142"); + // This starts the `current_thread` executor. // // Executors are responsible for scheduling many asynchronous tasks, driving @@ -447,19 +449,16 @@ pub fn main() { // // The `current_thread` executor multiplexes all scheduled tasks on the // current thread. This means that spawned tasks must not implement `Send`. - current_thread::run(|_| { - // Now, the server task must be spawned. - // - // It's important to note that all futures / tasks are lazy. No work - // will happen unless they are spawned onto an executor. - current_thread::spawn(server); - - println!("server running on localhost:6142"); - - // The `current_thread::run` function will now block until *all* spawned - // tasks complete. - // - // In our example, we have not defined a shutdown strategy, so - // this will block until `ctrl-c` is pressed at the terminal. - }); + // It's important to note that all futures / tasks are lazy. No work will + // happen unless they are spawned onto an executor. + // + // The executor will start running the `server` task, which, in turn, spawns + // new tasks for each incoming connection. + // + // The `current_thread::block_on_all` function will block until *all* + // spawned tasks complete. + // + // In our example, we have not defined a shutdown strategy, so this will + // block until `ctrl-c` is pressed at the terminal. + current_thread::block_on_all(server).unwrap(); } diff --git a/examples/hello_world.rs b/examples/hello_world.rs index 54a951289..1731f9423 100644 --- a/examples/hello_world.rs +++ b/examples/hello_world.rs @@ -55,6 +55,8 @@ pub fn main() { println!("accept error = {:?}", err); }); + println!("server running on localhost:6142"); + // This starts the `current_thread` executor. // // Executors are responsible for scheduling many asynchronous tasks, driving @@ -62,21 +64,17 @@ pub fn main() { // implementations, each providing different scheduling characteristics. // // The `current_thread` executor multiplexes all scheduled tasks on the - // current thread. This means that spawned tasks are not required to - // implement `Send`. - current_thread::run(|_| { - // Now, the server task must be spawned. - // - // It's important to note that all futures / tasks are lazy. No work - // will happen unless they are spawned onto an executor. - current_thread::spawn(server); - - println!("server running on localhost:6142"); - - // The `current_thread::run` function will now block until *all* spawned - // tasks complete. - // - // In our example, we have not defined a shutdown strategy, so - // this will block until `ctrl-c` is pressed at the terminal. - }); + // current thread. This means that spawned tasks must not implement `Send`. + // It's important to note that all futures / tasks are lazy. No work will + // happen unless they are spawned onto an executor. + // + // The executor will start running the `server` task, which, in turn, spawns + // new tasks for each incoming connection. + // + // The `current_thread::block_on_all` function will block until *all* + // spawned tasks complete. + // + // In our example, we have not defined a shutdown strategy, so this will + // block until `ctrl-c` is pressed at the terminal. + current_thread::block_on_all(server).unwrap(); } diff --git a/src/executor/current_thread.rs b/src/executor/current_thread.rs deleted file mode 100644 index cc0d3d2de..000000000 --- a/src/executor/current_thread.rs +++ /dev/null @@ -1,412 +0,0 @@ -//! Execute tasks on the current thread -//! -//! This module implements an executor that keeps futures on the same thread -//! that they are submitted on. This allows it to execute futures that are -//! not `Send`. -//! -//! Before being able to spawn futures with this module, an executor -//! context must be setup by calling [`run`]. From within that context [`spawn`] -//! may be called with the future to run in the background. -//! -//! ``` -//! # extern crate tokio; -//! # extern crate futures; -//! # use tokio::executor::current_thread; -//! use futures::future::lazy; -//! -//! // Calling execute here results in a panic -//! // current_thread::spawn(my_future); -//! -//! # pub fn main() { -//! current_thread::run(|_| { -//! // The execution context is setup, futures may be executed. -//! current_thread::spawn(lazy(|| { -//! println!("called from the current thread executor"); -//! Ok(()) -//! })); -//! }); -//! # } -//! ``` -//! -//! # Execution model -//! -//! When an execution context is setup with `run` the current thread will block -//! and all the futures managed by the executor are driven to completion. -//! Whenever a future receives a notification, it is pushed to the end of a -//! scheduled list. The executor will drain this list, advancing the state of -//! each future. -//! -//! All futures managed by this module will remain on the current thread, -//! as such, this module is able to safely execute futures that are not `Send`. -//! -//! Once a future is complete, it is dropped. Once all futures are completed, -//! [`run`] will unblock and return. -//! -//! This module makes a best effort to fairly schedule futures that it manages. -//! -//! [`spawn`]: fn.spawn.html -//! [`run`]: fn.run.html - -use super::{scheduler}; -use super::sleep::{self, Sleep, Wakeup}; - -use futures::Async; -use futures::executor::{self, Spawn}; -use futures::future::{Future, Executor, ExecuteError, ExecuteErrorKind}; - -use std::{fmt, thread}; -use std::cell::Cell; -use std::rc::Rc; - -/// Executes futures on the current thread. -/// -/// All futures executed using this executor will be executed on the current -/// thread. As such, `run` will wait for these futures to complete before -/// returning. -/// -/// For more details, see the [module level](index.html) documentation. -#[derive(Debug, Clone)] -pub struct TaskExecutor { - // Prevent the handle from moving across threads. - _p: ::std::marker::PhantomData>, -} - -/// A context yielded to the closure provided to `run`. -/// -/// This context is mostly a future-proofing of the library to add future -/// contextual information into it. Currently it only contains the `Enter` -/// instance used to reserve the current thread for blocking on futures. -#[derive(Debug)] -pub struct Context<'a> { - cancel: &'a Cell, -} - -/// Implements the "blocking" logic for the current thread executor. A -/// `TaskRunner` will be created during `run` and will sit on the stack until -/// execution is complete. -#[derive(Debug)] -struct TaskRunner { - /// Executes futures. - scheduler: Scheduler, -} - -struct CurrentRunner { - /// When set to true, the executor should return immediately, even if there - /// still futures to run. - cancel: Cell, - - /// Number of futures currently being executed by the runner. - num_futures: Cell, - - /// Raw pointer to the current scheduler pusher. - /// - /// The raw pointer is required in order to store it in a thread-local slot. - schedule: Cell>, -} - -type Scheduler = scheduler::Scheduler; -type Schedule = scheduler::Schedule; - -struct Task(Spawn>>); - -/// Current thread's task runner. This is set in `TaskRunner::with` -thread_local!(static CURRENT: CurrentRunner = CurrentRunner { - cancel: Cell::new(false), - num_futures: Cell::new(0), - schedule: Cell::new(None), -}); - -/// Calls the given closure, then block until all futures submitted for -/// execution complete. -/// -/// In more detail, this function will block until: -/// - All executing futures are complete, or -/// - `cancel_all_spawned` is invoked. -pub fn run(f: F) -> R -where F: FnOnce(&mut Context) -> R -{ - sleep::BlockThread::with_current(|mut sleep| { - TaskRunner::enter(&mut sleep, f) - }) -} - -#[deprecated(since = "0.1.1", note = "this was never supposed to be public")] -#[doc(hidden)] -pub fn run_with_sleep(_: &mut S, _: F) -> R -where F: FnOnce(&mut Context) -> R, - S: Sleep, -{ - // This could never be called publically because `Sleep` is not public. - unimplemented!(); -} - -/// Executes a future on the current thread. -/// -/// The provided future must complete or be canceled before `run` will return. -/// -/// # Panics -/// -/// This function can only be invoked from the context of a `run` call; any -/// other use will result in a panic. -pub fn spawn(future: F) -where F: Future + 'static -{ - execute(future).unwrap_or_else(|_| { - panic!("cannot call `execute` unless the thread is already \ - in the context of a call to `run`") - }) -} - -/// Returns an executor that executes futures on the current thread. -/// -/// The user of `TaskExecutor` must ensure that when a future is submitted, -/// that it is done within the context of a call to `run`. -/// -/// For more details, see the [module level](index.html) documentation. -pub fn task_executor() -> TaskExecutor { - TaskExecutor { - _p: ::std::marker::PhantomData, - } -} - -impl Executor for TaskExecutor -where F: Future + 'static -{ - fn execute(&self, future: F) -> Result<(), ExecuteError> { - execute(future) - } -} - -impl<'a> Context<'a> { - /// Cancels *all* executing futures. - pub fn cancel_all_spawned(&self) { - self.cancel.set(true); - } -} - -/// Submits a future to the current executor. This is done by -/// checking the thread-local variable tracking the current executor. -/// -/// If this function is not called in context of an executor, i.e. outside of -/// `run`, then `Err` is returned. -/// -/// This function does not panic. -fn execute(future: F) -> Result<(), ExecuteError> -where F: Future + 'static, -{ - CURRENT.with(|current| { - match current.schedule.get() { - Some(schedule) => { - let spawned = Task::new(future); - - let num_futures = current.num_futures.get(); - current.num_futures.set(num_futures + 1); - - unsafe { (*schedule).schedule(spawned); } - - Ok(()) - } - None => { - Err(ExecuteError::new(ExecuteErrorKind::Shutdown, future)) - } - } - }) -} - -impl TaskRunner -where T: Wakeup, -{ - /// Return a new `TaskRunner` - fn new(wakeup: T) -> TaskRunner { - let scheduler = scheduler::Scheduler::new(wakeup); - - TaskRunner { - scheduler: scheduler, - } - } - - /// Enter a new `TaskRunner` context - /// - /// This function handles advancing the scheduler state and blocking while - /// listening for notified futures. - /// - /// First, a new task runner is created backed by the current - /// `sleep::BlockThread` handle. Passing `sleep::BlockThread` into the - /// scheduler is how scheduled futures unblock the thread, signalling that - /// there is more work to do. - /// - /// Before any future is polled, the scheduler must be set to a thread-local - /// variable so that `execute` is able to submit new futures to the current - /// executor. Because `Scheduler::schedule` requires `&mut self`, this - /// introduces a mutability hazard. This hazard is minimized with some - /// indirection. See `set_schedule` for more details. - /// - /// Once all context is setup, the init closure is invoked. This is the - /// "boostrapping" process that executes the initial futures into the - /// scheduler. After this, the function loops and advances the scheduler - /// state until all futures complete. When no scheduled futures are ready to - /// be advanced, the thread is blocked using `S: Sleep`. - fn enter(sleep: &mut S, f: F) -> R - where F: FnOnce(&mut Context) -> R, - S: Sleep, - { - let mut runner = TaskRunner::new(sleep.wakeup()); - - CURRENT.with(|current| { - // Make sure that another task runner is not set. - // - // This should not be ever possible due to how `set_schedule` - // is setup, but better safe than sorry! - assert!(current.schedule.get().is_none()); - - // Enter an execution scope - let mut ctx = Context { - cancel: ¤t.cancel, - }; - - // Set the scheduler to the TLS and perform setup work, - // returning a future to execute. - // - // This could possibly suubmit other futures for execution. - let ret = current.set_schedule(&mut runner.scheduler as &mut Schedule, || { - f(&mut ctx) - }); - - // Execute the runner. - // - // This function will not return until either - // - // a) All futures have completed execution - // b) `cancel_all_spawned` is called, forcing the executor to - // return. - runner.run(sleep, current); - - // Not technically required, but this makes the fact that `ctx` - // needs to live until this point explicit. - drop(ctx); - - ret - }) - } - - fn run(&mut self, sleep: &mut S, current: &CurrentRunner) - where S: Sleep, - { - use super::scheduler::Tick; - - while current.is_running() { - // Try to advance the scheduler state - let res = self.scheduler.tick(|scheduler, spawned, notify| { - // `scheduler` is a `&mut Scheduler` reference returned back - // from the scheduler to us, but only within the context of this - // closure. - // - // This lets us push new futures into the scheduler. It also - // lets us pass the scheduler mutable reference into - // `set_schedule`, which sets the thread-local variable that - // `spawn` uses for submitting new futures to the - // "current" executor. - // - // See `set_schedule` documentation for more details on how we - // guard against mutable pointer aliasing. - current.set_schedule(scheduler as &mut Schedule, || { - match spawned.0.poll_future_notify(notify, 0) { - Ok(Async::Ready(_)) | Err(_) => { - Async::Ready(()) - } - Ok(Async::NotReady) => Async::NotReady, - } - }) - }); - - // Process the result of ticking the scheduler - match res { - // A future completed. `is_daemon` is true when the future was - // submitted as a daemon future. - Tick::Data(_) => { - let num_futures = current.num_futures.get(); - debug_assert!(num_futures > 0); - current.num_futures.set(num_futures - 1); - }, - Tick::Empty => { - // The scheduler did not have any work to process. - // - // At this point, the scheduler is currently running given - // that the `while` condition was true and no user code has - // been executed. - - debug_assert!(current.is_running()); - - // Block the current thread until a future managed by the scheduler - // receives a readiness notification. - sleep.sleep(); - } - Tick::Inconsistent => { - // Yield the thread and loop - thread::yield_now(); - } - } - } - } -} - -impl CurrentRunner { - /// Set the provided schedule handle to the TLS slot for the duration of the - /// closure. - /// - /// `spawn` will access the CURRENT thread-local variable in - /// order to push a future into the scheduler. This requires a `&mut` - /// reference, introducing mutability hazards. - /// - /// Rust requires that `&mut` references are not aliases, i.e. there are - /// never two "live" mutable references to the same piece of data. In order - /// to store a `&mut` reference in a thread-local variable, we must ensure - /// that one can not access the scheduler anywhere else. - /// - /// To do this, we only allow access to the thread local variable from - /// within the closure passed to `set_schedule`. This function also takes a - /// &mut reference to the scheduler, which is essentially holding a "lock" - /// on that reference, preventing any other location in the code from - /// also getting that &mut reference. - /// - /// When `set_schedule` returns, the thread-local variable containing the - /// mut reference is set to null. This is done even if the closure panics. - /// - /// This reduces the odds of introducing pointer aliasing. - fn set_schedule(&self, schedule: &mut Schedule, f: F) -> R - where F: FnOnce() -> R - { - // Ensure that the runner is removed from the thread-local context - // when leaving the scope. This handles cases that involve panicking. - struct Reset<'a>(&'a CurrentRunner); - - impl<'a> Drop for Reset<'a> { - fn drop(&mut self) { - self.0.schedule.set(None); - } - } - - let _reset = Reset(self); - - self.schedule.set(Some(schedule as *mut Schedule)); - - f() - } - - fn is_running(&self) -> bool { - self.num_futures.get() > 0 && !self.cancel.get() - } -} - -impl Task { - fn new + 'static>(f: T) -> Self { - Task(executor::spawn(Box::new(f))) - } -} - -impl fmt::Debug for Task { - fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { - fmt.debug_struct("Task") - .finish() - } -} diff --git a/src/executor/current_thread/mod.rs b/src/executor/current_thread/mod.rs new file mode 100644 index 000000000..b6dc86aab --- /dev/null +++ b/src/executor/current_thread/mod.rs @@ -0,0 +1,722 @@ +//! Execute many tasks concurrently on the current thread. +//! +//! [`CurrentThread`] is an executor that keeps tasks on the same thread that +//! they were spawned from. This allows it to execute futures that are not +//! `Send`. +//! +//! A single [`CurrentThread`] instance is able to efficiently manage a large +//! number of tasks and will attempt to schedule all tasks fairly. +//! +//! All tasks that are being managed by a [`CurrentThread`] executor are able to +//! spawn additional tasks by calling [`spawn`]. This function only works from +//! within the context of a running [`CurrentThread`] instance. +//! +//! The easiest way to start a new [`CurrentThread`] executor is to call +//! [`block_on_all`] with an initial task to seed the executor. +//! +//! For example: +//! +//! ``` +//! # extern crate tokio; +//! # extern crate futures; +//! # use tokio::executor::current_thread; +//! use futures::future::lazy; +//! +//! // Calling execute here results in a panic +//! // current_thread::spawn(my_future); +//! +//! # pub fn main() { +//! current_thread::block_on_all(lazy(|| { +//! // The execution context is setup, futures may be executed. +//! current_thread::spawn(lazy(|| { +//! println!("called from the current thread executor"); +//! Ok(()) +//! })); +//! +//! Ok::<_, ()>(()) +//! })); +//! # } +//! ``` +//! +//! The `block_on_all` function will block the current thread until **all** +//! tasks that have been spawned onto the [`CurrentThread`] instance have +//! completed. +//! +//! More fine-grain control can be achieved by using [`CurrentThread`] directly. +//! +//! ``` +//! # extern crate tokio; +//! # extern crate futures; +//! # use tokio::executor::current_thread::CurrentThread; +//! use futures::future::{lazy, empty}; +//! use std::time::Duration; +//! +//! // Calling execute here results in a panic +//! // current_thread::spawn(my_future); +//! +//! # pub fn main() { +//! let mut current_thread = CurrentThread::new(); +//! +//! // Spawn a task, the task is not executed yet. +//! current_thread.spawn(lazy(|| { +//! println!("Spawning a task"); +//! Ok(()) +//! })); +//! +//! // Spawn a task that never completes +//! current_thread.spawn(empty()); +//! +//! // Run the executor, but only until the provided future completes. This +//! // provides the opportunity to start executing previously spawned tasks. +//! let res = current_thread.block_on(lazy(|| { +//! Ok::<_, ()>("Hello") +//! })).unwrap(); +//! +//! // Now, run the executor for *at most* 1 second. Since a task was spawned +//! // that never completes, this function will return with an error. +//! current_thread.run_timeout(Duration::from_secs(1)).unwrap_err(); +//! # } +//! ``` +//! +//! # Execution model +//! +//! Internally, [`CurrentThread`] maintains a queue. When one of its tasks is +//! notified, the task gets added to the queue. The executor will pop tasks from +//! the queue and call [`Future::poll`]. If the task gets notified while it is +//! being executed, it won't get re-executed until all other tasks currently in +//! the queue get polled. +//! +//! Before the task is polled, a thread-local variable referencing the current +//! [`CurrentThread`] instance is set. This enables [`spawn`] to spawn new tasks +//! onto the same executor without having to thread through a handle value. +//! +//! If the [`CurrentThread`] instance still has uncompleted tasks, but none of +//! these tasks are ready to be polled, the current thread is put to sleep. When +//! a task is notified, the thread is woken up and processing resumes. +//! +//! All tasks managed by [`CurrentThread`] remain on the current thread. When a +//! task completes, it is dropped. +//! +//! [`spawn`]: fn.spawn.html +//! [`block_on_all`]: fn.block_on_all.html +//! [`CurrentThread`]: struct.CurrentThread.html +//! [`Future::poll`]: https://docs.rs/futures/0.1/futures/future/trait.Future.html#tymethod.poll + +#![allow(deprecated)] + +mod scheduler; +use self::scheduler::Scheduler; + +use tokio_executor::{self, Enter, SpawnError}; +use tokio_executor::park::{Park, Unpark, ParkThread}; + +use futures::{executor, Async, Future}; +use futures::future::{self, Executor, ExecuteError, ExecuteErrorKind}; + +use std::fmt; +use std::cell::Cell; +use std::marker::PhantomData; +use std::rc::Rc; +use std::time::{Duration, Instant}; + +/// Executes tasks on the current thread +pub struct CurrentThread { + /// Execute futures and receive unpark notifications. + scheduler: Scheduler, + + /// Current number of futures being executed + num_futures: usize, + + /// Thread park handle + park: P, +} + +/// Executes futures on the current thread. +/// +/// All futures executed using this executor will be executed on the current +/// thread. As such, `run` will wait for these futures to complete before +/// returning. +/// +/// For more details, see the [module level](index.html) documentation. +#[derive(Debug, Clone)] +pub struct TaskExecutor { + // Prevent the handle from moving across threads. + _p: ::std::marker::PhantomData>, +} + +/// Returned by the `turn` function +#[derive(Debug)] +pub struct Turn(()); + +/// A `CurrentThread` instance bound to a supplied execution conext. +pub struct Entered<'a, P: Park + 'a> { + executor: &'a mut CurrentThread

, + enter: &'a mut Enter, +} + +#[deprecated(since = "0.1.2", note = "use block_on_all instead")] +#[doc(hidden)] +#[derive(Debug)] +pub struct Context<'a> { + cancel: Cell, + _p: PhantomData<&'a ()>, +} + +/// Error returned by the `run` function. +#[derive(Debug)] +pub struct RunError { + _p: (), +} + +/// Error returned by the `run_timeout` function. +#[derive(Debug)] +pub struct RunTimeoutError { + timeout: bool, +} + +/// Error returned by the `turn` function. +#[derive(Debug)] +pub struct TurnError { + _p: (), +} + +/// Error returned by the `block_on` function. +#[derive(Debug)] +pub struct BlockError { + inner: Option, +} + +/// This is mostly split out to make the borrow checker happy. +struct Borrow<'a, U: 'a> { + scheduler: &'a mut Scheduler, + num_futures: &'a mut usize, +} + +trait SpawnLocal { + fn spawn_local(&mut self, future: Box>); +} + +struct CurrentRunner { + spawn: Cell>, +} + +/// Current thread's task runner. This is set in `TaskRunner::with` +thread_local!(static CURRENT: CurrentRunner = CurrentRunner { + spawn: Cell::new(None), +}); + +#[deprecated(since = "0.1.2", note = "use block_on_all instead")] +#[doc(hidden)] +#[allow(deprecated)] +pub fn run(f: F) -> R +where F: FnOnce(&mut Context) -> R +{ + let mut context = Context { + cancel: Cell::new(false), + _p: PhantomData, + }; + + let mut current_thread = CurrentThread::new(); + + let ret = current_thread + .block_on(future::lazy(|| Ok::<_, ()>(f(&mut context)))) + .unwrap(); + + if context.cancel.get() { + return ret; + } + + current_thread.run().unwrap(); + ret +} + +/// Run the executor bootstrapping the execution with the provided future. +/// +/// This creates a new [`CurrentThread`] executor, spawns the provided future, +/// and blocks the current thread until the provided future and **all** +/// subsequently spawned futures complete. In other words: +/// +/// * If the provided boostrap future does **not** spawn any additional tasks, +/// `block_on_all` returns once `future` completes. +/// * If the provided bootstrap future **does** spawn additional tasks, then +/// `block_on_all` returns once **all** spawned futures complete. +/// +/// See [module level][mod] documentation for more details. +/// +/// [`CurrentThread`]: struct.CurrentThread.html +/// [mod]: index.html +pub fn block_on_all(future: F) -> Result +where F: Future, +{ + let mut current_thread = CurrentThread::new(); + + let ret = current_thread.block_on(future); + current_thread.run().unwrap(); + + ret.map_err(|e| e.into_inner().expect("unexpected execution error")) +} + +/// Executes a future on the current thread. +/// +/// The provided future must complete or be canceled before `run` will return. +/// +/// Unlike [`tokio::spawn`], this function will always spawn on a +/// `CurrentThread` executor and is able to spawn futures that are not `Send`. +/// +/// # Panics +/// +/// This function can only be invoked from the context of a `run` call; any +/// other use will result in a panic. +/// +/// [`tokio::spawn`]: ../fn.spawn.html +pub fn spawn(future: F) +where F: Future + 'static +{ + TaskExecutor::current() + .spawn_local(Box::new(future)) + .unwrap(); +} + +// ===== impl CurrentThread ===== + +impl CurrentThread { + /// Create a new instance of `CurrentThread`. + pub fn new() -> Self { + CurrentThread::new_with_park(ParkThread::new()) + } +} + +impl CurrentThread

{ + /// Create a new instance of `CurrentThread` backed by the given park + /// handle. + pub fn new_with_park(park: P) -> Self { + let unpark = park.unpark(); + + CurrentThread { + scheduler: Scheduler::new(unpark), + num_futures: 0, + park, + } + } + + /// Returns `true` if the executor is currently idle. + /// + /// An idle executor is defined by not currently having any spawned tasks. + pub fn is_idle(&self) -> bool { + self.num_futures == 0 + } + + /// Spawn the future on the executor. + /// + /// This internally queues the future to be executed once `run` is called. + pub fn spawn(&mut self, future: F) -> &mut Self + where F: Future + 'static, + { + self.borrow().spawn_local(Box::new(future)); + self + } + + /// Synchronously waits for the provided `future` to complete. + /// + /// This function can be used to synchronously block the current thread + /// until the provided `future` has resolved either successfully or with an + /// error. The result of the future is then returned from this function + /// call. + /// + /// Note that this function will **also** execute any spawned futures on the + /// current thread, but will **not** block until these other spawned futures + /// have completed. + /// + /// The caller is responsible for ensuring that other spawned futures + /// complete execution. + pub fn block_on(&mut self, future: F) + -> Result> + where F: Future + { + let mut enter = tokio_executor::enter().unwrap(); + self.enter(&mut enter).block_on(future) + } + + /// Run the executor to completion, blocking the thread until **all** + /// spawned futures have completed. + pub fn run(&mut self) -> Result<(), RunError> { + let mut enter = tokio_executor::enter().unwrap(); + self.enter(&mut enter).run() + } + + /// Run the executor to completion, blocking the thread until all + /// spawned futures have completed **or** `duration` time has elapsed. + pub fn run_timeout(&mut self, duration: Duration) + -> Result<(), RunTimeoutError> + { + let mut enter = tokio_executor::enter().unwrap(); + self.enter(&mut enter).run_timeout(duration) + } + + /// Perform a single iteration of the event loop + pub fn turn(&mut self, duration: Option) + -> Result + { + let mut enter = tokio_executor::enter().unwrap(); + self.enter(&mut enter).turn(duration) + } + + /// Bind `CurrentThread` instance with an execution context. + pub fn enter<'a>(&'a mut self, enter: &'a mut Enter) -> Entered<'a, P> { + Entered { + executor: self, + enter, + } + } + + fn borrow(&mut self) -> Borrow { + Borrow { + scheduler: &mut self.scheduler, + num_futures: &mut self.num_futures, + } + } +} + +impl tokio_executor::Executor for CurrentThread { + fn spawn(&mut self, future: Box + Send>) + -> Result<(), SpawnError> + { + self.borrow().spawn_local(future); + Ok(()) + } +} + +impl fmt::Debug for CurrentThread

{ + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + fmt.debug_struct("CurrentThread") + .field("scheduler", &self.scheduler) + .field("num_futures", &self.num_futures) + .finish() + } +} + +// ===== impl Entered ===== + +impl<'a, P: Park> Entered<'a, P> { + /// Spawn the future on the executor. + /// + /// This internally queues the future to be executed once `run` is called. + pub fn spawn(&mut self, future: F) -> &mut Self + where F: Future + 'static, + { + self.executor.borrow().spawn_local(Box::new(future)); + self + } + + /// Synchronously waits for the provided `future` to complete. + /// + /// This function can be used to synchronously block the current thread + /// until the provided `future` has resolved either successfully or with an + /// error. The result of the future is then returned from this function + /// call. + /// + /// Note that this function will **also** execute any spawned futures on the + /// current thread, but will **not** block until these other spawned futures + /// have completed. + /// + /// The caller is responsible for ensuring that other spawned futures + /// complete execution. + pub fn block_on(&mut self, future: F) + -> Result> + where F: Future + { + let mut future = executor::spawn(future); + let notify = self.executor.scheduler.notify(); + + loop { + let res = self.executor.borrow().enter(self.enter, || { + future.poll_future_notify(¬ify, 0) + }); + + match res { + Ok(Async::Ready(e)) => return Ok(e), + Err(e) => return Err(BlockError { inner: Some(e) }), + Ok(Async::NotReady) => {} + } + + self.tick(); + + if let Err(_) = self.executor.park.park() { + return Err(BlockError { inner: None }); + } + } + } + + /// Run the executor to completion, blocking the thread until **all** + /// spawned futures have completed. + pub fn run(&mut self) -> Result<(), RunError> { + self.run_timeout2(None) + .map_err(|_| RunError { _p: () }) + } + + /// Run the executor to completion, blocking the thread until all + /// spawned futures have completed **or** `duration` time has elapsed. + pub fn run_timeout(&mut self, duration: Duration) + -> Result<(), RunTimeoutError> + { + self.run_timeout2(Some(duration)) + } + + /// Perform a single iteration of the event loop + pub fn turn(&mut self, duration: Option) + -> Result + { + if !self.tick() { + let res = match duration { + Some(duration) => self.executor.park.park_timeout(duration), + None => self.executor.park.park(), + }; + + if res.is_err() { + return Err(TurnError { _p: () }); + } + + self.tick(); + } + + Ok(Turn(())) + } + + fn run_timeout2(&mut self, dur: Option) + -> Result<(), RunTimeoutError> + { + if self.executor.is_idle() { + // Nothing to do + return Ok(()); + } + + let mut time = dur.map(|dur| (Instant::now() + dur, dur)); + + loop { + self.tick(); + + if self.executor.is_idle() { + return Ok(()); + } + + match time { + Some((until, rem)) => { + if let Err(_) = self.executor.park.park_timeout(rem) { + return Err(RunTimeoutError::new(false)); + } + + let now = Instant::now(); + + if now >= until { + return Err(RunTimeoutError::new(true)); + } + + time = Some((until, until - now)); + } + None => { + if let Err(_) = self.executor.park.park() { + return Err(RunTimeoutError::new(false)); + } + } + } + } + } + + /// Returns `true` if any futures were processed + fn tick(&mut self) -> bool { + let num_futures = &mut self.executor.num_futures; + let enter = &mut *self.enter; + + // work the scheduler + self.executor.scheduler.tick(|scheduler, scheduled| { + let mut borrow = Borrow { + scheduler, + num_futures, + }; + + // A future completed, decrement the future count + if borrow.enter(enter, || scheduled.tick()) { + debug_assert!(*borrow.num_futures > 0); + *borrow.num_futures -= 1; + } + }) + } +} + +impl<'a, P: Park> fmt::Debug for Entered<'a, P> { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + fmt.debug_struct("Entered") + .field("executor", &self.executor) + .field("enter", &self.enter) + .finish() + } +} + +// ===== impl TaskExecutor ===== + +#[deprecated(since = "0.1.2", note = "use TaskExecutor::current instead")] +#[doc(hidden)] +pub fn task_executor() -> TaskExecutor { + TaskExecutor { + _p: ::std::marker::PhantomData, + } +} + +impl TaskExecutor { + /// Returns an executor that executes futures on the current thread. + /// + /// The user of `TaskExecutor` must ensure that when a future is submitted, + /// that it is done within the context of a call to `run`. + /// + /// For more details, see the [module level](index.html) documentation. + pub fn current() -> TaskExecutor { + TaskExecutor { + _p: ::std::marker::PhantomData, + } + } + + /// Spawn a future onto the current `CurrentThread` instance. + pub fn spawn_local(&mut self, future: Box>) + -> Result<(), SpawnError> + { + CURRENT.with(|current| { + match current.spawn.get() { + Some(spawn) => { + unsafe { (*spawn).spawn_local(future) }; + Ok(()) + } + None => { + Err(SpawnError::shutdown()) + } + } + }) + } +} + +impl tokio_executor::Executor for TaskExecutor { + fn spawn(&mut self, future: Box + Send>) + -> Result<(), SpawnError> + { + self.spawn_local(future) + } + + fn status(&self) -> Result<(), SpawnError> { + CURRENT.with(|current| { + if current.spawn.get().is_some() { + Ok(()) + } else { + Err(SpawnError::shutdown()) + } + }) + } +} + +impl Executor for TaskExecutor +where F: Future + 'static +{ + fn execute(&self, future: F) -> Result<(), ExecuteError> { + CURRENT.with(|current| { + match current.spawn.get() { + Some(spawn) => { + unsafe { (*spawn).spawn_local(Box::new(future)) }; + Ok(()) + } + None => { + Err(ExecuteError::new(ExecuteErrorKind::Shutdown, future)) + } + } + }) + } +} + +// ===== impl Context ===== + +impl<'a> Context<'a> { + /// Cancels *all* executing futures. + pub fn cancel_all_spawned(&self) { + self.cancel.set(true); + } +} + +// ===== impl Borrow ===== + +impl<'a, U: Unpark> Borrow<'a, U> { + fn enter(&mut self, _: &mut Enter, f: F) -> R + where F: FnOnce() -> R, + { + CURRENT.with(|current| { + current.set_spawn(self, || { + f() + }) + }) + } +} + +impl<'a, U: Unpark> SpawnLocal for Borrow<'a, U> { + fn spawn_local(&mut self, future: Box>) { + *self.num_futures += 1; + self.scheduler.schedule(future); + } +} + +// ===== impl CurrentRunner ===== + +impl CurrentRunner { + fn set_spawn(&self, spawn: &mut SpawnLocal, f: F) -> R + where F: FnOnce() -> R + { + struct Reset<'a>(&'a CurrentRunner); + + impl<'a> Drop for Reset<'a> { + fn drop(&mut self) { + self.0.spawn.set(None); + } + } + + let _reset = Reset(self); + + let spawn = unsafe { hide_lt(spawn as *mut SpawnLocal) }; + self.spawn.set(Some(spawn)); + + f() + } +} + +unsafe fn hide_lt<'a>(p: *mut (SpawnLocal + 'a)) -> *mut (SpawnLocal + 'static) { + use std::mem; + mem::transmute(p) +} + +// ===== impl RunTimeoutError ===== + +impl RunTimeoutError { + fn new(timeout: bool) -> Self { + RunTimeoutError { timeout } + } + + /// Returns `true` if the error was caused by the operation timeing out. + pub fn is_timeout(&self) -> bool { + self.timeout + } +} + +impl From for RunTimeoutError { + fn from(_: tokio_executor::EnterError) -> Self { + RunTimeoutError::new(false) + } +} + +// ===== impl BlockError ===== + +impl BlockError { + /// Returns the error yielded by the future being blocked on + pub fn into_inner(self) -> Option { + self.inner + } +} + +impl From for BlockError { + fn from(_: tokio_executor::EnterError) -> Self { + BlockError { inner: None } + } +} diff --git a/src/executor/scheduler.rs b/src/executor/current_thread/scheduler.rs similarity index 73% rename from src/executor/scheduler.rs rename to src/executor/current_thread/scheduler.rs index e9e6e1a31..2b05637aa 100644 --- a/src/executor/scheduler.rs +++ b/src/executor/current_thread/scheduler.rs @@ -1,46 +1,36 @@ -//! An unbounded set of futures. +use tokio_executor::park::Unpark; -use super::sleep::Wakeup; - -use futures::Async; -use futures::executor::{self, UnsafeNotify, NotifyHandle}; +use futures::{Future, Async}; +use futures::executor::{self, Spawn, UnsafeNotify, NotifyHandle}; use std::cell::UnsafeCell; use std::fmt::{self, Debug}; -use std::marker::PhantomData; use std::mem; use std::ptr; use std::sync::atomic::Ordering::{Relaxed, SeqCst, Acquire, Release, AcqRel}; -use std::sync::atomic::{AtomicPtr, AtomicBool}; +use std::sync::atomic::{AtomicPtr, AtomicBool, AtomicUsize}; use std::sync::{Arc, Weak}; use std::usize; +use std::thread; +use std::marker::PhantomData; /// A generic task-aware scheduler. /// /// This is used both by `FuturesUnordered` and the current-thread executor. -pub struct Scheduler { - inner: Arc>, - nodes: List, +pub struct Scheduler { + inner: Arc>, + nodes: List, } -/// Schedule new futures -pub trait Schedule { - /// Schedule a new future. - fn schedule(&mut self, item: T); -} - -pub struct Notify<'a, T: 'a, W: 'a>(&'a Arc>); +pub struct Notify<'a, U: 'a>(&'a Arc>); // A linked-list of nodes -struct List { +struct List { len: usize, - head: *const Node, - tail: *const Node, + head: *const Node, + tail: *const Node, } -unsafe impl Send for Scheduler {} -unsafe impl Sync for Scheduler {} - // Scheduler is implemented using two linked lists. The first linked list tracks // all items managed by a `Scheduler`. This list is stored on the `Scheduler` // struct and is **not** thread safe. The second linked list is an @@ -70,45 +60,51 @@ unsafe impl Sync for Scheduler {} // decremented. Once the node is popped from the mpsc channel, then the final // arc reference count can be decremented, thus freeing the node. -#[allow(missing_debug_implementations)] -struct Inner { - // The task using `Scheduler`. - wakeup: W, +struct Inner { + // Thread unpark handle + unpark: U, + + // Tick number + tick_num: AtomicUsize, // Head/tail of the readiness queue - head_readiness: AtomicPtr>, - tail_readiness: UnsafeCell<*const Node>, + head_readiness: AtomicPtr>, + tail_readiness: UnsafeCell<*const Node>, // Used as part of the MPSC queue algorithm - stub: Arc>, + stub: Arc>, } -struct Node { +unsafe impl Send for Inner {} +unsafe impl Sync for Inner {} + +impl executor::Notify for Inner { + fn notify(&self, _: usize) { + self.unpark.unpark(); + } +} + +struct Node { // The item - item: UnsafeCell>, + item: UnsafeCell>, + + // The tick at which this node was notified + notified_at: AtomicUsize, // Next pointer for linked list tracking all active nodes - next_all: UnsafeCell<*const Node>, + next_all: UnsafeCell<*const Node>, // Previous node in linked list tracking all active nodes - prev_all: UnsafeCell<*const Node>, + prev_all: UnsafeCell<*const Node>, // Next pointer in readiness queue - next_readiness: AtomicPtr>, + next_readiness: AtomicPtr>, // Whether or not this node is currently in the mpsc queue. queued: AtomicBool, // Queue that we'll be enqueued to when notified - queue: Weak>, -} - -/// Returned by the `Scheduler::tick` function, allowing the caller to decide -/// what action to take next. -pub enum Tick { - Data(T), - Empty, - Inconsistent, + queue: Weak>, } /// Returned by `Inner::dequeue`, representing either a dequeue success (with @@ -119,31 +115,43 @@ pub enum Tick { /// the future and the caller should try again soon. /// /// [1024cores]: http://www.1024cores.net/home/lock-free-algorithms/queues/intrusive-mpsc-node-based-queue -enum Dequeue { - Data(*const Node), +enum Dequeue { + Data(*const Node), Empty, Inconsistent, } -impl Scheduler -where W: Wakeup, +/// Wraps a spawned boxed future +struct Task(Spawn>>); + +/// A task that is scheduled. `turn` must be called +pub struct Scheduled<'a, U: 'a> { + task: &'a mut Task, + notify: &'a Notify<'a, U>, + done: &'a mut bool, +} + +impl Scheduler +where U: Unpark, { /// Constructs a new, empty `Scheduler` /// /// The returned `Scheduler` does not contain any items and, in this /// state, `Scheduler::poll` will return `Ok(Async::Ready(None))`. - pub fn new(wakeup: W) -> Self { + pub fn new(unpark: U) -> Self { let stub = Arc::new(Node { item: UnsafeCell::new(None), + notified_at: AtomicUsize::new(0), next_all: UnsafeCell::new(ptr::null()), prev_all: UnsafeCell::new(ptr::null()), next_readiness: AtomicPtr::new(ptr::null_mut()), queued: AtomicBool::new(true), queue: Weak::new(), }); - let stub_ptr = &*stub as *const Node; + let stub_ptr = &*stub as *const Node; let inner = Arc::new(Inner { - wakeup: wakeup, + unpark, + tick_num: AtomicUsize::new(0), head_readiness: AtomicPtr::new(stub_ptr as *mut _), tail_readiness: UnsafeCell::new(stub_ptr), stub: stub, @@ -154,27 +162,59 @@ where W: Wakeup, nodes: List::new(), } } -} -impl Scheduler { - /// Advance the scheduler state. + pub fn notify(&self) -> NotifyHandle { + self.inner.clone().into() + } + + pub fn schedule(&mut self, item: Box>) { + let node = Arc::new(Node { + item: UnsafeCell::new(Some(Task::new(item))), + notified_at: AtomicUsize::new(0), + next_all: UnsafeCell::new(ptr::null_mut()), + prev_all: UnsafeCell::new(ptr::null_mut()), + next_readiness: AtomicPtr::new(ptr::null_mut()), + queued: AtomicBool::new(true), + queue: Arc::downgrade(&self.inner), + }); + + // Right now our node has a strong reference count of 1. We transfer + // ownership of this reference count to our internal linked list + // and we'll reclaim ownership through the `unlink` function below. + let ptr = self.nodes.push_back(node); + + // We'll need to get the item "into the system" to start tracking it, + // e.g. getting its unpark notifications going to us tracking which + // items are ready. To do that we unconditionally enqueue it for + // polling here. + self.inner.enqueue(ptr); + } + + /// Advance the scheduler state, returning `true` if any futures were + /// processed. /// /// This function should be called whenever the caller is notified via a /// wakeup. - pub fn tick(&mut self, mut f: F) -> Tick - where F: FnMut(&mut Self, &mut T, &Notify) -> Async + pub fn tick(&mut self, mut f: F) -> bool + where F: FnMut(&mut Self, &mut Scheduled), { + let mut ret = false; + let tick = self.inner.tick_num.fetch_add(1, SeqCst); + loop { - let node = match unsafe { self.inner.dequeue() } { + let node = match unsafe { self.inner.dequeue(Some(tick)) } { Dequeue::Empty => { - return Tick::Empty; + return ret; } Dequeue::Inconsistent => { - return Tick::Inconsistent; + thread::yield_now(); + continue; } Dequeue::Data(node) => node, }; + ret = true; + debug_assert!(node != self.inner.stub()); unsafe { @@ -203,12 +243,12 @@ impl Scheduler { // assume is is complete (will return Ready or panic), in // which case we'll want to discard it regardless. // - struct Bomb<'a, T: 'a, W: 'a> { - queue: &'a mut Scheduler, - node: Option>>, + struct Bomb<'a, U: 'a> { + queue: &'a mut Scheduler, + node: Option>>, } - impl<'a, T, W> Drop for Bomb<'a, T, W> { + impl<'a, U> Drop for Bomb<'a, U> { fn drop(&mut self) { if let Some(node) = self.node.take() { release_node(node); @@ -221,10 +261,12 @@ impl Scheduler { queue: self, }; + let mut done = false; + // Now that the bomb holds the node, create a new scope. This // scope ensures that the borrow will go out of scope before we // mutate the node pointer in `bomb` again - let res = { + { let node = bomb.node.as_ref().unwrap(); // Get a reference to the inner future. We already ensured @@ -241,65 +283,65 @@ impl Scheduler { // Poll the underlying item with the appropriate `notify` // implementation. This is where a large bit of the unsafety // starts to stem from internally. The `notify` instance itself - // is basically just our `Arc>` and tracks the mpsc + // is basically just our `Arc` and tracks the mpsc // queue of ready items. // - // Critically though `Node` won't actually access `T`, the + // Critically though `Node` won't actually access `Task`, the // item, while it's floating around inside of `Task` // instances. These structs will basically just use `T` to size // the internal allocation, appropriately accessing fields and // deallocating the node if need be. let queue = &mut *bomb.queue; let notify = Notify(bomb.node.as_ref().unwrap()); - f(queue, item, ¬ify) - }; - let ret = match res { - Async::NotReady => { - // The future is not done, push it back into the "all - // node" list. - let node = bomb.node.take().unwrap(); - bomb.queue.nodes.push_back(node); - continue; - } - Async::Ready(v) => { - // `bomb` will take care of unlinking and releasing the - // node. - Tick::Data(v) - } - }; + let mut scheduled = Scheduled { + task: item, + notify: ¬ify, + done: &mut done, + }; - return ret + f(queue, &mut scheduled); + } + + if !done { + // The future is not done, push it back into the "all + // node" list. + let node = bomb.node.take().unwrap(); + bomb.queue.nodes.push_back(node); + } } } } } -impl Schedule for Scheduler { - fn schedule(&mut self, item: T) { - let node = Arc::new(Node { - item: UnsafeCell::new(Some(item)), - next_all: UnsafeCell::new(ptr::null_mut()), - prev_all: UnsafeCell::new(ptr::null_mut()), - next_readiness: AtomicPtr::new(ptr::null_mut()), - queued: AtomicBool::new(true), - queue: Arc::downgrade(&self.inner), - }); +impl<'a, U: Unpark> Scheduled<'a, U> { + /// Polls the task, returns `true` if the task has completed. + pub fn tick(&mut self) -> bool { + // Tick the future + let ret = match self.task.0.poll_future_notify(self.notify, 0) { + Ok(Async::Ready(_)) | Err(_) => true, + Ok(Async::NotReady) => false, + }; - // Right now our node has a strong reference count of 1. We transfer - // ownership of this reference count to our internal linked list - // and we'll reclaim ownership through the `unlink` function below. - let ptr = self.nodes.push_back(node); - - // We'll need to get the item "into the system" to start tracking it, - // e.g. getting its unpark notifications going to us tracking which - // items are ready. To do that we unconditionally enqueue it for - // polling here. - self.inner.enqueue(ptr); + *self.done = ret; + ret } } -fn release_node(node: Arc>) { +impl Task { + pub fn new(future: Box + 'static>) -> Self { + Task(executor::spawn(future)) + } +} + +impl fmt::Debug for Task { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + fmt.debug_struct("Task") + .finish() + } +} + +fn release_node(node: Arc>) { // The item is done, try to reset the queued flag. This will prevent // `notify` from doing any work in the item let prev = node.queued.swap(true, SeqCst); @@ -327,17 +369,17 @@ fn release_node(node: Arc>) { } } -impl Debug for Scheduler { +impl Debug for Scheduler { fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { write!(fmt, "Scheduler {{ ... }}") } } -impl Drop for Scheduler { +impl Drop for Scheduler { fn drop(&mut self) { // When a `Scheduler` is dropped we want to drop all items associated // with it. At the same time though there may be tons of `Task` handles - // flying around which contain `Node` references inside them. We'll + // flying around which contain `Node` references inside them. We'll // let those naturally get deallocated when the `Task` itself goes out // of scope or gets notified. while let Some(node) = self.nodes.pop_front() { @@ -348,7 +390,7 @@ impl Drop for Scheduler { // mpsc queue. None of those nodes, however, have items associated // with them so they're safe to destroy on any thread. At this point // the `Scheduler` struct, the owner of the one strong reference - // to `Inner` will drop the strong reference. At that point + // to `Inner` will drop the strong reference. At that point // whichever thread releases the strong refcount last (be it this // thread or some other thread as part of an `upgrade`) will clear out // the mpsc queue and free all remaining nodes. @@ -359,9 +401,9 @@ impl Drop for Scheduler { } } -impl Inner { +impl Inner { /// The enqueue function from the 1024cores intrusive MPSC queue algorithm. - fn enqueue(&self, node: *const Node) { + fn enqueue(&self, node: *const Node) { unsafe { debug_assert!((*node).queued.load(Relaxed)); @@ -379,7 +421,7 @@ impl Inner { /// /// Note that this unsafe as it required mutual exclusion (only one thread /// can call this) to be guaranteed elsewhere. - unsafe fn dequeue(&self) -> Dequeue { + unsafe fn dequeue(&self, tick: Option) -> Dequeue { let mut tail = *self.tail_readiness.get(); let mut next = (*tail).next_readiness.load(Acquire); @@ -393,6 +435,13 @@ impl Inner { next = (*next).next_readiness.load(Acquire); } + if let Some(tick) = tick { + // Only dequeue if the node matches the tick num + if (*tail).notified_at.load(SeqCst) != tick { + return Dequeue::Empty; + } + } + if !next.is_null() { *self.tail_readiness.get() = next; debug_assert!(tail != self.stub()); @@ -415,14 +464,14 @@ impl Inner { Dequeue::Inconsistent } - fn stub(&self) -> *const Node { + fn stub(&self) -> *const Node { &*self.stub } } -impl Drop for Inner { +impl Drop for Inner { fn drop(&mut self) { - // Once we're in the destructor for `Inner` we need to clear out the + // Once we're in the destructor for `Inner` we need to clear out the // mpsc queue of nodes if there's anything left in there. // // Note that each node has a strong reference count associated with it @@ -431,7 +480,7 @@ impl Drop for Inner { // so we're just pulling out nodes and dropping their refcounts. unsafe { loop { - match self.dequeue() { + match self.dequeue(None) { Dequeue::Empty => break, Dequeue::Inconsistent => abort("inconsistent in drop"), Dequeue::Data(ptr) => drop(ptr2arc(ptr)), @@ -441,7 +490,7 @@ impl Drop for Inner { } } -impl List { +impl List { fn new() -> Self { List { len: 0, @@ -451,7 +500,7 @@ impl List { } /// Prepends an element to the back of the list - fn push_back(&mut self, node: Arc>) -> *const Node { + fn push_back(&mut self, node: Arc>) -> *const Node { let ptr = arc2ptr(node); unsafe { @@ -475,7 +524,7 @@ impl List { } /// Pop an element from the front of the list - fn pop_front(&mut self) -> Option>> { + fn pop_front(&mut self) -> Option>> { if self.head.is_null() { // The list is empty return None; @@ -502,7 +551,7 @@ impl List { } /// Remove a specific node - unsafe fn remove(&mut self, node: *const Node) -> Arc> { + unsafe fn remove(&mut self, node: *const Node) -> Arc> { let node = ptr2arc(node); let next = *node.next_all.get(); let prev = *node.prev_all.get(); @@ -527,69 +576,67 @@ impl List { } } -impl<'a, T, W> Clone for Notify<'a, T, W> { +impl<'a, U> Clone for Notify<'a, U> { fn clone(&self) -> Self { Notify(self.0) } } -impl<'a, T: fmt::Debug, W: fmt::Debug> fmt::Debug for Notify<'a, T, W> { +impl<'a, U> fmt::Debug for Notify<'a, U> { fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { fmt.debug_struct("Notiy").finish() } } -impl<'a, T, W: Wakeup> From> for NotifyHandle { - fn from(handle: Notify<'a, T, W>) -> NotifyHandle { +impl<'a, U: Unpark> From> for NotifyHandle { + fn from(handle: Notify<'a, U>) -> NotifyHandle { unsafe { let ptr = handle.0.clone(); - let ptr = mem::transmute::>, *mut ArcNode>(ptr); + let ptr = mem::transmute::>, *mut ArcNode>(ptr); NotifyHandle::new(hide_lt(ptr)) } } } -struct ArcNode(PhantomData<(T, W)>); +struct ArcNode(PhantomData); -// We should never touch `T` on any thread other than the one owning +// We should never touch `Task` on any thread other than the one owning // `Scheduler`, so this should be a safe operation. -// -// `W` already requires `Sync + Send` -unsafe impl Send for ArcNode {} -unsafe impl Sync for ArcNode {} +unsafe impl Send for ArcNode {} +unsafe impl Sync for ArcNode {} -impl executor::Notify for ArcNode { +impl executor::Notify for ArcNode { fn notify(&self, _id: usize) { unsafe { - let me: *const ArcNode = self; - let me: *const *const ArcNode = &me; - let me = me as *const Arc>; + let me: *const ArcNode = self; + let me: *const *const ArcNode = &me; + let me = me as *const Arc>; Node::notify(&*me) } } } -unsafe impl UnsafeNotify for ArcNode { +unsafe impl UnsafeNotify for ArcNode { unsafe fn clone_raw(&self) -> NotifyHandle { - let me: *const ArcNode = self; - let me: *const *const ArcNode = &me; - let me = &*(me as *const Arc>); + let me: *const ArcNode = self; + let me: *const *const ArcNode = &me; + let me = &*(me as *const Arc>); Notify(me).into() } unsafe fn drop_raw(&self) { - let mut me: *const ArcNode = self; - let me = &mut me as *mut *const ArcNode as *mut Arc>; + let mut me: *const ArcNode = self; + let me = &mut me as *mut *const ArcNode as *mut Arc>; ptr::drop_in_place(me); } } -unsafe fn hide_lt(p: *mut ArcNode) -> *mut UnsafeNotify { +unsafe fn hide_lt(p: *mut ArcNode) -> *mut UnsafeNotify { mem::transmute(p as *mut UnsafeNotify) } -impl Node { - fn notify(me: &Arc>) { +impl Node { + fn notify(me: &Arc>) { let inner = match me.queue.upgrade() { Some(inner) => inner, None => return, @@ -611,15 +658,19 @@ impl Node { // still. let prev = me.queued.swap(true, SeqCst); if !prev { + // Get the current scheduler tick + let tick_num = inner.tick_num.load(SeqCst); + me.notified_at.store(tick_num, SeqCst); + inner.enqueue(&**me); - inner.wakeup.wakeup(); + inner.unpark.unpark(); } } } -impl Drop for Node { +impl Drop for Node { fn drop(&mut self) { - // Currently a `Node` is sent across all threads for any lifetime, + // Currently a `Node` is sent across all threads for any lifetime, // regardless of `T`. This means that for memory safety we can't // actually touch `T` at any time except when we have a reference to the // `Scheduler` itself. diff --git a/src/executor/mod.rs b/src/executor/mod.rs index 896676c81..9b2b6615b 100644 --- a/src/executor/mod.rs +++ b/src/executor/mod.rs @@ -1,8 +1,207 @@ //! Task execution utilities. //! -//! This module only contains `current_thread`, an executor for multiplexing -//! many tasks on a single thread. +//! In the Tokio execution model, futures are lazy. When a future is created, no +//! work is performed. In order for the work defined by the future to happen, +//! the future must be submitted to an executor. A future that is submitted to +//! an executor is called a "task". +//! +//! The executor executor is responsible for ensuring that [`Future::poll`] is +//! called whenever the task is [notified]. Notification happens when the +//! internal state of a task transitions from "not ready" to ready. For +//! example, a socket might have received data and a call to `read` will now be +//! able to succeed. +//! +//! The specific strategy used to manage the tasks is left up to the +//! executor. There are two main flavors of executors: single-threaded and +//! multithreaded. This module provides both. +//! +//! * **[`current_thread`]**: A single-threaded executor that support spawning +//! tasks that are not `Send`. It guarantees that tasks will be executed on +//! the same thread from which they are spawned. +//! +//! * **[`thread_pool`]**: A multi-threaded executor that maintains a pool of +//! threads. Tasks are spawned to one of the threads in the pool and executed. +//! The pool employes a [work-stealing] strategy for optimizing how tasks get +//! spread across the available threads. +//! +//! # `Executor` trait. +//! +//! This module provides the [`Executor`] trait (re-exported from +//! [`tokio-executor`]), which describes the API that all executors must +//! implement. +//! +//! A free [`spawn`] function is provided that allows spawning futures onto the +//! default executor (tracked via a thread-local variable) without referencing a +//! handle. It is expected that all executors will set a value for the default +//! executor. This value will often be set to the executor itself, but it is +//! possible that the default executor might be set to a different executor. +//! +//! For example, the [`current_thread`] executor might set the default executor +//! to a thread pool instead of itself, allowing futures to spawn new tasks onto +//! the thread pool when those tasks are `Send`. +//! +//! [`Future::poll`]: https://docs.rs/futures/0.1/futures/future/trait.Future.html#tymethod.poll +//! [notified]: https://docs.rs/futures/0.1/futures/executor/trait.Notify.html#tymethod.notify +//! [`current_thread`]: current_thread/index.html +//! [`thread_pool`]: thread_pool/index.html +//! [work-stealing]: https://en.wikipedia.org/wiki/Work_stealing +//! [`tokio-executor`]: # +//! [`Executor`]: # +//! [`spawn`]: # + pub mod current_thread; -mod scheduler; -mod sleep; + +pub mod thread_pool { + //! Maintains a pool of threads across which the set of spawned tasks are + //! executed. + //! + //! [`ThreadPool`] is an executor that uses a thread pool for executing + //! tasks concurrently across multiple cores. It uses a thread pool that is + //! optimized for use cases that involve multiplexing large number of + //! independent tasks that perform short(ish) amounts of computation and are + //! mainly waiting on I/O, i.e. the Tokio use case. + //! + //! Usually, users of [`ThreadPool`] will not create pool instances. + //! Instead, they will create a [`Runtime`] instance, which comes with a + //! pre-configured thread pool. + //! + //! At the core, [`ThreadPool`] uses a work-stealing based scheduling + //! strategy. When spawning a task while *external* to the thread pool + //! (i.e., from a thread that is not part of the thread pool), the task is + //! randomly assigned to a worker thread. When spawning a task while + //! *internal* to the thread pool, the task is assigned to the current + //! worker. + //! + //! Each worker maintains its own queue and first focuses on processing all + //! tasks in its queue. When the worker's queue is empty, the worker will + //! attempt to *steal* tasks from other worker queues. This strategy helps + //! ensure that work is evenly distributed across threads while minimizing + //! synchronization between worker threads. + //! + //! # Usage + //! + //! Thread pool instances are created using [`ThreadPool::new`] or + //! [`Builder::new`]. The first option returns a thread pool with default + //! configuration values. The second option allows configuring the thread + //! pool before instantiating it. + //! + //! Once an instance is obtained, futures may be spawned onto it using the + //! [`spawn`] function. + //! + //! A handle to the thread pool is obtained using [`ThreadPool::sender`]. + //! This handle is **only** able to spawn futures onto the thread pool. It + //! is unable to affect the lifecycle of the thread pool in any way. This + //! handle can be passed into functions or stored in structs as a way to + //! grant the capability of spawning futures. + //! + //! # Examples + //! + //! ```rust + //! # extern crate tokio; + //! # extern crate futures; + //! # use tokio::executor::thread_pool::ThreadPool; + //! use futures::future::{Future, lazy}; + //! + //! # pub fn main() { + //! // Create a thread pool with default configuration values + //! let thread_pool = ThreadPool::new(); + //! + //! thread_pool.spawn(lazy(|| { + //! println!("called from a worker thread"); + //! Ok(()) + //! })); + //! + //! // Gracefully shutdown the threadpool + //! thread_pool.shutdown().wait().unwrap(); + //! # } + //! ``` + //! + //! [`ThreadPool`]: struct.ThreadPool.html + //! [`ThreadPool::new`]: struct.ThreadPool.html#method.new + //! [`ThreadPool::sender`]: struct.ThreadPool.html#method.sender + //! [`spawn`]: struct.ThreadPool.html#method.spawn + //! [`Builder::new`]: struct.Builder.html#method.new + //! [`Runtime`]: ../../runtime/struct.Runtime.html + + pub use tokio_threadpool::{ + Builder, + Sender, + Shutdown, + ThreadPool, + }; +} + +pub use tokio_executor::{Executor, DefaultExecutor, SpawnError}; + +use futures::{Future, Poll, Async}; + +/// Future, returned by `spawn`, that completes once the future is spawned. +/// +/// See [`spawn`] for more details. +/// +/// [`spawn`]: fn.spawn.html +#[derive(Debug)] +#[must_use = "Spawn does nothing unless polled"] +pub struct Spawn(Option); + +/// Spawns a future on the default executor. +/// +/// In order for a future to do work, it must be spawned on an executor. The +/// `spawn` function is the easiest way to do this. It spawns a future on the +/// [default executor] for the current execution context (tracked using a +/// thread-local variable). +/// +/// The default executor is **usually** a thread pool. +/// +/// Note that the function doesn't immediately spawn the future. Instead, it +/// returns `Spawn`, which itself is a future that completes once the spawn has +/// succeeded. +/// +/// # Examples +/// +/// In this example, a server is started and `spawn` is used to start a new task +/// that processes each received connection. +/// +/// ```rust +/// # extern crate tokio; +/// # extern crate futures; +/// # use futures::{Future, Stream}; +/// use tokio::net::TcpListener; +/// +/// # fn process(_: T) -> Box + Send> { +/// # unimplemented!(); +/// # } +/// # fn dox() { +/// # let addr = "127.0.0.1:8080".parse().unwrap(); +/// let listener = TcpListener::bind(&addr).unwrap(); +/// +/// let server = listener.incoming() +/// .map_err(|e| println!("error = {:?}", e)) +/// .for_each(|socket| { +/// tokio::spawn(process(socket)) +/// }); +/// +/// tokio::run(server); +/// # } +/// # pub fn main() {} +/// ``` +/// +/// [default executor]: struct.DefaultExecutor.html +pub fn spawn(f: F) -> Spawn +where F: Future + 'static + Send +{ + Spawn(Some(f)) +} + +impl Future for Spawn +where F: Future + Send + 'static +{ + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + ::tokio_executor::spawn(self.0.take().unwrap()); + Ok(Async::Ready(())) + } +} diff --git a/src/executor/sleep.rs b/src/executor/sleep.rs deleted file mode 100644 index 7058d81a2..000000000 --- a/src/executor/sleep.rs +++ /dev/null @@ -1,169 +0,0 @@ -use futures::executor::Notify; - -use std::fmt; -use std::sync::{Arc, Mutex, Condvar}; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::{Duration, Instant}; - -/// Puts the current thread to sleep. -pub trait Sleep { - /// Wake up handle. - type Wakeup: Wakeup; - - /// Get a new `Wakeup` handle. - fn wakeup(&self) -> Self::Wakeup; - - /// Put the current thread to sleep. - fn sleep(&mut self); - - /// Put the current thread to sleep for at most `duration`. - fn sleep_timeout(&mut self, duration: Duration); -} - -/// Wake up a sleeping thread. -pub trait Wakeup: Clone + Send + 'static { - /// Wake up the sleeping thread. - fn wakeup(&self); -} - -/// Blocks the current thread -pub struct BlockThread { - state: AtomicUsize, - mutex: Mutex<()>, - condvar: Condvar, -} - -const IDLE: usize = 0; -const NOTIFY: usize = 1; -const SLEEP: usize = 2; - -thread_local! { - static CURRENT_THREAD_NOTIFY: Arc = Arc::new(BlockThread { - state: AtomicUsize::new(IDLE), - mutex: Mutex::new(()), - condvar: Condvar::new(), - }); -} - -// ===== impl BlockThread ===== - -impl BlockThread { - pub fn with_current(f: F) -> R - where F: FnOnce(&Arc) -> R, - { - CURRENT_THREAD_NOTIFY.with(|notify| f(notify)) - } - - pub fn park(&self) { - self.park_timeout(None); - } - - pub fn park_timeout(&self, dur: Option) { - // If currently notified, then we skip sleeping. This is checked outside - // of the lock to avoid acquiring a mutex if not necessary. - match self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) { - NOTIFY => return, - IDLE => {}, - _ => unreachable!(), - } - - // The state is currently idle, so obtain the lock and then try to - // transition to a sleeping state. - let mut m = self.mutex.lock().unwrap(); - - // Transition to sleeping - match self.state.compare_and_swap(IDLE, SLEEP, Ordering::SeqCst) { - NOTIFY => { - // Notified before we could sleep, consume the notification and - // exit - self.state.store(IDLE, Ordering::SeqCst); - return; - } - IDLE => {}, - _ => unreachable!(), - } - - // Track (until, remaining) - let mut time = dur.map(|dur| (Instant::now() + dur, dur)); - - loop { - m = match time { - Some((until, rem)) => { - let (guard, _) = self.condvar.wait_timeout(m, rem).unwrap(); - let now = Instant::now(); - - if now >= until { - // Timed out... exit sleep state - self.state.store(IDLE, Ordering::SeqCst); - return; - } - - time = Some((until, until - now)); - guard - } - None => self.condvar.wait(m).unwrap(), - }; - - // Transition back to idle, loop otherwise - if NOTIFY == self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) { - return; - } - } - } - - fn unpark(&self) { - // First, try transitioning from IDLE -> NOTIFY, this does not require a - // lock. - match self.state.compare_and_swap(IDLE, NOTIFY, Ordering::SeqCst) { - IDLE | NOTIFY => return, - SLEEP => {} - _ => unreachable!(), - } - - // The other half is sleeping, this requires a lock - let _m = self.mutex.lock().unwrap(); - - // Transition from SLEEP -> NOTIFY - match self.state.compare_and_swap(SLEEP, NOTIFY, Ordering::SeqCst) { - SLEEP => {} - _ => return, - } - - // Wakeup the sleeper - self.condvar.notify_one(); - } -} - -impl Notify for BlockThread { - fn notify(&self, _unpark_id: usize) { - self.unpark(); - } -} - -impl<'a> Sleep for &'a Arc { - type Wakeup = Arc; - - fn wakeup(&self) -> Self::Wakeup { - (*self).clone() - } - - fn sleep(&mut self) { - self.park(); - } - - fn sleep_timeout(&mut self, duration: Duration) { - self.park_timeout(Some(duration)); - } -} - -impl Wakeup for Arc { - fn wakeup(&self) { - self.unpark(); - } -} - -impl fmt::Debug for BlockThread { - fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { - fmt.debug_struct("BlockThread").finish() - } -} diff --git a/src/lib.rs b/src/lib.rs index aaf765608..11dfc007c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -39,50 +39,45 @@ //! //! ```no_run //! extern crate futures; -//! extern crate futures_cpupool; //! extern crate tokio; //! extern crate tokio_io; //! //! use futures::prelude::*; -//! use futures::future::Executor; -//! use futures_cpupool::CpuPool; //! use tokio_io::AsyncRead; //! use tokio_io::io::copy; //! use tokio::net::TcpListener; //! //! fn main() { -//! let pool = CpuPool::new_num_cpus(); -//! //! // Bind the server's socket. //! let addr = "127.0.0.1:12345".parse().unwrap(); //! let listener = TcpListener::bind(&addr) //! .expect("unable to bind TCP listener"); //! //! // Pull out a stream of sockets for incoming connections -//! let server = listener.incoming().for_each(|sock| { -//! // Split up the reading and writing parts of the -//! // socket. -//! let (reader, writer) = sock.split(); +//! let server = listener.incoming() +//! .map_err(|e| println!("accept failed = {:?}", e)) +//! .for_each(|sock| { +//! // Split up the reading and writing parts of the +//! // socket. +//! let (reader, writer) = sock.split(); //! -//! // A future that echos the data and returns how -//! // many bytes were copied... -//! let bytes_copied = copy(reader, writer); +//! // A future that echos the data and returns how +//! // many bytes were copied... +//! let bytes_copied = copy(reader, writer); //! -//! // ... after which we'll print what happened. -//! let handle_conn = bytes_copied.map(|amt| { -//! println!("wrote {:?} bytes", amt) -//! }).map_err(|err| { -//! eprintln!("IO error {:?}", err) +//! // ... after which we'll print what happened. +//! let handle_conn = bytes_copied.map(|amt| { +//! println!("wrote {:?} bytes", amt) +//! }).map_err(|err| { +//! eprintln!("IO error {:?}", err) +//! }); +//! +//! // Spawn the future as a concurrent task. +//! tokio::spawn(handle_conn) //! }); //! -//! // Spawn the future as a concurrent task. -//! pool.execute(handle_conn).unwrap(); -//! -//! Ok(()) -//! }); -//! -//! // Spin up the server on this thread -//! server.wait().unwrap(); +//! // Start the Tokio runtime +//! tokio::run(server); //! } //! ``` @@ -99,6 +94,8 @@ extern crate mio; extern crate slab; #[macro_use] extern crate tokio_io; +extern crate tokio_executor; +extern crate tokio_threadpool; #[macro_use] extern crate log; @@ -106,3 +103,7 @@ extern crate log; pub mod executor; pub mod net; pub mod reactor; +pub mod runtime; + +pub use executor::spawn; +pub use runtime::run; diff --git a/src/reactor/background.rs b/src/reactor/background.rs new file mode 100644 index 000000000..6d4e83115 --- /dev/null +++ b/src/reactor/background.rs @@ -0,0 +1,209 @@ +use std::io; +use std::thread; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::SeqCst; + +use reactor::{Reactor, Handle}; +use futures::{Future, Async, Poll}; +use futures::task::AtomicTask; + +/// Handle to the reactor running on a background thread. +#[derive(Debug)] +pub struct Background { + /// When `None`, the reactor thread will run until the process terminates. + inner: Option, +} + +/// Future that resolves when the reactor thread has shutdown. +#[derive(Debug)] +pub struct Shutdown { + inner: Inner, +} + +/// Actual Background handle. +#[derive(Debug)] +struct Inner { + /// Handle to the reactor + handle: Handle, + + /// Shared state between the background handle and the reactor thread. + shared: Arc, +} + +#[derive(Debug)] +struct Shared { + /// Signal the reactor thread to shutdown. + shutdown: AtomicUsize, + + /// Task to notify when the reactor thread enters a shutdown state. + shutdown_task: AtomicTask, +} + +/// Notifies the reactor thread to shutdown once the reactor becomes idle. +const SHUTDOWN_IDLE: usize = 1; + +/// Notifies the reactor thread to shutdown immediately. +const SHUTDOWN_NOW: usize = 2; + +/// The reactor is currently shutdown. +const SHUTDOWN: usize = 3; + +// ===== impl Background ===== + +impl Background { + /// Launch a reactor in the background and return a handle to the thread. + pub fn new(reactor: Reactor) -> io::Result { + // Grab a handle to the reactor + let handle = reactor.handle().clone(); + + // Create the state shared between the background handle and the reactor + // thread. + let shared = Arc::new(Shared { + shutdown: AtomicUsize::new(0), + shutdown_task: AtomicTask::new(), + }); + + // For the reactor thread + let shared2 = shared.clone(); + + // Start the reactor thread + thread::Builder::new() + .spawn(move || run(reactor, shared2))?; + + Ok(Background { + inner: Some(Inner { + handle, + shared, + }), + }) + } + + /// Returns a reference to the reactor handle. + pub fn handle(&self) -> &Handle { + &self.inner.as_ref().unwrap().handle + } + + /// Shutdown the reactor on idle. + /// + /// Returns a future that completes once the reactor thread has shutdown. + pub fn shutdown_on_idle(mut self) -> Shutdown { + let inner = self.inner.take().unwrap(); + inner.shutdown_on_idle(); + + Shutdown { inner } + } + + /// Shutdown the reactor immediately + /// + /// Returns a future that completes once the reactor thread has shutdown. + pub fn shutdown_now(mut self) -> Shutdown { + let inner = self.inner.take().unwrap(); + inner.shutdown_now(); + + Shutdown { inner } + } + + /// Run the reactor on its thread until the process terminates. + pub fn forget(mut self) { + drop(self.inner.take()); + } +} + +impl Drop for Background { + fn drop(&mut self) { + let inner = match self.inner.take() { + Some(i) => i, + None => return, + }; + + let shutdown = Shutdown { inner }; + let _ = shutdown.wait(); + } +} + +// ===== impl Shutdown ===== + +impl Future for Shutdown { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + self.inner.shared.shutdown_task.register(); + + if !self.inner.is_shutdown() { + return Ok(Async::NotReady); + } + + Ok(().into()) + } +} + +// ===== impl Inner ===== + +impl Inner { + /// Returns true if the reactor thread is shutdown. + fn is_shutdown(&self) -> bool { + self.shared.shutdown.load(SeqCst) == SHUTDOWN + } + + /// Notify the reactor thread to shutdown once the reactor transitions to an + /// idle state. + fn shutdown_on_idle(&self) { + self.shared.shutdown + .compare_and_swap(0, SHUTDOWN_IDLE, SeqCst); + self.handle.wakeup(); + } + + /// Notify the reactor thread to shutdown immediately. + fn shutdown_now(&self) { + let mut curr = self.shared.shutdown.load(SeqCst); + + loop { + if curr >= SHUTDOWN_NOW { + return; + } + + let act = self.shared.shutdown + .compare_and_swap(curr, SHUTDOWN_NOW, SeqCst); + + if act == curr { + self.handle.wakeup(); + return; + } + + curr = act; + } + } +} + +// ===== impl Reactor thread ===== + +fn run(mut reactor: Reactor, shared: Arc) { + debug!("starting background reactor"); + loop { + let shutdown = shared.shutdown.load(SeqCst); + + if shutdown == SHUTDOWN_NOW { + debug!("shutting background reactor down NOW"); + break; + } + + if shutdown == SHUTDOWN_IDLE && reactor.is_idle() { + debug!("shutting background reactor on idle"); + break; + } + + reactor.turn(None).unwrap(); + } + + drop(reactor); + + // Transition the state to shutdown + shared.shutdown.store(SHUTDOWN, SeqCst); + + // Notify any waiters + shared.shutdown_task.notify(); + + debug!("background reactor has shutdown"); +} diff --git a/src/reactor/global.rs b/src/reactor/global.rs deleted file mode 100644 index b295326c9..000000000 --- a/src/reactor/global.rs +++ /dev/null @@ -1,54 +0,0 @@ -use std::io; -use std::thread; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; - -use reactor::{Reactor, Handle}; - -pub struct HelperThread { - thread: Option>, - reactor: Handle, - done: Arc, -} - -impl HelperThread { - pub fn new() -> io::Result { - let reactor = Reactor::new()?; - let reactor_handle = reactor.handle().clone(); - let done = Arc::new(AtomicBool::new(false)); - let done2 = done.clone(); - let thread = thread::Builder::new().spawn(move || run(reactor, done))?; - - Ok(HelperThread { - thread: Some(thread), - reactor: reactor_handle, - done: done2, - }) - } - - pub fn handle(&self) -> &Handle { - &self.reactor - } - - pub fn forget(mut self) { - drop(self.thread.take()); - } -} - -impl Drop for HelperThread { - fn drop(&mut self) { - let thread = match self.thread.take() { - Some(thread) => thread, - None => return - }; - self.done.store(true, Ordering::SeqCst); - self.reactor.wakeup(); - thread.join().unwrap(); - } -} - -fn run(mut reactor: Reactor, done: Arc) { - while !done.load(Ordering::SeqCst) { - reactor.turn(None).unwrap(); - } -} diff --git a/src/reactor/mod.rs b/src/reactor/mod.rs index 56758dfb3..f359f89f2 100644 --- a/src/reactor/mod.rs +++ b/src/reactor/mod.rs @@ -16,9 +16,13 @@ //! [`PollEvented`]: struct.PollEvented.html //! [`TcpStream`]: ../net/struct.TcpStream.html +use tokio_executor::Enter; +use tokio_executor::park::{Park, Unpark}; + use std::{fmt, usize}; use std::io::{self, ErrorKind}; use std::mem; +use std::cell::RefCell; use std::sync::atomic::Ordering::{Relaxed, SeqCst}; use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT}; use std::sync::{Arc, Weak, RwLock}; @@ -30,7 +34,8 @@ use mio; use mio::event::Evented; use slab::Slab; -mod global; +pub(crate) mod background; +use self::background::Background; mod poll_evented; pub use self::poll_evented::PollEvented; @@ -51,6 +56,33 @@ pub struct Reactor { _wakeup_registration: mio::Registration, } +/// A handle to an event loop. +/// +/// A `Handle` is used for associating I/O objects with an event loop +/// explicitly. Typically though you won't end up using a `Handle` that often +/// and will instead use an implicitly configured handle for your thread. +#[derive(Clone)] +pub struct Handle { + inner: Weak, +} + +/// Return value from the `turn` method on `Reactor`. +/// +/// Currently this value doesn't actually provide any functionality, but it may +/// in the future give insight into what happened during `turn`. +#[derive(Debug)] +pub struct Turn { + _priv: (), +} + +/// Error returned from `Handle::set_fallback`. +#[derive(Clone, Debug)] +pub struct SetFallbackError(()); + +#[deprecated(since = "0.1.2", note = "use SetFallbackError instead")] +#[doc(hidden)] +pub type SetDefaultError = SetFallbackError; + struct Inner { /// The underlying system event queue. io: mio::Poll, @@ -62,16 +94,6 @@ struct Inner { wakeup: mio::SetReadiness } -/// A handle to an event loop. -/// -/// A `Handle` is used for associating I/O objects with an event loop -/// explicitly. Typically though you won't end up using a `Handle` that often -/// and will instead use and implicitly configured handle for your thread. -#[derive(Clone)] -pub struct Handle { - inner: Weak, -} - struct ScheduledIo { readiness: AtomicUsize, reader: AtomicTask, @@ -83,6 +105,12 @@ enum Direction { Write, } +/// The global fallback reactor. +static HANDLE_FALLBACK: AtomicUsize = ATOMIC_USIZE_INIT; + +/// Tracks the reactor for the current execution context. +thread_local!(static CURRENT_REACTOR: RefCell> = RefCell::new(None)); + const TOKEN_WAKEUP: mio::Token = mio::Token(0); const TOKEN_START: usize = 1; @@ -95,6 +123,45 @@ fn _assert_kinds() { _assert::(); } +// ===== impl Reactor ===== + +/// Set the default reactor for the duration of the closure +/// +/// # Panics +/// +/// This function panics if there already is a default reactor set. +pub(crate) fn with_default(handle: &Handle, enter: &mut Enter, f: F) -> R +where F: FnOnce(&mut Enter) -> R +{ + // Ensure that the executor is removed from the thread-local context + // when leaving the scope. This handles cases that involve panicking. + struct Reset; + + impl Drop for Reset { + fn drop(&mut self) { + CURRENT_REACTOR.with(|current| { + let mut current = current.borrow_mut(); + *current = None; + }); + } + } + + // This ensures the value for the current reactor gets reset even if there + // is a panic. + let _r = Reset; + + CURRENT_REACTOR.with(|current| { + { + let mut current = current.borrow_mut(); + assert!(current.is_none(), "default Tokio reactor already set \ + for execution context"); + *current = Some(handle.clone()); + } + + f(enter) + }) +} + impl Reactor { /// Creates a new event loop, returning any error that happened during the /// creation. @@ -118,7 +185,7 @@ impl Reactor { }) } - /// Returns a handle to this event loop which can be sent across threads + /// Returns a handle to this event loop which can be sent across threads /// and can be used as a proxy to the event loop itself. /// /// Handles are cloneable and clones always refer to the same event loop. @@ -153,7 +220,7 @@ impl Reactor { /// Additionally if the global reactor thread has already been initialized /// then this function will also return an error. (aka if `Handle::default` /// has been called previously in this program). - pub fn set_fallback(&self) -> Result<(), SetDefaultError> { + pub fn set_fallback(&self) -> Result<(), SetFallbackError> { set_fallback(self.handle()) } @@ -188,6 +255,18 @@ impl Reactor { Ok(Turn { _priv: () }) } + /// Returns true if the reactor is currently idle. + pub(crate) fn is_idle(&self) -> bool { + self.inner.io_dispatch + .read().unwrap() + .is_empty() + } + + /// Run the reactor in the background + pub(crate) fn background(self) -> io::Result { + Background::new(self) + } + fn poll(&mut self, max_wait: Option) -> io::Result<()> { // Block waiting for an event to happen, peeling out how many events // happened. @@ -244,13 +323,23 @@ impl Reactor { } } -/// Return value from the `turn` method on `Reactor`. -/// -/// Currently this value doesn't actually provide any functionality, but it may -/// in the future give insight into what happened during `turn`. -#[derive(Debug)] -pub struct Turn { - _priv: (), +impl Park for Reactor { + type Unpark = Handle; + type Error = io::Error; + + fn unpark(&self) -> Self::Unpark { + self.handle() + } + + fn park(&mut self) -> io::Result<()> { + self.turn(None)?; + Ok(()) + } + + fn park_timeout(&mut self, duration: Duration) -> io::Result<()> { + self.turn(Some(duration))?; + Ok(()) + } } impl fmt::Debug for Reactor { @@ -259,19 +348,133 @@ impl fmt::Debug for Reactor { } } -impl Drop for Inner { - fn drop(&mut self) { - // When a reactor is dropped it needs to wake up all blocked tasks as - // they'll never receive a notification, and all connected I/O objects - // will start returning errors pretty quickly. - let io = self.io_dispatch.read().unwrap(); - for (_, io) in io.iter() { - io.writer.notify(); - io.reader.notify(); +// ===== impl Handle ===== + +impl Handle { + /// Returns a handle to the current reactor. + pub fn current() -> Handle { + Handle::default() + } + + /// Returns a handle to the fallback reactor. + fn fallback() -> Handle { + let mut fallback = HANDLE_FALLBACK.load(SeqCst); + + // If the fallback hasn't been previously initialized then let's spin + // up a helper thread and try to initialize with that. If we can't + // actually create a helper thread then we'll just return a "defunct" + // handle which will return errors when I/O objects are attempted to be + // associated. + if fallback == 0 { + let reactor = match Reactor::new() { + Ok(reactor) => reactor, + Err(_) => return Handle { inner: Weak::new() }, + }; + + // If we successfully set ourselves as the actual fallback then we + // want to `forget` the helper thread to ensure that it persists + // globally. If we fail to set ourselves as the fallback that means + // that someone was racing with this call to `Handle::default`. + // They ended up winning so we'll destroy our helper thread (which + // shuts down the thread) and reload the fallback. + if set_fallback(reactor.handle().clone()).is_ok() { + let ret = reactor.handle().clone(); + + match reactor.background() { + Ok(bg) => bg.forget(), + // The global handle is fubar, but y'all probably got bigger + // problems if a thread can't spawn. + Err(_) => {} + } + + return ret + } + + fallback = HANDLE_FALLBACK.load(SeqCst); + } + + // At this point our fallback handle global was configured so we use + // its value to reify a handle, clone it, and then forget our reified + // handle as we don't actually have an owning reference to it. + assert!(fallback != 0); + + unsafe { + let handle = Handle::from_usize(fallback); + let ret = handle.clone(); + drop(handle.into_usize()); + return ret + } + } + + /// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise + /// makes the next call to `turn` return immediately. + /// + /// This method is intended to be used in situations where a notification + /// needs to otherwise be sent to the main reactor. If the reactor is + /// currently blocked inside of `turn` then it will wake up and soon return + /// after this method has been called. If the reactor is not currently + /// blocked in `turn`, then the next call to `turn` will not block and + /// return immediately. + fn wakeup(&self) { + if let Some(inner) = self.inner() { + inner.wakeup.set_readiness(mio::Ready::readable()).unwrap(); + } + } + + fn into_usize(self) -> usize { + unsafe { + mem::transmute::, usize>(self.inner) + } + } + + unsafe fn from_usize(val: usize) -> Handle { + let inner = mem::transmute::>(val);; + Handle { inner } + } + + fn inner(&self) -> Option> { + self.inner.upgrade() + } +} + +impl Unpark for Handle { + fn unpark(&self) { + self.wakeup(); + } +} + +impl Default for Handle { + fn default() -> Handle { + CURRENT_REACTOR.with(|current| { + match *current.borrow() { + Some(ref handle) => handle.clone(), + None => Handle::fallback(), + } + }) + } +} + +impl fmt::Debug for Handle { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "Handle") + } +} + +fn set_fallback(handle: Handle) -> Result<(), SetFallbackError> { + unsafe { + let val = handle.into_usize(); + match HANDLE_FALLBACK.compare_exchange(0, val, SeqCst, SeqCst) { + Ok(_) => Ok(()), + Err(_) => { + drop(Handle::from_usize(val)); + Err(SetFallbackError(())) + } } } } +// ===== impl Inner ===== + impl Inner { /// Register an I/O resource with the reactor. /// @@ -330,104 +533,20 @@ impl Inner { } } -static HANDLE_FALLBACK: AtomicUsize = ATOMIC_USIZE_INIT; - -/// Error returned from `Handle::set_fallback`. -#[derive(Clone, Debug)] -pub struct SetDefaultError(()); - -impl Handle { - /// Forces a reactor blocked in a call to `turn` to wakeup, or otherwise - /// makes the next call to `turn` return immediately. - /// - /// This method is intended to be used in situations where a notification - /// needs to otherwise be sent to the main reactor. If the reactor is - /// currently blocked inside of `turn` then it will wake up and soon return - /// after this method has been called. If the reactor is not currently - /// blocked in `turn`, then the next call to `turn` will not block and - /// return immediately. - fn wakeup(&self) { - if let Some(inner) = self.inner() { - inner.wakeup.set_readiness(mio::Ready::readable()).unwrap(); - } - } - - fn into_usize(self) -> usize { - unsafe { - mem::transmute::, usize>(self.inner) - } - } - - unsafe fn from_usize(val: usize) -> Handle { - let inner = mem::transmute::>(val);; - Handle { inner } - } - - fn inner(&self) -> Option> { - self.inner.upgrade() - } -} - -impl Default for Handle { - fn default() -> Handle { - let mut fallback = HANDLE_FALLBACK.load(SeqCst); - - // If the fallback hasn't been previously initialized then let's spin - // up a helper thread and try to initialize with that. If we can't - // actually create a helper thread then we'll just return a "defunkt" - // handle which will return errors when I/O objects are attempted to be - // associated. - if fallback == 0 { - let helper = match global::HelperThread::new() { - Ok(helper) => helper, - Err(_) => return Handle { inner: Weak::new() }, - }; - - // If we successfully set ourselves as the actual fallback then we - // want to `forget` the helper thread to ensure that it persists - // globally. If we fail to set ourselves as the fallback that means - // that someone was racing with this call to `Handle::default`. - // They ended up winning so we'll destroy our helper thread (which - // shuts down the thread) and reload the fallback. - if set_fallback(helper.handle().clone()).is_ok() { - let ret = helper.handle().clone(); - helper.forget(); - return ret - } - fallback = HANDLE_FALLBACK.load(SeqCst); - } - - // At this point our fallback handle global was configured so we use - // its value to reify a handle, clone it, and then forget our reified - // handle as we don't actually have an owning reference to it. - assert!(fallback != 0); - unsafe { - let handle = Handle::from_usize(fallback); - let ret = handle.clone(); - drop(handle.into_usize()); - return ret +impl Drop for Inner { + fn drop(&mut self) { + // When a reactor is dropped it needs to wake up all blocked tasks as + // they'll never receive a notification, and all connected I/O objects + // will start returning errors pretty quickly. + let io = self.io_dispatch.read().unwrap(); + for (_, io) in io.iter() { + io.writer.notify(); + io.reader.notify(); } } } -impl fmt::Debug for Handle { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "Handle") - } -} - -fn set_fallback(handle: Handle) -> Result<(), SetDefaultError> { - unsafe { - let val = handle.into_usize(); - match HANDLE_FALLBACK.compare_exchange(0, val, SeqCst, SeqCst) { - Ok(_) => Ok(()), - Err(_) => { - drop(Handle::from_usize(val)); - Err(SetDefaultError(())) - } - } - } -} +// ===== misc ===== fn read_ready() -> mio::Ready { mio::Ready::readable() | platform::hup() diff --git a/src/runtime.rs b/src/runtime.rs new file mode 100644 index 000000000..700a1222a --- /dev/null +++ b/src/runtime.rs @@ -0,0 +1,362 @@ +//! A batteries included runtime for applications using Tokio. +//! +//! Applications using Tokio require some runtime support in order to work: +//! +//! * A [reactor] to drive I/O resources. +//! * An [executor] to execute tasks that use these I/O resources. +//! +//! While it is possible to setup each component manually, this involves a bunch +//! of boilerplate. +//! +//! [`Runtime`] bundles all of these various runtime components into a single +//! handle that can be started and shutdown together, eliminating the necessary +//! boilerplate to run a Tokio application. +//! +//! Most applications wont need to use [`Runtime`] directly. Instead, they will +//! use the [`run`] function, which uses [`Runtime`] under the hood. +//! +//! Creating a [`Runtime`] does the following: +//! +//! * Spawn a background thread running a [`Reactor`] instance. +//! * Start a [`ThreadPool`] for executing futures. +//! +//! The thread pool uses a work-stealing strategy and is configured to start a +//! worker thread for each CPU core available on the system. This tends to be +//! the ideal setup for Tokio applications. +//! +//! # Usage +//! +//! Most applications will use the [`run`] function. This takes a future to +//! "seed" the application, blocking the thread until the runtime becomes +//! [idle]. +//! +//! ```rust +//! # extern crate tokio; +//! # extern crate futures; +//! # use futures::{Future, Stream}; +//! use tokio::net::TcpListener; +//! +//! # fn process(_: T) -> Box + Send> { +//! # unimplemented!(); +//! # } +//! # fn dox() { +//! # let addr = "127.0.0.1:8080".parse().unwrap(); +//! let listener = TcpListener::bind(&addr).unwrap(); +//! +//! let server = listener.incoming() +//! .map_err(|e| println!("error = {:?}", e)) +//! .for_each(|socket| { +//! tokio::spawn(process(socket)) +//! }); +//! +//! tokio::run(server); +//! # } +//! # pub fn main() {} +//! ``` +//! +//! In this function, the `run` function blocks until the runtime becomes idle. +//! See [`shutdown_on_idle`][idle] for more shutdown details. +//! +//! From within the context of the runtime, additional tasks are spawned using +//! the [`tokio::spawn`] function. Futures spawned using this function will be +//! executed on the same thread pool used by the [`Runtime`]. +//! +//! A [`Runtime`] instance can also be used directly. +//! +//! ```rust +//! # extern crate tokio; +//! # extern crate futures; +//! # use futures::{Future, Stream}; +//! use tokio::runtime::Runtime; +//! use tokio::net::TcpListener; +//! +//! # fn process(_: T) -> Box + Send> { +//! # unimplemented!(); +//! # } +//! # fn dox() { +//! # let addr = "127.0.0.1:8080".parse().unwrap(); +//! let listener = TcpListener::bind(&addr).unwrap(); +//! +//! let server = listener.incoming() +//! .map_err(|e| println!("error = {:?}", e)) +//! .for_each(|socket| { +//! tokio::spawn(process(socket)) +//! }); +//! +//! // Create the runtime +//! let mut rt = Runtime::new().unwrap(); +//! +//! // Spawn the server task +//! rt.spawn(server); +//! +//! // Wait until the runtime becomes idle and shut it down. +//! rt.shutdown_on_idle() +//! .wait().unwrap(); +//! # } +//! # pub fn main() {} +//! ``` +//! +//! [reactor]: ../reactor/struct.Reactor.html +//! [executor]: https://tokio.rs/docs/getting-started/runtime-model/#executors +//! [`Runtime`]: struct.Runtime.html +//! [`ThreadPool`]: ../executor/thread_pool/struct.ThreadPool.html +//! [`run`]: fn.run.html +//! [idle]: struct.Runtime.html#method.shutdown_on_idle +//! [`tokio::spawn`]: ../executor/fn.spawn.html + +use reactor::{self, Reactor, Handle}; +use reactor::background::Background; + +use tokio_threadpool::{self as threadpool, ThreadPool}; +use futures::Poll; +use futures::future::Future; + +use std::{fmt, io}; + +/// Handle to the Tokio runtime. +/// +/// The Tokio runtime includes a reactor as well as an executor for running +/// tasks. +/// +/// See [module level][mod] documentation for more details. +/// +/// [mod]: index.html +#[derive(Debug)] +pub struct Runtime { + inner: Option, +} + +/// A future that resolves when the Tokio `Runtime` is shut down. +pub struct Shutdown { + inner: Box + Send>, +} + +#[derive(Debug)] +struct Inner { + /// Reactor running on a background thread. + reactor: Background, + + /// Task execution pool. + pool: ThreadPool, +} + +// ===== impl Runtime ===== + +/// Start the Tokio runtime using the supplied future to bootstrap execution. +/// +/// This function is used to bootstrap the execution of a Tokio application. It +/// does the following: +/// +/// * Start the Tokio runtime using a default configuration. +/// * Spawn the given future onto the thread pool. +/// * Block the çurrent thread until the runtime shuts down. +/// +/// Note that the function will not return immediately once `future` has +/// completed. Instead it waits for the entire runtime to become idle. +/// +/// See [module level][mod] documentation for more details. +/// +/// # Examples +/// +/// ```rust +/// # extern crate tokio; +/// # extern crate futures; +/// # use futures::{Future, Stream}; +/// use tokio::net::TcpListener; +/// +/// # fn process(_: T) -> Box + Send> { +/// # unimplemented!(); +/// # } +/// # fn dox() { +/// # let addr = "127.0.0.1:8080".parse().unwrap(); +/// let listener = TcpListener::bind(&addr).unwrap(); +/// +/// let server = listener.incoming() +/// .map_err(|e| println!("error = {:?}", e)) +/// .for_each(|socket| { +/// tokio::spawn(process(socket)) +/// }); +/// +/// tokio::run(server); +/// # } +/// # pub fn main() {} +/// ``` +/// +/// # Panics +/// +/// This function panics if called from the context of an executor. +/// +/// [mod]: ../index.html +pub fn run(future: F) +where F: Future + Send + 'static, +{ + let mut runtime = Runtime::new().unwrap(); + runtime.spawn(future); + runtime.shutdown_on_idle().wait().unwrap(); +} + +impl Runtime { + /// Create a new runtime instance with default configuration values. + /// + /// See [module level][mod] documentation for more details. + /// + /// [mod]: index.html + pub fn new() -> io::Result { + // Spawn a reactor on a background thread. + let reactor = Reactor::new()?.background()?; + + // Get a handle to the reactor. + let handle = reactor.handle().clone(); + + let pool = threadpool::Builder::new() + .around_worker(move |w, enter| { + reactor::with_default(&handle, enter, |_| { + w.run(); + }); + }) + .build(); + + Ok(Runtime { + inner: Some(Inner { + reactor, + pool, + }), + }) + } + + /// Return a reference to the reactor handle for this runtime instance. + pub fn handle(&self) -> &Handle { + self.inner.as_ref().unwrap().reactor.handle() + } + + /// Spawn a future onto the Tokio runtime. + /// + /// This spawns the given future onto the runtime's executor, usually a + /// thread pool. The thread pool is then responsible for polling the future + /// until it completes. + /// + /// See [module level][mod] documentation for more details. + /// + /// [mod]: index.html + /// + /// # Examples + /// + /// ```rust + /// # extern crate tokio; + /// # extern crate futures; + /// # use futures::{future, Future, Stream}; + /// use tokio::runtime::Runtime; + /// + /// # fn dox() { + /// // Create the runtime + /// let mut rt = Runtime::new().unwrap(); + /// + /// // Spawn a future onto the runtime + /// rt.spawn(future::lazy(|| { + /// println!("now running on a worker thread"); + /// Ok(()) + /// })); + /// # } + /// # pub fn main() {} + /// ``` + /// + /// # Panics + /// + /// This function panics if the spawn fails. Failure occurs if the executor + /// is currently at capacity and is unable to spawn a new future. + pub fn spawn(&mut self, future: F) -> &mut Self + where F: Future + Send + 'static, + { + self.inner_mut().pool.sender().spawn(future).unwrap(); + self + } + + /// Signals the runtime to shutdown once it becomes idle. + /// + /// Returns a future that completes once the shutdown operation has + /// completed. + /// + /// This function can be used to perform a graceful shutdown of the runtime. + /// + /// The runtime enters an idle state once **all** of the following occur. + /// + /// * The thread pool has no tasks to execute, i.e., all tasks that were + /// spawned have completed. + /// * The reactor is not managing any I/O resources. + /// + /// See [module level][mod] documentation for more details. + /// + /// [mod]: index.html + pub fn shutdown_on_idle(mut self) -> Shutdown { + let inner = self.inner.take().unwrap(); + + let inner = Box::new({ + let pool = inner.pool; + let reactor = inner.reactor; + + pool.shutdown_on_idle().and_then(|_| { + reactor.shutdown_on_idle() + }) + }); + + Shutdown { inner } + } + + /// Signals the runtime to shutdown immediately. + /// + /// Returns a future that completes once the shutdown operation has + /// completed. + /// + /// This function will forcibly shutdown the runtime, causing any + /// in-progress work to become canceled. The shutdown steps are: + /// + /// * Drain any scheduled work queues. + /// * Drop any futures that have not yet completed. + /// * Drop the reactor. + /// + /// Once the reactor has dropped, any outstanding I/O resources bound to + /// that reactor will no longer function. Calling any method on them will + /// result in an error. + /// + /// See [module level][mod] documentation for more details. + /// + /// [mod]: index.html + pub fn shutdown_now(mut self) -> Shutdown { + let inner = self.inner.take().unwrap(); + + let inner = Box::new({ + let pool = inner.pool; + let reactor = inner.reactor; + + pool.shutdown_now().and_then(|_| { + reactor.shutdown_now() + }) + }); + + Shutdown { inner } + } + + fn inner_mut(&mut self) -> &mut Inner { + self.inner.as_mut().unwrap() + } +} + +// ===== impl Shutdown ===== + +impl Future for Shutdown { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + try_ready!(self.inner.poll()); + Ok(().into()) + } +} + +impl fmt::Debug for Shutdown { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + fmt.debug_struct("Shutdown") + .field("inner", &"Box>") + .finish() + } +} diff --git a/tests/current_thread.rs b/tests/current_thread.rs new file mode 100644 index 000000000..926eee63a --- /dev/null +++ b/tests/current_thread.rs @@ -0,0 +1,262 @@ +extern crate tokio; +extern crate tokio_executor; +extern crate futures; + +use tokio::executor::current_thread::{self, block_on_all, CurrentThread}; + +use std::cell::{Cell, RefCell}; +use std::rc::Rc; +use std::thread; +use std::time::Duration; + +use futures::task; +use futures::future::{self, lazy}; +use futures::prelude::*; +use futures::sync::oneshot; + +#[test] +fn spawn_from_block_on_all() { + let cnt = Rc::new(Cell::new(0)); + let c = cnt.clone(); + + let msg = current_thread::block_on_all(lazy(move || { + c.set(1 + c.get()); + + // Spawn! + current_thread::spawn(lazy(move || { + c.set(1 + c.get()); + Ok::<(), ()>(()) + })); + + Ok::<_, ()>("hello") + })).unwrap(); + + assert_eq!(2, cnt.get()); + assert_eq!(msg, "hello"); +} + +#[test] +fn block_waits() { + let cnt = Rc::new(Cell::new(0)); + let cnt2 = cnt.clone(); + + let (tx, rx) = oneshot::channel(); + + thread::spawn(|| { + thread::sleep(Duration::from_millis(1000)); + tx.send(()).unwrap(); + }); + + block_on_all(rx.then(move |_| { + cnt.set(1 + cnt.get()); + Ok::<_, ()>(()) + })).unwrap(); + + assert_eq!(1, cnt2.get()); +} + +#[test] +fn spawn_many() { + const ITER: usize = 200; + + let cnt = Rc::new(Cell::new(0)); + let mut current_thread = CurrentThread::new(); + + for _ in 0..ITER { + let cnt = cnt.clone(); + current_thread.spawn(lazy(move || { + cnt.set(1 + cnt.get()); + Ok::<(), ()>(()) + })); + } + + current_thread.run().unwrap(); + + assert_eq!(cnt.get(), ITER); +} + +#[test] +fn does_not_set_global_executor_by_default() { + use tokio_executor::Executor; + + block_on_all(lazy(|| { + tokio_executor::DefaultExecutor::current() + .spawn(Box::new(lazy(|| ok()))) + .unwrap_err(); + + ok() + })).unwrap(); +} + +#[test] +fn spawn_from_block_on_future() { + let cnt = Rc::new(Cell::new(0)); + + let mut current_thread = CurrentThread::new(); + + current_thread.block_on(lazy(|| { + let cnt = cnt.clone(); + + current_thread::spawn(lazy(move || { + cnt.set(1 + cnt.get()); + Ok(()) + })); + + Ok::<_, ()>(()) + })).unwrap(); + + current_thread.run().unwrap(); + + assert_eq!(1, cnt.get()); +} + +struct Never(Rc<()>); + +impl Future for Never { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + Ok(Async::NotReady) + } +} + +#[test] +fn outstanding_tasks_are_dropped_when_executor_is_dropped() { + let mut rc = Rc::new(()); + + let mut current_thread = CurrentThread::new(); + current_thread.spawn(Never(rc.clone())); + + drop(current_thread); + + // Ensure the daemon is dropped + assert!(Rc::get_mut(&mut rc).is_some()); + + // Using the global spawn fn + + let mut rc = Rc::new(()); + + let mut current_thread = CurrentThread::new(); + + current_thread.block_on(lazy(|| { + current_thread::spawn(Never(rc.clone())); + Ok::<_, ()>(()) + })).unwrap(); + + drop(current_thread); + + // Ensure the daemon is dropped + assert!(Rc::get_mut(&mut rc).is_some()); +} + +#[test] +#[should_panic] +fn nesting_run() { + block_on_all(lazy(|| { + block_on_all(lazy(|| { + ok() + })).unwrap(); + + ok() + })).unwrap(); +} + +#[test] +#[should_panic] +fn run_in_future() { + block_on_all(lazy(|| { + current_thread::spawn(lazy(|| { + block_on_all(lazy(|| { + ok() + })).unwrap(); + ok() + })); + ok() + })).unwrap(); +} + +#[test] +fn tick_on_infini_future() { + let num = Rc::new(Cell::new(0)); + + struct Infini { + num: Rc>, + } + + impl Future for Infini { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + self.num.set(1 + self.num.get()); + task::current().notify(); + Ok(Async::NotReady) + } + } + + CurrentThread::new() + .spawn(Infini { + num: num.clone(), + }) + .turn(None) + .unwrap(); + + assert_eq!(1, num.get()); +} + +#[test] +fn tasks_are_scheduled_fairly() { + let state = Rc::new(RefCell::new([0, 0])); + + struct Spin { + state: Rc>, + idx: usize, + } + + impl Future for Spin { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + let mut state = self.state.borrow_mut(); + + if self.idx == 0 { + let diff = state[0] - state[1]; + + assert!(diff.abs() <= 1); + + if state[0] >= 50 { + return Ok(().into()); + } + } + + state[self.idx] += 1; + + if state[self.idx] >= 100 { + return Ok(().into()); + } + + task::current().notify(); + Ok(Async::NotReady) + } + } + + block_on_all(lazy(|| { + current_thread::spawn(Spin { + state: state.clone(), + idx: 0, + }); + + current_thread::spawn(Spin { + state: state, + idx: 1, + }); + + ok() + })).unwrap(); +} + +fn ok() -> future::FutureResult<(), ()> { + future::ok(()) +} diff --git a/tests/global.rs b/tests/global.rs index bf5682fa0..a863176b1 100644 --- a/tests/global.rs +++ b/tests/global.rs @@ -1,5 +1,6 @@ extern crate futures; extern crate tokio; +extern crate env_logger; use std::thread; @@ -15,6 +16,8 @@ macro_rules! t { #[test] fn hammer() { + let _ = env_logger::init(); + let threads = (0..10).map(|_| { thread::spawn(|| { let srv = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap())); diff --git a/tests/runtime.rs b/tests/runtime.rs new file mode 100644 index 000000000..ddda48607 --- /dev/null +++ b/tests/runtime.rs @@ -0,0 +1,52 @@ +extern crate futures; +extern crate tokio; +extern crate tokio_io; +extern crate env_logger; + +use futures::prelude::*; +use tokio::net::{TcpStream, TcpListener}; +use tokio_io::io; + +macro_rules! t { + ($e:expr) => (match $e { + Ok(e) => e, + Err(e) => panic!("{} failed with {:?}", stringify!($e), e), + }) +} + +#[test] +fn basic_runtime_usage() { + let _ = env_logger::init(); + + // TODO: Don't require the lazy wrapper + tokio::run(::futures::future::lazy(|| { + let server = t!(TcpListener::bind(&"127.0.0.1:0".parse().unwrap())); + let addr = t!(server.local_addr()); + let client = TcpStream::connect(&addr); + + let server = server.incoming().take(1) + .map_err(|e| println!("accept err = {:?}", e)) + .for_each(|socket| { + tokio::spawn({ + io::write_all(socket, b"hello") + .map(|_| println!("write done")) + .map_err(|e| println!("write err = {:?}", e)) + }) + }) + .map(|_| println!("accept done")); + + let client = client + .map_err(|e| println!("connect err = {:?}", e)) + .and_then(|client| { + // Read all + io::read_to_end(client, vec![]) + .map(|_| println!("read done")) + .map_err(|e| println!("read err = {:?}", e)) + }); + + tokio::spawn({ + server.join(client) + .map(|_| println!("done")) + }) + })); +} diff --git a/tokio-executor/Cargo.toml b/tokio-executor/Cargo.toml new file mode 100644 index 000000000..1d1fc7358 --- /dev/null +++ b/tokio-executor/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "tokio-executor" +version = "0.1.0" +documentation = "https://docs.rs/tokio-executor" +repository = "https://github.com/tokio-rs/tokio" +homepage = "https://github.com/tokio-rs/tokio" +license = "MIT/Apache-2.0" +authors = ["Carl Lerche "] +description = """ +Future execution primitives +""" +keywords = ["futures", "tokio"] +categories = ["concurrency", "asynchronous"] + +[dependencies] +futures = "0.1" diff --git a/tokio-executor/src/enter.rs b/tokio-executor/src/enter.rs new file mode 100644 index 000000000..3c9f3cfa1 --- /dev/null +++ b/tokio-executor/src/enter.rs @@ -0,0 +1,97 @@ +use std::prelude::v1::*; +use std::cell::Cell; +use std::fmt; + +thread_local!(static ENTERED: Cell = Cell::new(false)); + +/// Represents an executor context. +/// +/// For more details, see [`enter` documentation](fn.enter.html) +pub struct Enter { + on_exit: Vec>, + permanent: bool, +} + +/// An error returned by `enter` if an execution scope has already been +/// entered. +#[derive(Debug)] +pub struct EnterError { + _a: (), +} + +/// Marks the current thread as being within the dynamic extent of an +/// executor. +/// +/// Executor implementations should call this function before blocking the +/// thread. If `None` is returned, the executor should fail by panicking or +/// taking some other action without blocking the current thread. This prevents +/// deadlocks due to multiple executors competing for the same thread. +/// +/// # Error +/// +/// Returns an error if the current thread is already marked +pub fn enter() -> Result { + ENTERED.with(|c| { + if c.get() { + Err(EnterError { _a: () }) + } else { + c.set(true); + + Ok(Enter { + on_exit: Vec::new(), + permanent: false, + }) + } + }) +} + +impl Enter { + /// Register a callback to be invoked if and when the thread + /// ceased to act as an executor. + pub fn on_exit(&mut self, f: F) where F: FnOnce() + 'static { + self.on_exit.push(Box::new(f)); + } + + /// Treat the remainder of execution on this thread as part of an + /// executor; used mostly for thread pool worker threads. + /// + /// All registered `on_exit` callbacks are *dropped* without being + /// invoked. + pub fn make_permanent(mut self) { + self.permanent = true; + } +} + +impl fmt::Debug for Enter { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("Enter").finish() + } +} + +impl Drop for Enter { + fn drop(&mut self) { + ENTERED.with(|c| { + assert!(c.get()); + + if self.permanent { + return + } + + for callback in self.on_exit.drain(..) { + callback.call(); + } + + c.set(false); + }); + } +} + +trait Callback: 'static { + fn call(self: Box); +} + +impl Callback for F { + fn call(self: Box) { + (*self)() + } +} diff --git a/tokio-executor/src/global.rs b/tokio-executor/src/global.rs new file mode 100644 index 000000000..5c1435e11 --- /dev/null +++ b/tokio-executor/src/global.rs @@ -0,0 +1,154 @@ +use super::{Executor, Enter, SpawnError}; + +use futures::Future; + +use std::cell::Cell; +use std::marker::PhantomData; +use std::rc::Rc; + +/// Executes futures on the default executor for the current execution context. +/// +/// `DefaultExecutor` implements `Executor` and can be used to spawn futures +/// without referencing a specific executor. +/// +/// When an executor starts, it sets the `DefaultExecutor` handle to point to an +/// executor (usually itself) that is used to spawn new tasks. +/// +/// The current `DefaultExecutor` reference is tracked using a thread-local +/// variable and is set using `tokio_executor::with_default` +#[derive(Debug, Clone)] +pub struct DefaultExecutor { + // Prevent the handle from moving across threads. + _p: PhantomData>, +} + +impl DefaultExecutor { + /// Returns a handle to the default executor for the current context. + /// + /// Futures may be spawned onto the default executor using this handle. + /// + /// The returned handle will reference whichever executor is configured as + /// the default **at the time `spawn` is called`. This enables + /// `DefaultExecutor::current()` to be called before an execution context is + /// setup, then passed **into** an execution context before it is used. + pub fn current() -> DefaultExecutor { + DefaultExecutor { + _p: PhantomData, + } + } +} + +/// Thread-local tracking the current executor +thread_local!(static EXECUTOR: Cell> = Cell::new(None)); + +// ===== impl DefaultExecutor ===== + +impl super::Executor for DefaultExecutor { + fn spawn(&mut self, future: Box + Send>) + -> Result<(), SpawnError> + { + EXECUTOR.with(|current_executor| { + match current_executor.get() { + Some(executor) => { + let executor = unsafe { &mut *executor }; + executor.spawn(future) + } + None => { + Err(SpawnError::shutdown()) + } + } + }) + } +} + +// ===== global spawn fns ===== + +/// Submits a future for execution on the default executor -- usually a +/// threadpool. +/// +/// Futures are lazy constructs. When they are defined, no work happens. In +/// order for the logic defined by the future to be run, the future must be +/// spawned on an executor. This function is the easiest way to do so. +/// +/// This function must be called from an execution context, i.e. from a future +/// that has been already spawned onto an executor. +/// +/// Once spawned, the future will execute. The details of how that happens is +/// left up to the executor instance. If the executor is a thread pool, the +/// future will be pushed onto a queue that a worker thread polls from. If the +/// executor is a "current thread" executor, the future might be polled +/// immediately from within the call to `spawn` or it might be pushed onto an +/// internal queue. +/// +/// # Panics +/// +/// This function will panic if the default executor is not set or if spawning +/// onto the default executor returns an error. To avoid the panic, use the +/// `DefaultExecutor` handle directly. +/// +/// # Examples +/// +/// ```rust +/// # extern crate futures; +/// # extern crate tokio_executor; +/// # use tokio_executor::spawn; +/// # pub fn dox() { +/// use futures::future::lazy; +/// +/// spawn(lazy(|| { +/// println!("running on the default executor"); +/// Ok(()) +/// })); +/// # } +/// # pub fn main() {} +/// ``` +pub fn spawn(future: T) + where T: Future + Send + 'static, +{ + DefaultExecutor::current().spawn(Box::new(future)) + .unwrap() +} + +/// Set the default executor for the duration of the closure +/// +/// # Panics +/// +/// This function panics if there already is a default executor set. +pub fn with_default(executor: &mut T, enter: &mut Enter, f: F) -> R +where T: Executor, + F: FnOnce(&mut Enter) -> R +{ + EXECUTOR.with(|cell| { + assert!(cell.get().is_none(), "default executor already set for execution context"); + + // Ensure that the executor is removed from the thread-local context + // when leaving the scope. This handles cases that involve panicking. + struct Reset<'a>(&'a Cell>); + + impl<'a> Drop for Reset<'a> { + fn drop(&mut self) { + self.0.set(None); + } + } + + let _reset = Reset(cell); + + // While scary, this is safe. The function takes a + // `&mut Executor`, which guarantees that the reference lives for the + // duration of `with_default`. + // + // Because we are always clearing the TLS value at the end of the + // function, we can cast the reference to 'static which thread-local + // cells require. + let executor = unsafe { hide_lt(executor as &mut _ as *mut _) }; + + cell.set(Some(executor)); + + f(enter) + }) +} + +unsafe fn hide_lt<'a>(p: *mut (Executor + 'a)) -> *mut (Executor + 'static) { + use std::mem; + mem::transmute(p) +} diff --git a/tokio-executor/src/lib.rs b/tokio-executor/src/lib.rs new file mode 100644 index 000000000..299f197d5 --- /dev/null +++ b/tokio-executor/src/lib.rs @@ -0,0 +1,189 @@ +//! Task execution utilities. +//! +//! In the Tokio execution model, futures are lazy. When a future is created, no +//! work is performed. In order for the work defined by the future to happen, +//! the future must be submitted to an executor. A future that is submitted to +//! an executor is called a "task". +//! +//! The executor executor is responsible for ensuring that [`Future::poll`] is +//! called whenever the task is [notified]. Notification happens when the +//! internal state of a task transitions from "not ready" to ready. For +//! example, a socket might have received data and a call to `read` will now be +//! able to succeed. + +#![deny(missing_docs, missing_debug_implementations, warnings)] +#![doc(html_root_url = "https://docs.rs/tokio-executor/0.1")] + +extern crate futures; + +mod enter; +mod global; +pub mod park; + +pub use enter::{enter, Enter, EnterError}; +pub use global::{spawn, with_default, DefaultExecutor}; + +use futures::Future; + +/// A value that executes futures. +/// +/// The [`spawn`] function is used to submit a future to an executor. Once +/// submitted, the executor takes ownership of the future and becomes +/// responsible for driving the future to completion. +/// +/// The strategy employed by the executor to handle the future is less defined +/// and is left up to the `Executor` implementation. The `Executor` instance is +/// expected to call [`poll`] on the future once it has been notified, however +/// the "when" and "how" can vary greatly. +/// +/// For example, the executor might be a thread pool, in which case a set of +/// threads have already been spawned up and the future is inserted into a +/// queue. A thread will acquire the future and poll it. +/// +/// The `Executor` trait is only for futures that **are** `Send`. These are most +/// common. There currently is no trait that describes executors that operate +/// entirely on the current thread (i.e., are able to spawn futures that are not +/// `Send`). Note that single threaded executors can still implement `Executor`, +/// but only futures that are `Send` can be spawned via the trait. +/// +/// # Errors +/// +/// The [`spawn`] function returns `Result` with an error type of `SpawnError`. +/// This error type represents the reason that the executor was unable to spawn +/// the future. The two current represented scenarios are: +/// +/// * An executor being at capacity or full. As such, the executor is not able +/// to accept a new future. This error state is expected to be transient. +/// * An executor has been shutdown and can no longer accept new futures. This +/// error state is expected to be permanent. +/// +/// If a caller encounters an at capacity error, the caller should try to shed +/// load. This can be as simple as dropping the future that was spawned. +/// +/// If the caller encounters a shutdown error, the caller should attempt to +/// gracefully shutdown. +/// +/// # Examples +/// +/// ```rust +/// # extern crate futures; +/// # extern crate tokio_executor; +/// # use tokio_executor::Executor; +/// # fn docs(my_executor: &mut Executor) { +/// use futures::future::lazy; +/// my_executor.spawn(Box::new(lazy(|| { +/// println!("running on the executor"); +/// Ok(()) +/// }))).unwrap(); +/// # } +/// # fn main() {} +/// ``` +/// +/// [`spawn`]: #tymethod.spawn +/// [`poll`]: https://docs.rs/futures/0.1/futures/future/trait.Future.html#tymethod.poll +pub trait Executor { + /// Spawns a future object to run on this executor. + /// + /// `future` is passed to the executor, which will begin running it. The + /// future may run on the current thread or another thread at the discretion + /// of the `Executor` implementation. + /// + /// # Panics + /// + /// Implementors are encouraged to avoid panics. However, a panic is + /// permitted and the caller should check the implementation specific + /// documentation for more details on possible panics. + /// + /// # Examples + /// + /// ```rust + /// # extern crate futures; + /// # extern crate tokio_executor; + /// # use tokio_executor::Executor; + /// # fn docs(my_executor: &mut Executor) { + /// use futures::future::lazy; + /// my_executor.spawn(Box::new(lazy(|| { + /// println!("running on the executor"); + /// Ok(()) + /// }))).unwrap(); + /// # } + /// # fn main() {} + /// ``` + fn spawn(&mut self, future: Box + Send>) + -> Result<(), SpawnError>; + + /// Provides a best effort **hint** to whether or not `spawn` will succeed. + /// + /// This function may return both false positives **and** false negatives. + /// If `status` returns `Ok`, then a call to `spawn` will *probably* + /// succeed, but may fail. If `status` returns `Err`, a call to `spawn` will + /// *probably* fail, but may succeed. + /// + /// This allows a caller to avoid creating the task if the call to `spawn` + /// has a high likelihood of failing. + /// + /// # Panics + /// + /// This function must not panic. Implementors must ensure that panics do + /// not happen. + /// + /// # Examples + /// + /// ```rust + /// # extern crate futures; + /// # extern crate tokio_executor; + /// # use tokio_executor::Executor; + /// # fn docs(my_executor: &mut Executor) { + /// use futures::future::lazy; + /// + /// if my_executor.status().is_ok() { + /// my_executor.spawn(Box::new(lazy(|| { + /// println!("running on the executor"); + /// Ok(()) + /// }))).unwrap(); + /// } else { + /// println!("the executor is not in a good state"); + /// } + /// # } + /// # fn main() {} + /// ``` + fn status(&self) -> Result<(), SpawnError> { + Ok(()) + } +} + +/// Errors returned by `Executor::spawn`. +/// +/// Spawn errors should represent relatively rare scenarios. Currently, the two +/// scenarios represented by `SpawnError` are: +/// +/// * An executor being at capacity or full. As such, the executor is not able +/// to accept a new future. This error state is expected to be transient. +/// * An executor has been shutdown and can no longer accept new futures. This +/// error state is expected to be permanent. +#[derive(Debug)] +pub struct SpawnError { + is_shutdown: bool, +} + +impl SpawnError { + /// Return a new `SpawnError` reflecting a shutdown executor failure. + pub fn shutdown() -> Self { + SpawnError { is_shutdown: true } + } + + /// Return a new `SpawnError` reflecting an executor at capacity failure. + pub fn at_capacity() -> Self { + SpawnError { is_shutdown: false } + } + + /// Returns `true` if the error reflects a shutdown executor failure. + pub fn is_shutdown(&self) -> bool { + self.is_shutdown + } + + /// Returns `true` if the error reflects an executor at capacity failure. + pub fn is_at_capacity(&self) -> bool { + !self.is_shutdown + } +} diff --git a/tokio-executor/src/park.rs b/tokio-executor/src/park.rs new file mode 100644 index 000000000..7a60900f9 --- /dev/null +++ b/tokio-executor/src/park.rs @@ -0,0 +1,294 @@ +//! Abstraction over blocking and unblocking the current thread. +//! +//! Provides an abstraction over blocking the current thread. This is similar to +//! the park / unpark constructs provided by [`std`] but made generic. This +//! allows embedding custom functionality to perform when the thread is blocked. +//! +//! A blocked [`Park`][p] instance is unblocked by calling [`unpark`] on its +//! [`Unpark`][up] handle. +//! +//! The [`ParkThread`] struct implements [`Park`][p] using +//! [`thread::park`][`std`] to put the thread to sleep. The Tokio reactor also +//! implements park, but uses [`mio::Poll`][mio] to block the thread instead. +//! +//! The [`Park`][p] trait is composable. A timer implementation might decorate a +//! [`Park`][p] implementation by checking if any timeouts have elapsed after +//! the inner [`Park`][p] implementation unblocks. +//! +//! # Model +//! +//! Conceptually, each [`Park`][p] instance has an associated token, which is +//! initially not present: +//! +//! * The [`park`] method blocks the current thread unless or until the token +//! is available, at which point it atomically consumes the token. +//! * The [`unpark`] method atomically makes the token available if it wasn't +//! already. +//! +//! Some things to note: +//! +//! * If [`unpark`] is called before [`park`], the next call to [`park`] will +//! **not** block the thread. +//! * **Spurious** wakeups are permited, i.e., the [`park`] method may unblock +//! even if [`unpark`] was not called. +//! * [`park_timeout`] does the same as [`park`] but allows specifying a maximum +//! time to block the thread for. +//! +//! [`std`]: https://doc.rust-lang.org/std/thread/fn.park.html +//! [`thread::park`]: https://doc.rust-lang.org/std/thread/fn.park.html +//! [`ParkThread`]: struct.ParkThread.html +//! [p]: trait.Park.html +//! [`park`]: trait.Park.html#tymethod.park +//! [`park_timeout`]: trait.Park.html#tymethod.park_timeout +//! [`unpark`]: trait.Unpark.html#tymethod.unpark +//! [up]: trait.Unpark.html +//! [mio]: https://docs.rs/mio/0.6.13/mio/struct.Poll.html + +use std::marker::PhantomData; +use std::rc::Rc; +use std::sync::{Arc, Mutex, Condvar}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +/// Block the current thread. +/// +/// See [module documentation][mod] for more details. +/// +/// [mod]: ../index.html +pub trait Park { + /// Unpark handle type for the `Park` implementation. + type Unpark: Unpark; + + /// Error returned by `park` + type Error; + + /// Get a new `Unpark` handle associated with this `Park` instance. + fn unpark(&self) -> Self::Unpark; + + /// Block the current thread unless or until the token is available. + /// + /// A call to `park` does not guarantee that the thread will remain blocked + /// forever, and callers should be prepared for this possibility. This + /// function may wakeup spuriously for any reason. + /// + /// See [module documentation][mod] for more details. + /// + /// # Panics + /// + /// This function **should** not panic, but ultimiately, panics are left as + /// an implementation detail. Refer to the documentation for the specific + /// `Park` implementation + /// + /// [mod]: ../index.html + fn park(&mut self) -> Result<(), Self::Error>; + + /// Park the current thread for at most `duration`. + /// + /// This function is the same as `park` but allows specifying a maximum time + /// to block the thread for. + /// + /// Same as `park`, there is no guarantee that the thread will remain + /// blocked for any amount of time. Spurious wakeups are permitted for any + /// reason. + /// + /// See [module documentation][mod] for more details. + /// + /// # Panics + /// + /// This function **should** not panic, but ultimiately, panics are left as + /// an implementation detail. Refer to the documentation for the specific + /// `Park` implementation + /// + /// [mod]: ../index.html + fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error>; +} + +/// Unblock a thread blocked by the associated [`Park`] instance. +/// +/// See [module documentation][mod] for more details. +/// +/// [mod]: ../index.html +/// [`Park`]: trait.Park.html +pub trait Unpark: Sync + Send + 'static { + /// Unblock a thread that is blocked by the associated `Park` handle. + /// + /// Calling `unpark` atomically makes available the unpark token, if it is + /// not already available. + /// + /// See [module documentation][mod] for more details. + /// + /// # Panics + /// + /// This function **should** not panic, but ultimiately, panics are left as + /// an implementation detail. Refer to the documentation for the specific + /// `Unpark` implementation + /// + /// [mod]: ../index.html + fn unpark(&self); +} + +/// Blocks the current thread using a condition variable. +/// +/// Implements the [`Park`] functionality by using a condition variable. An +/// atomic variable is also used to avoid using the condition variable if +/// possible. +/// +/// The condition variable is cached in a thread-local variable and is shared +/// across all `ParkThread` instances created on the same thread. This also +/// means that an instance of `ParkThread` might be unblocked by a handle +/// associated with a different `ParkThread` instance. +#[derive(Debug)] +pub struct ParkThread { + _anchor: PhantomData>, +} + +/// 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: (), +} + +/// Unblocks a thread that was blocked by `ParkThread`. +#[derive(Clone, Debug)] +pub struct UnparkThread { + inner: Arc, +} + +#[derive(Debug)] +struct Inner { + state: AtomicUsize, + mutex: Mutex<()>, + condvar: Condvar, +} + +const IDLE: usize = 0; +const NOTIFY: usize = 1; +const SLEEP: usize = 2; + +thread_local! { + static CURRENT_PARK_THREAD: Arc = Arc::new(Inner { + state: AtomicUsize::new(IDLE), + mutex: Mutex::new(()), + condvar: Condvar::new(), + }); +} + +// ===== impl ParkThread ===== + +impl ParkThread { + /// Create a new `ParkThread` handle for the current thread. + /// + /// This type cannot be moved to other threads, so it should be created on + /// the thread that the caller intends to park. + pub fn new() -> ParkThread { + ParkThread { + _anchor: PhantomData, + } + } + + /// Get a reference to the `ParkThread` handle for this thread. + fn with_current(&self, f: F) -> R + where F: FnOnce(&Arc) -> R, + { + CURRENT_PARK_THREAD.with(|inner| f(inner)) + } +} + +impl Park for ParkThread { + type Unpark = UnparkThread; + type Error = ParkError; + + fn unpark(&self) -> Self::Unpark { + let inner = self.with_current(|inner| inner.clone()); + UnparkThread { inner } + } + + fn park(&mut self) -> Result<(), Self::Error> { + self.with_current(|inner| inner.park(None)) + } + + fn park_timeout(&mut self, duration: Duration) -> Result<(), Self::Error> { + self.with_current(|inner| inner.park(Some(duration))) + } +} + +// ===== impl UnparkThread ===== + +impl Unpark for UnparkThread { + fn unpark(&self) { + self.inner.unpark(); + } +} + +// ===== impl Inner ===== + +impl Inner { + /// Park the current thread for at most `dur`. + fn park(&self, timeout: Option) -> Result<(), ParkError> { + // If currently notified, then we skip sleeping. This is checked outside + // of the lock to avoid acquiring a mutex if not necessary. + match self.state.compare_and_swap(NOTIFY, IDLE, Ordering::SeqCst) { + NOTIFY => return Ok(()), + IDLE => {}, + _ => unreachable!(), + } + + // The state is currently idle, so obtain the lock and then try to + // transition to a sleeping state. + let mut m = self.mutex.lock().unwrap(); + + // Transition to sleeping + match self.state.compare_and_swap(IDLE, SLEEP, Ordering::SeqCst) { + NOTIFY => { + // Notified before we could sleep, consume the notification and + // exit + self.state.store(IDLE, Ordering::SeqCst); + return Ok(()); + } + IDLE => {}, + _ => unreachable!(), + } + + m = match timeout { + Some(timeout) => self.condvar.wait_timeout(m, timeout).unwrap().0, + None => self.condvar.wait(m).unwrap(), + }; + + // Transition back to idle. If the state has transitione dto `NOTIFY`, + // this will consume that notification + self.state.store(IDLE, Ordering::SeqCst); + + // Explicitly drop the mutex guard. There is no real point in doing it + // except that I find it helpful to make it explicit where we want the + // mutex to unlock. + drop(m); + + Ok(()) + } + + fn unpark(&self) { + // First, try transitioning from IDLE -> NOTIFY, this does not require a + // lock. + match self.state.compare_and_swap(IDLE, NOTIFY, Ordering::SeqCst) { + IDLE | NOTIFY => return, + SLEEP => {} + _ => unreachable!(), + } + + // The other half is sleeping, this requires a lock + let _m = self.mutex.lock().unwrap(); + + // Transition from SLEEP -> NOTIFY + match self.state.compare_and_swap(SLEEP, NOTIFY, Ordering::SeqCst) { + SLEEP => {} + _ => return, + } + + // Wakeup the sleeper + self.condvar.notify_one(); + } +} diff --git a/tokio-executor/tests/executor.rs b/tokio-executor/tests/executor.rs new file mode 100644 index 000000000..77436ec92 --- /dev/null +++ b/tokio-executor/tests/executor.rs @@ -0,0 +1,11 @@ +extern crate tokio_executor; +extern crate futures; + +use tokio_executor::*; +use futures::future::lazy; + +#[test] +fn spawn_out_of_executor_context() { + let res = DefaultExecutor::current().spawn(Box::new(lazy(|| Ok(())))); + assert!(res.is_err()); +} diff --git a/tokio-threadpool/Cargo.toml b/tokio-threadpool/Cargo.toml new file mode 100644 index 000000000..97e81030f --- /dev/null +++ b/tokio-threadpool/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "tokio-threadpool" +version = "0.1.0" +documentation = "https://docs.rs/tokio-threadpool" +repository = "https://github.com/tokio-rs/tokio" +homepage = "https://github.com/tokio-rs/tokio" +license = "MIT/Apache-2.0" +authors = ["Carl Lerche "] +description = """ +A Future aware thread pool based on work stealing. +""" +keywords = ["futures", "tokio"] +categories = ["concurrency", "asynchronous"] + +[dependencies] +tokio-executor = { version = "0.1", path = "../tokio-executor" } +futures = "0.1" +coco = "0.3" +num_cpus = "1.2" +rand = "0.3" +log = "0.3" + +[dev-dependencies] +tokio-timer = "0.1" +env_logger = "0.4" +futures-cpupool = "0.1.7" diff --git a/tokio-threadpool/README.md b/tokio-threadpool/README.md new file mode 100644 index 000000000..14edb236e --- /dev/null +++ b/tokio-threadpool/README.md @@ -0,0 +1,52 @@ +# Tokio Thread Pool + +A library for scheduling execution of futures concurrently across a pool of +threads. + +**Note**: This library isn't quite ready for use. + +### Why not Rayon? + +Rayon is designed to handle parallelizing single computations by breaking them +into smaller chunks. The scheduling for each individual chunk doesn't matter as +long as the root computation completes in a timely fashion. In other words, +Rayon does not provide any guarantees of fairness with regards to how each task +gets scheduled. + +On the other hand, `tokio-threadpool` is a general purpose scheduler and +attempts to schedule each task fairly. This is the ideal behavior when +scheduling a set of unrelated tasks. + +### Why not futures-cpupool? + +It's 10x slower. + +## Examples + +```rust +extern crate tokio_threadpool; +extern crate futures; + +use tokio_threadpool::*; +use futures::*; +use futures::sync::oneshot; + +pub fn main() { + let (tx, _pool) = ThreadPool::new(); + + let res = oneshot::spawn(future::lazy(|| { + println!("Running on the pool"); + Ok::<_, ()>("complete") + }), &tx); + + println!("Result: {:?}", res.wait()); +} +``` + +## License + +`tokio-threadpool` is primarily distributed under the terms of both the MIT +license and the Apache License (Version 2.0), with portions covered by various +BSD-like licenses. + +See LICENSE-APACHE, and LICENSE-MIT for details. diff --git a/tokio-threadpool/benches/basic.rs b/tokio-threadpool/benches/basic.rs new file mode 100644 index 000000000..7217ce0fb --- /dev/null +++ b/tokio-threadpool/benches/basic.rs @@ -0,0 +1,162 @@ +#![feature(test)] + +extern crate futures; +extern crate futures_pool; +extern crate futures_cpupool; +extern crate num_cpus; +extern crate test; + +const NUM_SPAWN: usize = 10_000; +const NUM_YIELD: usize = 1_000; +const TASKS_PER_CPU: usize = 50; + +mod us { + use futures::{task, Async}; + use futures::future::{self, Executor}; + use futures_pool::*; + use num_cpus; + use test; + use std::sync::{mpsc, Arc}; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering::SeqCst; + + #[bench] + fn spawn_many(b: &mut test::Bencher) { + let (sched_tx, _scheduler) = Pool::new(); + + let (tx, rx) = mpsc::sync_channel(10); + let rem = Arc::new(AtomicUsize::new(0)); + + b.iter(move || { + rem.store(super::NUM_SPAWN, SeqCst); + + for _ in 0..super::NUM_SPAWN { + let tx = tx.clone(); + let rem = rem.clone(); + + sched_tx.execute(future::lazy(move || { + if 1 == rem.fetch_sub(1, SeqCst) { + tx.send(()).unwrap(); + } + + Ok(()) + })).ok().unwrap(); + } + + let _ = rx.recv().unwrap(); + }); + } + + #[bench] + fn yield_many(b: &mut test::Bencher) { + let (sched_tx, _scheduler) = Pool::new(); + let tasks = super::TASKS_PER_CPU * num_cpus::get(); + + let (tx, rx) = mpsc::sync_channel(tasks); + + b.iter(move || { + for _ in 0..tasks { + let mut rem = super::NUM_YIELD; + let tx = tx.clone(); + + sched_tx.execute(future::poll_fn(move || { + rem -= 1; + + if rem == 0 { + tx.send(()).unwrap(); + Ok(Async::Ready(())) + } else { + // Notify the current task + task::current().notify(); + + // Not ready + Ok(Async::NotReady) + } + })).ok().unwrap(); + } + + for _ in 0..tasks { + let _ = rx.recv().unwrap(); + } + }); + } +} + +// In this case, CPU pool completes the benchmark faster, but this is due to how +// CpuPool currently behaves, starving other futures. This completes the +// benchmark quickly but results in poor runtime characteristics for a thread +// pool. +// +// See alexcrichton/futures-rs#617 +// +mod cpupool { + use futures::{task, Async}; + use futures::future::{self, Executor}; + use futures_cpupool::*; + use num_cpus; + use test; + use std::sync::{mpsc, Arc}; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering::SeqCst; + + #[bench] + fn spawn_many(b: &mut test::Bencher) { + let pool = CpuPool::new(num_cpus::get()); + + let (tx, rx) = mpsc::sync_channel(10); + let rem = Arc::new(AtomicUsize::new(0)); + + b.iter(move || { + rem.store(super::NUM_SPAWN, SeqCst); + + for _ in 0..super::NUM_SPAWN { + let tx = tx.clone(); + let rem = rem.clone(); + + pool.execute(future::lazy(move || { + if 1 == rem.fetch_sub(1, SeqCst) { + tx.send(()).unwrap(); + } + + Ok(()) + })).ok().unwrap(); + } + + let _ = rx.recv().unwrap(); + }); + } + + #[bench] + fn yield_many(b: &mut test::Bencher) { + let pool = CpuPool::new(num_cpus::get()); + let tasks = super::TASKS_PER_CPU * num_cpus::get(); + + let (tx, rx) = mpsc::sync_channel(tasks); + + b.iter(move || { + for _ in 0..tasks { + let mut rem = super::NUM_YIELD; + let tx = tx.clone(); + + pool.execute(future::poll_fn(move || { + rem -= 1; + + if rem == 0 { + tx.send(()).unwrap(); + Ok(Async::Ready(())) + } else { + // Notify the current task + task::current().notify(); + + // Not ready + Ok(Async::NotReady) + } + })).ok().unwrap(); + } + + for _ in 0..tasks { + let _ = rx.recv().unwrap(); + } + }); + } +} diff --git a/tokio-threadpool/benches/depth.rs b/tokio-threadpool/benches/depth.rs new file mode 100644 index 000000000..2e378beb2 --- /dev/null +++ b/tokio-threadpool/benches/depth.rs @@ -0,0 +1,72 @@ +#![feature(test)] + +extern crate futures; +extern crate futures_pool; +extern crate futures_cpupool; +extern crate num_cpus; +extern crate test; + +const ITER: usize = 20_000; + +mod us { + use futures::future::{self, Executor}; + use futures_pool::*; + use test; + use std::sync::mpsc; + + #[bench] + fn chained_spawn(b: &mut test::Bencher) { + let (sched_tx, _scheduler) = Pool::new(); + + fn spawn(sched_tx: Sender, res_tx: mpsc::Sender<()>, n: usize) { + if n == 0 { + res_tx.send(()).unwrap(); + } else { + let sched_tx2 = sched_tx.clone(); + sched_tx.execute(future::lazy(move || { + spawn(sched_tx2, res_tx, n - 1); + Ok(()) + })).ok().unwrap(); + } + } + + b.iter(move || { + let (res_tx, res_rx) = mpsc::channel(); + + spawn(sched_tx.clone(), res_tx, super::ITER); + res_rx.recv().unwrap(); + }); + } +} + +mod cpupool { + use futures::future::{self, Executor}; + use futures_cpupool::*; + use num_cpus; + use test; + use std::sync::mpsc; + + #[bench] + fn chained_spawn(b: &mut test::Bencher) { + let pool = CpuPool::new(num_cpus::get()); + + fn spawn(pool: CpuPool, res_tx: mpsc::Sender<()>, n: usize) { + if n == 0 { + res_tx.send(()).unwrap(); + } else { + let pool2 = pool.clone(); + pool.execute(future::lazy(move || { + spawn(pool2, res_tx, n - 1); + Ok(()) + })).ok().unwrap(); + } + } + + b.iter(move || { + let (res_tx, res_rx) = mpsc::channel(); + + spawn(pool.clone(), res_tx, super::ITER); + res_rx.recv().unwrap(); + }); + } +} diff --git a/tokio-threadpool/examples/depth.rs b/tokio-threadpool/examples/depth.rs new file mode 100644 index 000000000..7957f09ed --- /dev/null +++ b/tokio-threadpool/examples/depth.rs @@ -0,0 +1,46 @@ +extern crate futures; +extern crate tokio_threadpool; +extern crate env_logger; + +use tokio_threadpool::*; +use futures::future::{self, Executor}; + +use std::sync::mpsc; + +const ITER: usize = 2_000_000; +// const ITER: usize = 30; + +fn chained_spawn() { + let pool = ThreadPool::new(); + let tx = pool.sender().clone(); + + fn spawn(tx: Sender, res_tx: mpsc::Sender<()>, n: usize) { + if n == 0 { + res_tx.send(()).unwrap(); + } else { + let tx2 = tx.clone(); + tx.execute(future::lazy(move || { + spawn(tx2, res_tx, n - 1); + Ok(()) + })).ok().unwrap(); + } + } + + loop { + println!("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"); + let (res_tx, res_rx) = mpsc::channel(); + + for _ in 0..10 { + spawn(tx.clone(), res_tx.clone(), ITER); + } + + for _ in 0..10 { + res_rx.recv().unwrap(); + } + } +} + +pub fn main() { + let _ = ::env_logger::init(); + chained_spawn(); +} diff --git a/tokio-threadpool/examples/hello.rs b/tokio-threadpool/examples/hello.rs new file mode 100644 index 000000000..3324f862a --- /dev/null +++ b/tokio-threadpool/examples/hello.rs @@ -0,0 +1,21 @@ +extern crate futures; +extern crate tokio_threadpool; +extern crate env_logger; + +use tokio_threadpool::*; +use futures::*; +use futures::sync::oneshot; + +pub fn main() { + let _ = ::env_logger::init(); + + let pool = ThreadPool::new(); + let tx = pool.sender().clone(); + + let res = oneshot::spawn(future::lazy(|| { + println!("Running on the pool"); + Ok::<_, ()>("complete") + }), &tx); + + println!("Result: {:?}", res.wait()); +} diff --git a/tokio-threadpool/examples/smoke.rs b/tokio-threadpool/examples/smoke.rs new file mode 100644 index 000000000..8ab144917 --- /dev/null +++ b/tokio-threadpool/examples/smoke.rs @@ -0,0 +1,34 @@ +extern crate futures; +extern crate tokio_threadpool; +extern crate tokio_timer; +extern crate env_logger; + +use tokio_threadpool::*; +use tokio_timer::Timer; + +use futures::*; +use futures::sync::oneshot::spawn; + +use std::thread; +use std::time::Duration; + +pub fn main() { + let _ = ::env_logger::init(); + + let timer = Timer::default(); + { + let pool = ThreadPool::new(); + let tx = pool.sender().clone(); + + let fut = timer.interval(Duration::from_millis(300)) + .for_each(|_| { + println!("~~~~~ Hello ~~~"); + Ok(()) + }) + .map_err(|_| unimplemented!()); + + spawn(fut, &tx).wait().unwrap(); + } + + thread::sleep(Duration::from_millis(100)); +} diff --git a/tokio-threadpool/src/lib.rs b/tokio-threadpool/src/lib.rs new file mode 100644 index 000000000..6a6afe347 --- /dev/null +++ b/tokio-threadpool/src/lib.rs @@ -0,0 +1,2098 @@ +//! A work-stealing based thread pool for executing futures. + +#![deny(warnings, missing_docs, missing_debug_implementations)] + +extern crate tokio_executor; +extern crate futures; +extern crate coco; +extern crate num_cpus; +extern crate rand; + +#[macro_use] +extern crate log; + +mod task; + +use tokio_executor::{Enter, SpawnError}; + +use coco::deque; +use task::Task; + +use futures::{future, Future, Poll, Async}; +use futures::executor::Notify; +use futures::task::AtomicTask; + +use rand::{Rng, SeedableRng, XorShiftRng}; + +use std::{fmt, mem, thread, usize}; +use std::cell::{Cell, UnsafeCell}; +use std::marker::PhantomData; +use std::rc::Rc; +use std::sync::{Arc, Weak, Mutex, Condvar}; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::{AcqRel, Acquire, Release, Relaxed}; +use std::time::{Instant, Duration}; + +/// Work-stealing based thread pool for executing futures. +/// +/// Create `ThreadPool` instances using `Builder`. +#[derive(Debug)] +pub struct ThreadPool { + inner: Sender, +} + +/// 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 { + inner: Arc, +} + +/// 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: ThreadPool, +} + +/// Builds a thread pool with custom configuration values. +/// +/// Methods can be chanined 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 +/// +/// ``` +/// # extern crate tokio_threadpool; +/// # extern crate futures; +/// # use tokio_threadpool::Builder; +/// use futures::future::{Future, lazy}; +/// use std::time::Duration; +/// +/// # pub fn main() { +/// // Create a thread pool with default configuration values +/// let thread_pool = Builder::new() +/// .pool_size(4) +/// .keep_alive(Some(Duration::from_secs(30))) +/// .build(); +/// +/// thread_pool.spawn(lazy(|| { +/// println!("called from a worker thread"); +/// Ok(()) +/// })); +/// +/// // Gracefully shutdown the threadpool +/// thread_pool.shutdown().wait().unwrap(); +/// # } +/// ``` +#[derive(Debug)] +pub struct Builder { + /// Thread pool specific configuration values + config: Config, + + /// Number of workers to spawn + pool_size: usize, +} + +/// Thread pool specific configuration values +#[derive(Debug, Clone)] +struct Config { + keep_alive: Option, + // Used to configure a worker thread + name_prefix: Option, + stack_size: Option, + around_worker: Option, +} + +#[derive(Debug)] +struct Inner { + // ThreadPool state + state: AtomicUsize, + + // Stack tracking sleeping workers. + sleep_stack: AtomicUsize, + + // Number of workers who haven't reached the final state of shutdown + // + // This is only used to know when to single `shutdown_task` once the + // shutdown process has completed. + num_workers: AtomicUsize, + + // Used to generate a thread local RNG seed + next_thread_id: AtomicUsize, + + // Storage for workers + // + // This will *usually* be a small number + workers: Box<[WorkerEntry]>, + + // Task notified when the worker shuts down + shutdown_task: AtomicTask, + + // Configuration + config: Config, +} + +#[derive(Clone)] +struct Callback { + f: Arc, +} + +/// Implements the future `Notify` API. +/// +/// This is how external events are able to signal the task, informing it to try +/// to poll the future again. +#[derive(Debug)] +struct Notifier { + inner: Weak, +} + +/// 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)] +struct State(usize); + +/// Flag used to track if the pool is running +const SHUTDOWN_ON_IDLE: usize = 1; +const SHUTDOWN_NOW: usize = 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. +const MAX_FUTURES: usize = usize::MAX >> NUM_FUTURES_OFFSET; + +/// 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)] +struct SleepStack(usize); + +/// Extracts the head of the worker stack from the scheduler state +const STACK_MASK: usize = ((1 << 16) - 1); + +/// 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. +const MAX_WORKERS: usize = 1 << 15; + +/// Used to mark the stack as empty +const EMPTY: usize = MAX_WORKERS; + +/// Used to mark the stack as terminated +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; + +// Some constants used to work with State +// const A: usize: 0; + +// TODO: This should be split up between what is accessed by each thread and +// what is concurrent. The bits accessed by each thread should be sized to +// exactly one cache line. +#[derive(Debug)] +struct WorkerEntry { + // Worker state. This is mutated when notifying the worker. + state: AtomicUsize, + + // Next entry in the parked Trieber stack + next_sleeper: UnsafeCell, + + // Worker half of deque + deque: deque::Worker, + + // Stealer half of deque + steal: deque::Stealer, + + // Park mutex + park_mutex: Mutex<()>, + + // Park condvar + park_condvar: Condvar, + + // MPSC queue of jobs submitted to the worker from an external source. + inbound: task::Queue, +} + +/// Tracks worker state +#[derive(Clone, Copy, Eq, PartialEq)] +struct WorkerState(usize); + +/// Set when the worker is pushed onto the scheduler's stack of sleeping +/// threads. +const PUSHED_MASK: usize = 0b001; + +/// Manages the worker lifecycle part of the state +const WORKER_LIFECYCLE_MASK: usize = 0b1110; +const WORKER_LIFECYCLE_SHIFT: usize = 1; + +/// The worker does not currently have an associated thread. +const WORKER_SHUTDOWN: usize = 0; + +/// The worker is currently processing its task. +const WORKER_RUNNING: usize = 1; + +/// The worker is currently asleep in the condvar +const WORKER_SLEEPING: usize = 2; + +/// The worker has been notified it should process more work. +const WORKER_NOTIFIED: usize = 3; + +/// 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. +const WORKER_SIGNALED: usize = 4; + +/// Thread worker +/// +/// This is passed to the `around_worker` callback set on `Builder`. This +/// callback is only expected to call `run` on it. +#[derive(Debug)] +pub struct Worker { + // Shared scheduler data + inner: Arc, + + // WorkerEntry index + idx: usize, + + // Set when the worker should finalize on drop + should_finalize: Cell, + + // Keep the value on the current thread. + _p: PhantomData>, +} + +// Pointer to the current worker info +thread_local!(static CURRENT_WORKER: Cell<*const Worker> = Cell::new(0 as *const _)); + +// ===== impl Builder ===== + +impl Builder { + /// Returns a new thread pool builder initialized with default configuration + /// values. + /// + /// Configuration methods can be chained on the return value. + /// + /// # Examples + /// + /// ``` + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::Builder; + /// use std::time::Duration; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// let thread_pool = Builder::new() + /// .pool_size(4) + /// .keep_alive(Some(Duration::from_secs(30))) + /// .build(); + /// # } + /// ``` + pub fn new() -> Builder { + let num_cpus = num_cpus::get(); + + Builder { + pool_size: num_cpus, + config: Config { + keep_alive: None, + name_prefix: None, + stack_size: None, + around_worker: None, + }, + } + } + + /// 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 + /// + /// ``` + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::Builder; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// 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 {}", 32768); + + self.pool_size = val; + self + } + + /// Set the worker thread keep alive duration + /// + /// If set, a worker thread will wait for up to the specified duration for + /// work, at which point the thread will shutdown. When work becomes + /// available, a new thread will eventually be spawned to replace the one + /// that shut down. + /// + /// When the value is `None`, the thread will wait for work forever. + /// + /// The default value is `None`. + /// + /// # Examples + /// + /// ``` + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::Builder; + /// use std::time::Duration; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// let thread_pool = Builder::new() + /// .keep_alive(Some(Duration::from_secs(30))) + /// .build(); + /// # } + /// ``` + pub fn keep_alive(&mut self, val: Option) -> &mut Self { + self.config.keep_alive = val; + 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 + /// + /// ``` + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::Builder; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// let thread_pool = Builder::new() + /// .name_prefix("my-pool-") + /// .build(); + /// # } + /// ``` + pub fn name_prefix>(&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 + /// + /// ``` + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::Builder; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// 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 + /// + /// ``` + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::Builder; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// let thread_pool = Builder::new() + /// .around_worker(|worker, _| { + /// println!("worker is starting up"); + /// worker.run(); + /// println!("worker is shutting down"); + /// }) + /// .build(); + /// # } + /// ``` + pub fn around_worker(&mut self, f: F) -> &mut Self + where F: Fn(&Worker, &mut Enter) + Send + Sync + 'static + { + self.config.around_worker = Some(Callback::new(f)); + self + } + + /// Create the configured `ThreadPool`. + /// + /// The returned `ThreadPool` instance is ready to spawn tasks. + /// + /// # Examples + /// + /// ``` + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::Builder; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// let thread_pool = Builder::new() + /// .build(); + /// # } + /// ``` + pub fn build(&self) -> ThreadPool { + let mut workers = vec![]; + + trace!("build; num-workers={}", self.pool_size); + + for _ in 0..self.pool_size { + workers.push(WorkerEntry::new()); + } + + let inner = Arc::new(Inner { + state: AtomicUsize::new(State::new().into()), + sleep_stack: AtomicUsize::new(SleepStack::new().into()), + num_workers: AtomicUsize::new(self.pool_size), + next_thread_id: AtomicUsize::new(0), + workers: workers.into_boxed_slice(), + shutdown_task: AtomicTask::new(), + config: self.config.clone(), + }); + + // Now, we prime the sleeper stack + for i in 0..self.pool_size { + inner.push_sleeper(i).unwrap(); + } + + let inner = Sender { inner }; + + ThreadPool { inner } + } +} + +// ===== impl ThreadPool ===== + +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() + } + + /// 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 + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::ThreadPool; + /// use futures::future::{Future, lazy}; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// let thread_pool = ThreadPool::new(); + /// + /// thread_pool.spawn(lazy(|| { + /// println!("called from a worker thread"); + /// Ok(()) + /// })); + /// + /// // Gracefully shutdown the threadpool + /// thread_pool.shutdown().wait().unwrap(); + /// # } + /// ``` + /// + /// # Panics + /// + /// This function panics if the spawn fails. Use [`Sender::spawn`] for a + /// version that returns a `Result` instead of panicking. + pub fn spawn(&self, future: F) + where F: Future + Send + 'static, + { + self.sender().spawn(future).unwrap(); + } + + /// 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 + } + + /// Return a mutable reference to the sender handle + pub fn sender_mut(&mut self) -> &mut Sender { + &mut self.inner + } + + /// 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(self) -> Shutdown { + self.inner().shutdown(false, false); + Shutdown { inner: self } + } + + /// 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(self) -> Shutdown { + self.inner().shutdown(true, false); + Shutdown { inner: self } + } + + /// 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(self) -> Shutdown { + self.inner().shutdown(true, true); + Shutdown { inner: self } + } + + fn inner(&self) -> &Inner { + &*self.inner.inner + } +} + +// ===== impl Sender ====== + +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 + /// # extern crate tokio_threadpool; + /// # extern crate futures; + /// # use tokio_threadpool::ThreadPool; + /// use futures::future::{Future, lazy}; + /// + /// # pub fn main() { + /// // Create a thread pool with default configuration values + /// let thread_pool = ThreadPool::new(); + /// + /// thread_pool.sender().spawn(lazy(|| { + /// println!("called from a worker thread"); + /// Ok(()) + /// })).unwrap(); + /// + /// // Gracefully shutdown the threadpool + /// thread_pool.shutdown().wait().unwrap(); + /// # } + /// ``` + pub fn spawn(&self, future: F) -> Result<(), SpawnError> + where F: Future + Send + 'static, + { + let mut s = self; + tokio_executor::Executor::spawn(&mut s, Box::new(future)) + } + + /// Logic to prepare for spawning + fn prepare_for_spawn(&self) -> Result<(), SpawnError> { + let mut state: State = self.inner.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() == SHUTDOWN_NOW { + // Cannot execute the future, executor is shutdown. + return Err(SpawnError::shutdown()); + } + + next.inc_num_futures(); + + let actual = self.inner.state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + trace!("execute; count={:?}", next.num_futures()); + break; + } + + state = actual; + } + + Ok(()) + } +} + +impl tokio_executor::Executor for Sender { + fn status(&self) -> Result<(), tokio_executor::SpawnError> { + let s = self; + tokio_executor::Executor::status(&s) + } + + fn spawn(&mut self, future: Box + Send>) + -> Result<(), SpawnError> + { + let mut s = &*self; + tokio_executor::Executor::spawn(&mut s, future) + } +} + +impl<'a> tokio_executor::Executor for &'a Sender { + fn status(&self) -> Result<(), tokio_executor::SpawnError> { + let state: State = self.inner.state.load(Acquire).into(); + + if state.num_futures() == MAX_FUTURES { + // No capacity + return Err(SpawnError::at_capacity()); + } + + if state.lifecycle() == SHUTDOWN_NOW { + // Cannot execute the future, executor is shutdown. + return Err(SpawnError::shutdown()); + } + + Ok(()) + } + + fn spawn(&mut self, future: Box + 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 = Task::new(future); + + self.inner.submit(task, &self.inner); + + Ok(()) + } +} + +impl future::Executor for Sender +where T: Future + Send + 'static, +{ + fn execute(&self, future: T) -> Result<(), future::ExecuteError> { + if let Err(e) = tokio_executor::Executor::status(self) { + let kind = if e.is_at_capacity() { + future::ExecuteErrorKind::NoCapacity + } else { + future::ExecuteErrorKind::Shutdown + }; + + return Err(future::ExecuteError::new(kind, future)); + } + + let _ = self.spawn(future); + Ok(()) + } +} + +impl Clone for Sender { + #[inline] + fn clone(&self) -> Sender { + let inner = self.inner.clone(); + Sender { inner } + } +} + +// ===== impl Shutdown ===== + +impl Shutdown { + fn inner(&self) -> &Inner { + self.inner.inner() + } +} + +impl Future for Shutdown { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + trace!("Shutdown::poll"); + + self.inner().shutdown_task.register(); + + if 0 != self.inner().num_workers.load(Acquire) { + return Ok(Async::NotReady); + } + + Ok(().into()) + } +} + +// ===== impl Inner ===== + +impl Inner { + /// Start shutting down the pool. This means that no new futures will be + /// accepted. + 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() >= SHUTDOWN_NOW { + // 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. + SHUTDOWN_NOW + } else { + SHUTDOWN_ON_IDLE + }); + + 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(); + } + + fn terminate_sleeping_workers(&self) { + 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.pop_sleeper(WORKER_SIGNALED, TERMINATED) { + trace!(" -> shutdown worker; idx={:?}; state={:?}", idx, worker_state); + self.signal_stop(idx, worker_state); + } + } + + /// Signals to the worker that it should stop + fn signal_stop(&self, idx: usize, mut state: WorkerState) { + let worker = &self.workers[idx]; + + // Transition the worker state to signaled + loop { + let mut next = state; + + match state.lifecycle() { + WORKER_SHUTDOWN => { + trace!("signal_stop -- WORKER_SHUTDOWN; idx={}", idx); + // If the worker is in the shutdown state, then it will never be + // started again. + self.worker_terminated(); + + return; + } + WORKER_RUNNING | WORKER_SLEEPING => {} + _ => { + trace!("signal_stop -- skipping; idx={}; state={:?}", idx, state); + // All other states will naturally converge to a state of + // shutdown. + return; + } + } + + next.set_lifecycle(WORKER_SIGNALED); + + let actual = worker.state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + break; + } + + state = actual; + } + + // Wakeup the worker + worker.wakeup(); + } + + fn worker_terminated(&self) { + let prev = self.num_workers.fetch_sub(1, AcqRel); + + trace!("worker_terminated; num_workers={}", prev - 1); + + if 1 == prev { + trace!("notifying shutdown task"); + self.shutdown_task.notify(); + } + } + + /// 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. + fn submit(&self, task: Task, inner: &Arc) { + Worker::with_current(|worker| { + match worker { + Some(worker) => { + let idx = worker.idx; + + trace!(" -> submit internal; idx={}", idx); + + worker.inner.workers[idx].submit_internal(task); + worker.inner.signal_work(inner); + } + None => { + self.submit_external(task, inner); + } + } + }); + } + + /// Submit a task to the scheduler from off worker + /// + /// Called from outside of the scheduler, this function is how new tasks + /// enter the system. + fn submit_external(&self, task: Task, inner: &Arc) { + // First try to get a handle to a sleeping worker. This ensures that + // sleeping tasks get woken up + if let Some((idx, state)) = self.pop_sleeper(WORKER_NOTIFIED, EMPTY) { + trace!("submit to existing worker; idx={}; state={:?}", idx, state); + self.submit_to_external(idx, task, state, inner); + return; + } + + // All workers are active, so pick a random worker and submit the + // task to it. + let len = self.workers.len(); + let idx = self.rand_usize() % len; + + trace!(" -> submitting to random; idx={}", idx); + + let state: WorkerState = self.workers[idx].state.load(Acquire).into(); + self.submit_to_external(idx, task, state, inner); + } + + fn submit_to_external(&self, + idx: usize, + task: Task, + state: WorkerState, + inner: &Arc) + { + let entry = &self.workers[idx]; + + if !entry.submit_external(task, state) { + Worker::spawn(idx, inner); + } + } + + /// 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. + fn signal_work(&self, inner: &Arc) { + if let Some((idx, mut state)) = self.pop_sleeper(WORKER_SIGNALED, EMPTY) { + let entry = &self.workers[idx]; + + // Transition the worker state to signaled + loop { + let mut next = state; + + // pop_sleeper should skip these + debug_assert!(state.lifecycle() != WORKER_SIGNALED); + next.set_lifecycle(WORKER_SIGNALED); + + let actual = entry.state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + break; + } + + state = actual; + } + + // The state has been transitioned to signal, now we need to wake up + // the worker if necessary. + match state.lifecycle() { + WORKER_SLEEPING => { + trace!("signal_work -- wakeup; idx={}", idx); + self.workers[idx].wakeup(); + } + WORKER_SHUTDOWN => { + trace!("signal_work -- spawn; idx={}", idx); + Worker::spawn(idx, inner); + } + _ => {} + } + } + } + + /// Push a worker on the sleep stack + /// + /// Returns `Err` if the pool has been terminated + fn push_sleeper(&self, idx: usize) -> Result<(), ()> { + let mut state: SleepStack = self.sleep_stack.load(Acquire).into(); + + debug_assert!(WorkerState::from(self.workers[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(()); + } + + self.workers[idx].set_next_sleeper(head); + next.set_head(idx); + + let actual = self.sleep_stack.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if state == actual { + return Ok(()); + } + + state = actual; + } + } + + /// Pop a worker from the sleep stack + fn pop_sleeper(&self, max_lifecycle: usize, terminal: usize) + -> Option<(usize, WorkerState)> + { + debug_assert!(terminal == EMPTY || terminal == TERMINATED); + + let mut state: SleepStack = self.sleep_stack.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.sleep_stack.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 = self.workers[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.sleep_stack.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + // The worker has been removed from the stack, so the pushed bit + // can be unset. Release ordering is used to ensure that this + // operation happens after actually popping the task. + debug_assert_eq!(1, PUSHED_MASK); + + // Unset the PUSHED flag and get the current state. + let state: WorkerState = self.workers[head].state + .fetch_sub(PUSHED_MASK, Release).into(); + + 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; + } + } + + /// Generates a random number + /// + /// Uses a thread-local seeded XorShift. + fn rand_usize(&self) -> usize { + // Use a thread-local random number generator. If the thread does not + // have one yet, then seed a new one + thread_local!(static THREAD_RNG_KEY: UnsafeCell> = UnsafeCell::new(None)); + + THREAD_RNG_KEY.with(|t| { + #[cfg(target_pointer_width = "32")] + fn new_rng(thread_id: usize) -> XorShiftRng { + XorShiftRng::from_seed([ + thread_id as u32, + 0x00000000, + 0xa8a7d469, + 0x97830e05]) + } + + #[cfg(target_pointer_width = "64")] + fn new_rng(thread_id: usize) -> XorShiftRng { + XorShiftRng::from_seed([ + thread_id as u32, + (thread_id >> 32) as u32, + 0xa8a7d469, + 0x97830e05]) + } + + let thread_id = self.next_thread_id.fetch_add(1, Relaxed); + let rng = unsafe { &mut *t.get() }; + + if rng.is_none() { + *rng = Some(new_rng(thread_id)); + } + + rng.as_mut().unwrap().next_u32() as usize + }) + } +} + +impl Notify for Notifier { + fn notify(&self, id: usize) { + trace!("Notifier::notify; id=0x{:x}", id); + + let id = id as usize; + let task = unsafe { Task::from_notify_id_ref(&id) }; + + if !task.schedule() { + trace!(" -> task already scheduled"); + // task is already scheduled, there is nothing more to do + return; + } + + // TODO: Check if the pool is still running + + // Bump the ref count + let task = task.clone(); + + if let Some(inner) = self.inner.upgrade() { + let _ = inner.submit(task, &inner); + } + } + + fn clone_id(&self, id: usize) -> usize { + unsafe { + let handle = Task::from_notify_id_ref(&id); + mem::forget(handle.clone()); + } + + id + } + + fn drop_id(&self, id: usize) { + unsafe { + let _ = Task::from_notify_id(id); + } + } +} + +unsafe impl Send for Inner {} +unsafe impl Sync for Inner {} + +// ===== impl Worker ===== + +impl Worker { + fn spawn(idx: usize, inner: &Arc) { + trace!("spawning new worker thread; idx={}", idx); + + let mut th = thread::Builder::new(); + + if let Some(ref prefix) = inner.config.name_prefix { + th = th.name(format!("{}{}", prefix, idx)); + } + + if let Some(stack) = inner.config.stack_size { + th = th.stack_size(stack); + } + + let inner = inner.clone(); + + th.spawn(move || { + let worker = Worker { + inner: inner, + idx: idx, + should_finalize: Cell::new(false), + _p: PhantomData, + }; + + // Make sure the ref to the worker does not move + let wref = &worker; + + // Create another worker... It's ok, this is just a new type around + // `Inner` that is expected to stay on the current thread. + CURRENT_WORKER.with(|c| { + c.set(wref as *const _); + + let inner = wref.inner.clone(); + let mut sender = Sender { inner }; + + // Enter an execution context + let mut enter = tokio_executor::enter().unwrap(); + + tokio_executor::with_default(&mut sender, &mut enter, |enter| { + if let Some(ref callback) = wref.inner.config.around_worker { + callback.call(wref, enter); + } else { + wref.run(); + } + }); + }); + }).unwrap(); + } + + fn with_current) -> 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 })) + } + }) + } + + /// Run the worker + /// + /// This function blocks until the worker is shutting down. + pub fn run(&self) { + // Get the notifier. + let notify = Arc::new(Notifier { + inner: Arc::downgrade(&self.inner), + }); + + let mut first = true; + let mut spin_cnt = 0; + + while self.check_run_state(first) { + first = false; + + // Poll inbound until empty, transfering all tasks to the internal + // queue. + let consistent = self.drain_inbound(); + + // Run the next available task + if self.try_run_task(¬ify) { + spin_cnt = 0; + // As long as there is work, keep looping. + continue; + } + + // No work in this worker's queue, it is time to try stealing. + if self.try_steal_task(¬ify) { + spin_cnt = 0; + continue; + } + + if !consistent { + spin_cnt = 0; + continue; + } + + // Starting to get sleeeeepy + if spin_cnt < 32 { + spin_cnt += 1; + + // Don't do anything further + } else if spin_cnt < 256 { + spin_cnt += 1; + + // Yield the thread + thread::yield_now(); + } else { + if !self.sleep() { + return; + } + } + + // If there still isn't any work to do, shutdown the worker? + } + + self.should_finalize.set(true); + } + + /// 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 { + let mut state: WorkerState = self.entry().state.load(Acquire).into(); + + loop { + let pool_state: State = self.inner.state.load(Acquire).into(); + + if pool_state.is_terminated() { + return false; + } + + let mut next = state; + + match state.lifecycle() { + WORKER_RUNNING => break, + WORKER_NOTIFIED | WORKER_SIGNALED => { + // transition back to running + next.set_lifecycle(WORKER_RUNNING); + } + lifecycle => panic!("unexpected worker state; lifecycle={}", lifecycle), + } + + let actual = self.entry().state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + break; + } + + state = actual; + } + + // If this is the first iteration of the worker loop, then the state can + // be signaled. + 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.inner.signal_work(&self.inner); + } + + true + } + + /// Runs the next task on this worker's queue. + /// + /// Returns `true` if work was found. + #[inline] + fn try_run_task(&self, notify: &Arc) -> bool { + use coco::deque::Steal::*; + + // Poll the internal queue for a task to run + match self.entry().deque.steal_weak() { + Data(task) => { + self.run_task(task, notify); + true + } + Empty => false, + Inconsistent => true, + } + } + + /// Tries to steal a task from another worker. + /// + /// Returns `true` if work was found + #[inline] + fn try_steal_task(&self, notify: &Arc) -> bool { + use coco::deque::Steal::*; + + let len = self.inner.workers.len(); + let mut idx = self.inner.rand_usize() % len; + let mut found_work = false; + let start = idx; + + loop { + if idx < len { + match self.inner.workers[idx].steal.steal_weak() { + Data(task) => { + trace!("stole task"); + + self.run_task(task, notify); + + trace!("try_steal_task -- signal_work; self={}; from={}", + self.idx, idx); + + // Signal other workers that work is available + self.inner.signal_work(&self.inner); + + return true; + } + Empty => {} + Inconsistent => found_work = true, + } + + idx += 1; + } else { + idx = 0; + } + + if idx == start { + break; + } + } + + found_work + } + + fn run_task(&self, task: Task, notify: &Arc) { + use task::Run::*; + + match task.run(notify) { + Idle => {} + Schedule => { + self.entry().push_internal(task); + } + Complete => { + let mut state: State = self.inner.state.load(Acquire).into(); + + loop { + let mut next = state; + next.dec_num_futures(); + + let actual = self.inner.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.inner.terminate_sleeping_workers(); + } + } + + // The worker's run loop will detect the shutdown state + // next iteration. + return; + } + + state = actual; + } + } + } + } + + /// Drains all tasks on the extern queue and pushes them onto the internal + /// queue. + /// + /// Returns `true` if the operation was able to complete in a consistent + /// state. + #[inline] + fn drain_inbound(&self) -> bool { + use task::Poll::*; + + let mut found_work = false; + + loop { + let task = unsafe { self.entry().inbound.poll() }; + + match task { + Empty => { + if found_work { + trace!("found work while draining; signal_work"); + self.inner.signal_work(&self.inner); + } + + return true; + } + Inconsistent => { + if found_work { + trace!("found work while draining; signal_work"); + self.inner.signal_work(&self.inner); + } + + return false; + } + Data(task) => { + found_work = true; + self.entry().push_internal(task); + } + } + } + } + + /// Put the worker to sleep + /// + /// Returns `true` if woken up due to new work arriving. + #[inline] + fn sleep(&self) -> bool { + trace!("Worker::sleep; idx={}", self.idx); + + let mut state: WorkerState = 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. However, part of + // this process is also to do some final state checks to avoid entering + // the mutex if at all possible. + + loop { + let mut next = state; + + match state.lifecycle() { + WORKER_RUNNING => { + // Try setting the pushed state + next.set_pushed(); + } + WORKER_NOTIFIED | WORKER_SIGNALED => { + // No need to sleep, transition back to running and move on. + next.set_lifecycle(WORKER_RUNNING); + } + actual => panic!("unexpected worker state; {}", actual), + } + + 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.idx); + + // We obtained permission to push the worker into the + // sleeper queue. + if let Err(_) = self.inner.push_sleeper(self.idx) { + trace!(" sleeping -- push to stack failed; idx={}", self.idx); + // 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; + } + + // Acquire the sleep mutex, the state is transitioned to sleeping within + // the mutex in order to avoid losing wakeup notifications. + let mut lock = self.entry().park_mutex.lock().unwrap(); + + // Transition the state to sleeping, a CAS is still needed as other + // state transitions could happen unrelated to the sleep / wakeup + // process. We also have to redo the lifecycle check done above as + // the state could have been transitioned before entering the mutex. + loop { + let mut next = state; + + match state.lifecycle() { + WORKER_RUNNING => {} + WORKER_NOTIFIED | WORKER_SIGNALED => { + // Release the lock, sleep will not happen this call. + drop(lock); + + // Transition back to running + loop { + let mut next = state; + next.set_lifecycle(WORKER_RUNNING); + + let actual = self.entry().state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + return true; + } + + state = actual; + } + } + _ => unreachable!(), + } + + trace!(" sleeping -- set WORKER_SLEEPING; idx={}", self.idx); + + next.set_lifecycle(WORKER_SLEEPING); + + let actual = self.entry().state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + break; + } + + state = actual; + } + + trace!(" -> starting to sleep; idx={}", self.idx); + + let sleep_until = self.inner.config.keep_alive + .map(|dur| Instant::now() + dur); + + // The state has been transitioned to sleeping, we can now wait on the + // condvar. This is done in a loop as condvars can wakeup spuriously. + loop { + let mut drop_thread = false; + + lock = match sleep_until { + Some(when) => { + let now = Instant::now(); + + if when >= now { + drop_thread = true; + } + + let dur = when - now; + + self.entry().park_condvar + .wait_timeout(lock, dur) + .unwrap().0 + } + None => { + self.entry().park_condvar.wait(lock).unwrap() + } + }; + + trace!(" -> wakeup; idx={}", self.idx); + + // Reload the state + state = self.entry().state.load(Acquire).into(); + + loop { + match state.lifecycle() { + WORKER_SLEEPING => {} + WORKER_NOTIFIED | WORKER_SIGNALED => { + // Release the lock, done sleeping + drop(lock); + + // Transition back to running + loop { + let mut next = state; + next.set_lifecycle(WORKER_RUNNING); + + let actual = self.entry().state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + return true; + } + + state = actual; + } + } + _ => unreachable!(), + } + + if !drop_thread { + break; + } + + let mut next = state; + next.set_lifecycle(WORKER_SHUTDOWN); + + let actual = self.entry().state.compare_and_swap( + state.into(), next.into(), AcqRel).into(); + + if actual == state { + // Transitioned to a shutdown state + return false; + } + + state = actual; + } + + // The worker hasn't been notified, go back to sleep + } + } + + fn entry(&self) -> &WorkerEntry { + &self.inner.workers[self.idx] + } +} + +impl Drop for Worker { + fn drop(&mut self) { + trace!("shutting down thread; idx={}", self.idx); + + if self.should_finalize.get() { + // Drain all work + self.drain_inbound(); + + while let Some(_) = self.entry().deque.pop() { + } + + // TODO: Drain the work queue... + self.inner.worker_terminated(); + } + } +} + +// ===== impl State ===== + +impl State { + #[inline] + fn new() -> State { + State(0) + } + + /// Returns the number of futures still pending completion. + fn num_futures(&self) -> usize { + self.0 >> NUM_FUTURES_OFFSET + } + + /// Increment the number of futures pending completion. + /// + /// Returns false on failure. + fn inc_num_futures(&mut self) { + debug_assert!(self.num_futures() < MAX_FUTURES); + debug_assert!(self.lifecycle() < SHUTDOWN_NOW); + + self.0 += 1 << NUM_FUTURES_OFFSET; + } + + /// Decrement the number of futures pending completion. + 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() == SHUTDOWN_ON_IDLE && num_futures == 1 { + self.0 = SHUTDOWN_NOW; + } + } + + /// Set the number of futures pending completion to zero + fn clear_num_futures(&mut self) { + self.0 = self.0 & LIFECYCLE_MASK; + } + + fn lifecycle(&self) -> usize { + self.0 & LIFECYCLE_MASK + } + + fn set_lifecycle(&mut self, val: usize) { + self.0 = (self.0 & NUM_FUTURES_MASK) | val; + } + + fn is_terminated(&self) -> bool { + self.lifecycle() == SHUTDOWN_NOW && self.num_futures() == 0 + } +} + +impl From for State { + fn from(src: usize) -> Self { + State(src) + } +} + +impl From 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("State") + .field("lifecycle", &self.lifecycle()) + .field("num_futures", &self.num_futures()) + .finish() + } +} + +// ===== impl SleepStack ===== + +impl SleepStack { + #[inline] + fn new() -> SleepStack { + SleepStack(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 for SleepStack { + fn from(src: usize) -> Self { + SleepStack(src) + } +} + +impl From for usize { + fn from(src: SleepStack) -> Self { + src.0 + } +} + +impl fmt::Debug for SleepStack { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + let head = self.head(); + + let mut fmt = fmt.debug_struct("SleepStack"); + + 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() + } +} + +// ===== impl WorkerEntry ===== + +impl WorkerEntry { + fn new() -> Self { + let (w, s) = deque::new(); + + WorkerEntry { + state: AtomicUsize::new(WorkerState::default().into()), + next_sleeper: UnsafeCell::new(0), + deque: w, + steal: s, + inbound: task::Queue::new(), + park_mutex: Mutex::new(()), + park_condvar: Condvar::new(), + } + } + + #[inline] + fn submit_internal(&self, task: Task) { + self.push_internal(task); + } + + /// Submits a task to the worker. This assumes that the caller is external + /// to the worker. Internal submissions go through another path. + /// + /// Returns `false` if the worker needs to be spawned. + fn submit_external(&self, task: Task, mut state: WorkerState) -> bool { + // Push the task onto the external queue + self.push_external(task); + + 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() { + WORKER_SLEEPING => { + // The worker is currently sleeping, the condition variable must + // be signaled + self.wakeup(); + true + } + WORKER_SHUTDOWN => false, + _ => true, + } + } + + #[inline] + fn push_external(&self, task: Task) { + self.inbound.push(task); + } + + #[inline] + fn push_internal(&self, task: Task) { + self.deque.push(task); + } + + #[inline] + fn wakeup(&self) { + let _lock = self.park_mutex.lock().unwrap(); + self.park_condvar.notify_one(); + } + + #[inline] + fn next_sleeper(&self) -> usize { + unsafe { *self.next_sleeper.get() } + } + + #[inline] + fn set_next_sleeper(&self, val: usize) { + unsafe { *self.next_sleeper.get() = val; } + } +} + +// ===== impl WorkerState ===== + +impl WorkerState { + /// Returns true if the worker entry is pushed in the sleeper stack + fn is_pushed(&self) -> bool { + self.0 & PUSHED_MASK == PUSHED_MASK + } + + fn set_pushed(&mut self) { + self.0 |= PUSHED_MASK + } + + fn is_notified(&self) -> bool { + match self.lifecycle() { + WORKER_NOTIFIED | WORKER_SIGNALED => true, + _ => false, + } + } + + fn lifecycle(&self) -> usize { + (self.0 & WORKER_LIFECYCLE_MASK) >> WORKER_LIFECYCLE_SHIFT + } + + fn set_lifecycle(&mut self, val: usize) { + self.0 = (self.0 & !WORKER_LIFECYCLE_MASK) | + (val << WORKER_LIFECYCLE_SHIFT) + } + + fn is_signaled(&self) -> bool { + self.lifecycle() == WORKER_SIGNALED + } + + fn notify(&mut self) { + if self.lifecycle() != WORKER_SIGNALED { + self.set_lifecycle(WORKER_NOTIFIED) + } + } +} + +impl Default for WorkerState { + fn default() -> WorkerState { + // All workers will start pushed in the sleeping stack + WorkerState(PUSHED_MASK) + } +} + +impl From for WorkerState { + fn from(src: usize) -> Self { + WorkerState(src) + } +} + +impl From for usize { + fn from(src: WorkerState) -> Self { + src.0 + } +} + +impl fmt::Debug for WorkerState { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + fmt.debug_struct("WorkerState") + .field("lifecycle", &match self.lifecycle() { + WORKER_SHUTDOWN => "WORKER_SHUTDOWN", + WORKER_RUNNING => "WORKER_RUNNING", + WORKER_SLEEPING => "WORKER_SLEEPING", + WORKER_NOTIFIED => "WORKER_NOTIFIED", + WORKER_SIGNALED => "WORKER_SIGNALED", + _ => unreachable!(), + }) + .field("is_pushed", &self.is_pushed()) + .finish() + } +} + +// ===== impl Callback ===== + +impl Callback { + fn new(f: F) -> Self + where F: Fn(&Worker, &mut Enter) + Send + Sync + 'static + { + Callback { f: Arc::new(f) } + } + + pub fn call(&self, worker: &Worker, enter: &mut Enter) { + (self.f)(worker, enter) + } +} + +impl fmt::Debug for Callback { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + write!(fmt, "Fn") + } +} diff --git a/tokio-threadpool/src/task.rs b/tokio-threadpool/src/task.rs new file mode 100644 index 000000000..af8e8c13d --- /dev/null +++ b/tokio-threadpool/src/task.rs @@ -0,0 +1,429 @@ +use Notifier; + +use futures::{future, Future, Async}; +use futures::executor::{self, Spawn}; + +use std::{fmt, mem, ptr}; +use std::cell::Cell; +use std::sync::Arc; +use std::sync::atomic::{self, AtomicUsize, AtomicPtr}; +use std::sync::atomic::Ordering::{AcqRel, Acquire, Release, Relaxed}; + +pub(crate) struct Task { + ptr: *mut Inner, +} + +#[derive(Debug)] +pub(crate) struct Queue { + head: AtomicPtr, + tail: Cell<*mut Inner>, + stub: Box, +} + +#[derive(Debug)] +pub(crate) enum Poll { + Empty, + Inconsistent, + Data(Task), +} + +#[derive(Debug)] +pub(crate) enum Run { + Idle, + Schedule, + Complete, +} + +struct Inner { + // Next pointer in the queue that submits tasks to a worker. + next: AtomicPtr, + + // Task state + state: AtomicUsize, + + // Number of outstanding references to the task + ref_count: AtomicUsize, + + // Store the future at the head of the struct + // + // The future is dropped immediately when it transitions to Complete + future: Option>, +} + +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +enum State { + /// Task is currently idle + Idle, + /// Task is currently running + Running, + /// Task is currently running, but has been notified that it must run again. + Notified, + /// Task has been scheduled + Scheduled, + /// Task is complete + Complete, +} + +type BoxFuture = Box + Send + 'static>; + +// ===== impl Task ===== + +impl Task { + /// Create a new task handle + pub fn new(future: BoxFuture) -> Task { + let inner = Box::new(Inner { + next: AtomicPtr::new(ptr::null_mut()), + state: AtomicUsize::new(State::new().into()), + ref_count: AtomicUsize::new(1), + future: Some(executor::spawn(future)), + }); + + Task { ptr: Box::into_raw(inner) } + } + + /// Transmute a u64 to a Task + pub unsafe fn from_notify_id(unpark_id: usize) -> Task { + mem::transmute(unpark_id) + } + + /// Transmute a u64 to a task ref + pub unsafe fn from_notify_id_ref<'a>(unpark_id: &'a usize) -> &'a Task { + mem::transmute(unpark_id) + } + + /// Execute the task returning `Run::Schedule` if the task needs to be + /// scheduled again. + pub fn run(&self, unpark: &Arc) -> Run { + use self::State::*; + + // Transition task to running state. At this point, the task must be + // scheduled. + let actual: State = self.inner().state.compare_and_swap( + Scheduled.into(), Running.into(), AcqRel).into(); + + trace!("running; state={:?}", actual); + + match actual { + Scheduled => {}, + _ => panic!("unexpected task state; {:?}", actual), + } + + trace!("Task::run; state={:?}", State::from(self.inner().state.load(Relaxed))); + + let res = self.inner_mut().future.as_mut().unwrap() + .poll_future_notify(unpark, self.ptr as usize); + + match res { + Ok(Async::Ready(_)) | Err(_) => { + trace!(" -> task complete"); + + // Drop the future + self.inner_mut().drop_future(); + + // Transition to the completed state + self.inner().state.store(State::Complete.into(), Release); + + Run::Complete + } + _ => { + 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 = self.inner().state.compare_and_swap( + Running.into(), Idle.into(), AcqRel).into(); + + match prev { + Running => Run::Idle, + Notified => { + self.inner().state.store(Scheduled.into(), Release); + Run::Schedule + } + _ => unreachable!(), + } + } + } + } + + /// Transition the task state to scheduled. + /// + /// Returns `true` if the caller is permitted to schedule the task. + pub fn schedule(&self) -> bool { + use self::State::*; + + loop { + let actual = self.inner().state.compare_and_swap( + Idle.into(), + Scheduled.into(), + Relaxed).into(); + + match actual { + Idle => return true, + Running => { + let actual = self.inner().state.compare_and_swap( + Running.into(), Notified.into(), Relaxed).into(); + + match actual { + Idle => continue, + _ => return false, + } + } + Complete | Notified | Scheduled => return false, + } + } + } + + #[inline] + fn inner(&self) -> &Inner { + unsafe { &*self.ptr } + } + + #[inline] + fn inner_mut(&self) -> &mut Inner { + unsafe { &mut *self.ptr } + } +} + +impl fmt::Debug for Task { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + fmt.debug_struct("Task") + .field("inner", self.inner()) + .finish() + } +} + +impl Clone for Task { + fn clone(&self) -> Task { + use std::isize; + + const MAX_REFCOUNT: usize = (isize::MAX) as usize; + // Using a relaxed ordering is alright here, as knowledge of the + // original reference prevents other threads from erroneously deleting + // the object. + // + // As explained in the [Boost documentation][1], Increasing the + // reference counter can always be done with memory_order_relaxed: New + // references to an object can only be formed from an existing + // reference, and passing an existing reference from one thread to + // another must already provide any required synchronization. + // + // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) + let old_size = self.inner().ref_count.fetch_add(1, Relaxed); + + // However we need to guard against massive refcounts in case someone + // is `mem::forget`ing Arcs. If we don't do this the count can overflow + // and users will use-after free. We racily saturate to `isize::MAX` on + // the assumption that there aren't ~2 billion threads incrementing + // the reference count at once. This branch will never be taken in + // any realistic program. + // + // We abort because such a program is incredibly degenerate, and we + // don't care to support it. + if old_size > MAX_REFCOUNT { + // TODO: abort + panic!(); + } + + Task { ptr: self.ptr } + } +} + +impl Drop for Task { + fn drop(&mut self) { + // Because `fetch_sub` is already atomic, we do not need to synchronize + // with other threads unless we are going to delete the object. This + // same logic applies to the below `fetch_sub` to the `weak` count. + if self.inner().ref_count.fetch_sub(1, Release) != 1 { + return; + } + + // This fence is needed to prevent reordering of use of the data and + // deletion of the data. Because it is marked `Release`, the decreasing + // of the reference count synchronizes with this `Acquire` fence. This + // means that use of the data happens before decreasing the reference + // count, which happens before this fence, which happens before the + // deletion of the data. + // + // As explained in the [Boost documentation][1], + // + // > It is important to enforce any possible access to the object in one + // > thread (through an existing reference) to *happen before* deleting + // > the object in a different thread. This is achieved by a "release" + // > operation after dropping a reference (any access to the object + // > through this reference must obviously happened before), and an + // > "acquire" operation before deleting the object. + // + // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) + atomic::fence(Acquire); + + unsafe { + let _ = Box::from_raw(self.ptr); + } + } +} + +unsafe impl Send for Task {} + +// ===== impl Inner ===== + +impl Inner { + fn stub() -> Inner { + Inner { + next: AtomicPtr::new(ptr::null_mut()), + state: AtomicUsize::new(State::stub().into()), + ref_count: AtomicUsize::new(0), + future: Some(executor::spawn(Box::new(future::empty()))), + } + } + + fn drop_future(&mut self) { + let _ = self.future.take(); + } +} + +impl Drop for Inner { + fn drop(&mut self) { + self.drop_future(); + } +} + +impl fmt::Debug for Inner { + fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result { + fmt.debug_struct("Inner") + .field("next", &self.next) + .field("state", &self.state) + .field("ref_count", &self.ref_count) + .field("future", &"Spawn") + .finish() + } +} + +// ===== impl Queue ===== + +impl Queue { + pub fn new() -> Queue { + let stub = Box::new(Inner::stub()); + let ptr = &*stub as *const _ as *mut _; + + Queue { + head: AtomicPtr::new(ptr), + tail: Cell::new(ptr), + stub: stub, + } + } + + pub fn push(&self, handle: Task) { + unsafe { + self.push2(handle.ptr); + + // Forgetting the handle is necessary to avoid the ref dec + mem::forget(handle); + } + } + + unsafe fn push2(&self, handle: *mut Inner) { + // 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 + (*handle).next = AtomicPtr::new(ptr::null_mut()); + + // 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 `handle` + // to any other threads calling `push`. + let prev = self.head.swap(handle, AcqRel); + + // Release `handle` to the consume end. + (*prev).next.store(handle, Release); + } + + pub unsafe fn poll(&self) -> Poll { + let mut tail = self.tail.get(); + let mut next = (*tail).next.load(Acquire); + let stub = &*self.stub as *const _ as *mut _; + + if tail == stub { + if next.is_null() { + return Poll::Empty; + } + + self.tail.set(next); + tail = next; + next = (*next).next.load(Acquire); + } + + if !next.is_null() { + self.tail.set(next); + + // No ref_count inc is necessary here as this poll is paired + // with a `push` which "forgets" the handle. + return Poll::Data(Task { + ptr: tail, + }); + } + + if self.head.load(Acquire) != tail { + return Poll::Inconsistent; + } + + self.push2(stub); + + next = (*tail).next.load(Acquire); + + if !next.is_null() { + self.tail.set(next); + return Poll::Data(Task { + ptr: tail, + }); + } + + Poll::Inconsistent + } +} + +// ===== impl State ===== + +impl State { + /// Returns the initial task state. + /// + /// Tasks start in the scheduled state as they are immediately scheduled on + /// creation. + fn new() -> State { + State::Scheduled + } + + fn stub() -> State { + State::Idle + } +} + +impl From for State { + fn from(src: usize) -> Self { + use self::State::*; + + match src { + 0 => Idle, + 1 => Running, + 2 => Notified, + 3 => Scheduled, + 4 => Complete, + _ => unreachable!(), + } + } +} + +impl From for usize { + fn from(src: State) -> Self { + use self::State::*; + + match src { + Idle => 0, + Running => 1, + Notified => 2, + Scheduled => 3, + Complete => 4, + } + } +} diff --git a/tokio-threadpool/tests/threadpool.rs b/tokio-threadpool/tests/threadpool.rs new file mode 100644 index 000000000..d8b97997f --- /dev/null +++ b/tokio-threadpool/tests/threadpool.rs @@ -0,0 +1,331 @@ +extern crate tokio_threadpool; +extern crate tokio_executor; +extern crate futures; +extern crate env_logger; + +use tokio_threadpool::*; +use futures::{Poll, Sink, Stream, Async}; +use futures::future::{Future, lazy}; + +use std::cell::Cell; +use std::sync::{mpsc, Arc}; +use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT}; +use std::sync::atomic::Ordering::Relaxed; +use std::time::Duration; + +thread_local!(static FOO: Cell = Cell::new(0)); + +#[test] +fn natural_shutdown_simple_futures() { + let _ = ::env_logger::init(); + + for _ in 0..1_000 { + static NUM_INC: AtomicUsize = ATOMIC_USIZE_INIT; + static NUM_DEC: AtomicUsize = ATOMIC_USIZE_INIT; + + FOO.with(|f| { + f.set(1); + + let pool = Builder::new() + .around_worker(|w, _| { + NUM_INC.fetch_add(1, Relaxed); + w.run(); + NUM_DEC.fetch_add(1, Relaxed); + }) + .build(); + let tx = pool.sender().clone(); + + let a = { + let (t, rx) = mpsc::channel(); + tx.spawn(lazy(move || { + // Makes sure this runs on a worker thread + FOO.with(|f| assert_eq!(f.get(), 0)); + + t.send("one").unwrap(); + Ok(()) + })).unwrap(); + rx + }; + + let b = { + let (t, rx) = mpsc::channel(); + tx.spawn(lazy(move || { + // Makes sure this runs on a worker thread + FOO.with(|f| assert_eq!(f.get(), 0)); + + t.send("two").unwrap(); + Ok(()) + })).unwrap(); + rx + }; + + drop(tx); + + assert_eq!("one", a.recv().unwrap()); + assert_eq!("two", b.recv().unwrap()); + + // Wait for the pool to shutdown + pool.shutdown().wait().unwrap(); + + // Assert that at least one thread started + let num_inc = NUM_INC.load(Relaxed); + assert!(num_inc > 0); + + // Assert that all threads shutdown + let num_dec = NUM_DEC.load(Relaxed); + assert_eq!(num_inc, num_dec); + }); + } +} + +#[test] +fn force_shutdown_drops_futures() { + let _ = ::env_logger::init(); + + for _ in 0..1_000 { + let num_inc = Arc::new(AtomicUsize::new(0)); + let num_dec = Arc::new(AtomicUsize::new(0)); + let num_drop = Arc::new(AtomicUsize::new(0)); + + struct Never(Arc); + + impl Future for Never { + type Item = (); + type Error = (); + + fn poll(&mut self) -> Poll<(), ()> { + Ok(Async::NotReady) + } + } + + impl Drop for Never { + fn drop(&mut self) { + self.0.fetch_add(1, Relaxed); + } + } + + let a = num_inc.clone(); + let b = num_dec.clone(); + + let mut pool = Builder::new() + .around_worker(move |w, _| { + a.fetch_add(1, Relaxed); + w.run(); + b.fetch_add(1, Relaxed); + }) + .build(); + let mut tx = pool.sender().clone(); + + tx.spawn(Never(num_drop.clone())).unwrap(); + + // Wait for the pool to shutdown + pool.shutdown_now().wait().unwrap(); + + // Assert that only a single thread was spawned. + let a = num_inc.load(Relaxed); + assert!(a >= 1); + + // Assert that all threads shutdown + let b = num_dec.load(Relaxed); + assert_eq!(a, b); + + // Assert that the future was dropped + let c = num_drop.load(Relaxed); + assert_eq!(c, 1); + } +} + +#[test] +fn thread_shutdown_timeout() { + use std::sync::Mutex; + + let _ = ::env_logger::init(); + + let (shutdown_tx, shutdown_rx) = mpsc::channel(); + let (complete_tx, complete_rx) = mpsc::channel(); + + let t = Mutex::new(shutdown_tx); + + let pool = Builder::new() + .keep_alive(Some(Duration::from_millis(200))) + .around_worker(move |w, _| { + w.run(); + // There could be multiple threads here + let _ = t.lock().unwrap().send(()); + }) + .build(); + let tx = pool.sender().clone(); + + let t = complete_tx.clone(); + tx.spawn(lazy(move || { + t.send(()).unwrap(); + Ok(()) + })).unwrap(); + + // The future completes + complete_rx.recv().unwrap(); + + // The thread shuts down eventually + shutdown_rx.recv().unwrap(); + + // Futures can still be run + tx.spawn(lazy(move || { + complete_tx.send(()).unwrap(); + Ok(()) + })).unwrap(); + + complete_rx.recv().unwrap(); + + pool.shutdown().wait().unwrap(); +} + +#[test] +fn many_oneshot_futures() { + const NUM: usize = 10_000; + + let _ = ::env_logger::init(); + + for _ in 0..50 { + let pool = ThreadPool::new(); + let mut tx = pool.sender().clone(); + let cnt = Arc::new(AtomicUsize::new(0)); + + for _ in 0..NUM { + let cnt = cnt.clone(); + tx.spawn(lazy(move || { + cnt.fetch_add(1, Relaxed); + Ok(()) + })).unwrap(); + } + + // Wait for the pool to shutdown + pool.shutdown().wait().unwrap(); + + let num = cnt.load(Relaxed); + assert_eq!(num, NUM); + } +} + +#[test] +fn many_multishot_futures() { + use futures::sync::mpsc; + + const CHAIN: usize = 200; + const CYCLES: usize = 5; + const TRACKS: usize = 50; + + let _ = ::env_logger::init(); + + for _ in 0..50 { + let pool = ThreadPool::new(); + let mut pool_tx = pool.sender().clone(); + + let mut start_txs = Vec::with_capacity(TRACKS); + let mut final_rxs = Vec::with_capacity(TRACKS); + + for _ in 0..TRACKS { + let (start_tx, mut chain_rx) = mpsc::channel(10); + + for _ in 0..CHAIN { + let (next_tx, next_rx) = mpsc::channel(10); + + let rx = chain_rx + .map_err(|e| panic!("{:?}", e)); + + // Forward all the messages + pool_tx.spawn(next_tx + .send_all(rx) + .map(|_| ()) + .map_err(|e| panic!("{:?}", e)) + ).unwrap(); + + chain_rx = next_rx; + } + + // This final task cycles if needed + let (final_tx, final_rx) = mpsc::channel(10); + let cycle_tx = start_tx.clone(); + let mut rem = CYCLES; + + pool_tx.spawn(chain_rx.take(CYCLES as u64).for_each(move |msg| { + rem -= 1; + let send = if rem == 0 { + final_tx.clone().send(msg) + } else { + cycle_tx.clone().send(msg) + }; + + send.then(|res| { + res.unwrap(); + Ok(()) + }) + })).unwrap(); + + start_txs.push(start_tx); + final_rxs.push(final_rx); + } + + for start_tx in start_txs { + start_tx.send("ping").wait().unwrap(); + } + + for final_rx in final_rxs { + final_rx.wait().next().unwrap().unwrap(); + } + + // Shutdown the pool + pool.shutdown().wait().unwrap(); + } +} + +#[test] +fn global_executor_is_configured() { + let pool = ThreadPool::new(); + let tx = pool.sender().clone(); + + let (signal_tx, signal_rx) = mpsc::channel(); + + tx.spawn(lazy(move || { + tokio_executor::spawn(lazy(move || { + signal_tx.send(()).unwrap(); + Ok(()) + })); + + Ok(()) + })).unwrap(); + + signal_rx.recv().unwrap(); + + pool.shutdown().wait().unwrap(); +} + +#[test] +fn new_threadpool_is_idle() { + let pool = ThreadPool::new(); + pool.shutdown_on_idle().wait().unwrap(); +} + +#[test] +fn busy_threadpool_is_not_idle() { + use futures::sync::oneshot; + + let pool = ThreadPool::new(); + let tx = pool.sender().clone(); + + let (term_tx, term_rx) = oneshot::channel(); + + tx.spawn(term_rx.then(|_| { + Ok(()) + })).unwrap(); + + let mut idle = pool.shutdown_on_idle(); + + futures::lazy(|| { + assert!(idle.poll().unwrap().is_not_ready()); + Ok::<_, ()>(()) + }).wait().unwrap(); + + term_tx.send(()).unwrap(); + + idle.wait().unwrap(); +}