diff --git a/azure-pipelines.yml b/azure-pipelines.yml index 4743a66a4..cc50f3c88 100644 --- a/azure-pipelines.yml +++ b/azure-pipelines.yml @@ -44,13 +44,10 @@ jobs: displayName: Test build permutations rust: stable -# Run loom tests -- template: ci/azure-loom.yml +# Run miri tests +- template: ci/azure-miri.yml parameters: - name: loom - rust: stable - crates: - - tokio + name: miri # Try cross compiling - template: ci/azure-cross-compile.yml @@ -99,16 +96,25 @@ jobs: # name: tsan # rust: stable +# Run loom tests +- template: ci/azure-loom.yml + parameters: + name: loom + rust: stable + - template: ci/azure-deploy-docs.yml parameters: rust: stable dependsOn: - rustfmt + - docs - clippy - test_tokio - test_linux + - test_integration - test_build - loom + - miri - cross - minrust - check_features diff --git a/benches/Cargo.toml b/benches/Cargo.toml index 8c9ad1d39..f4a1d8fb7 100644 --- a/benches/Cargo.toml +++ b/benches/Cargo.toml @@ -17,3 +17,8 @@ harness = false name = "mpsc" path = "mpsc.rs" harness = false + +[[bench]] +name = "scheduler" +path = "scheduler.rs" +harness = false diff --git a/benches/scheduler.rs b/benches/scheduler.rs new file mode 100644 index 000000000..0562a1201 --- /dev/null +++ b/benches/scheduler.rs @@ -0,0 +1,152 @@ +//! Benchmark implementation details of the theaded scheduler. These benches are +//! intended to be used as a form of regression testing and not as a general +//! purpose benchmark demonstrating real-world performance. + +use tokio::runtime::{self, Runtime}; +use tokio::sync::oneshot; + +use bencher::{benchmark_group, benchmark_main, Bencher}; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::Relaxed; +use std::sync::{mpsc, Arc}; + +fn spawn_many(b: &mut Bencher) { + const NUM_SPAWN: usize = 10_000; + + let mut rt = rt(); + + let (tx, rx) = mpsc::sync_channel(1000); + let rem = Arc::new(AtomicUsize::new(0)); + + b.iter(|| { + rem.store(NUM_SPAWN, Relaxed); + + rt.block_on(async { + for _ in 0..NUM_SPAWN { + let tx = tx.clone(); + let rem = rem.clone(); + + tokio::spawn(async move { + if 1 == rem.fetch_sub(1, Relaxed) { + tx.send(()).unwrap(); + } + }); + } + + let _ = rx.recv().unwrap(); + }); + }); +} + +fn yield_many(b: &mut Bencher) { + const NUM_YIELD: usize = 1_000; + const TASKS: usize = 200; + + let rt = rt(); + + let (tx, rx) = mpsc::sync_channel(TASKS); + + b.iter(move || { + for _ in 0..TASKS { + let tx = tx.clone(); + + rt.spawn(async move { + for _ in 0..NUM_YIELD { + tokio::task::yield_now().await; + } + + tx.send(()).unwrap(); + }); + } + + for _ in 0..TASKS { + let _ = rx.recv().unwrap(); + } + }); +} + +fn ping_pong(b: &mut Bencher) { + const NUM_PINGS: usize = 1_000; + + let mut rt = rt(); + + let (done_tx, done_rx) = mpsc::sync_channel(1000); + let rem = Arc::new(AtomicUsize::new(0)); + + b.iter(|| { + let done_tx = done_tx.clone(); + let rem = rem.clone(); + rem.store(NUM_PINGS, Relaxed); + + rt.block_on(async { + tokio::spawn(async move { + for _ in 0..NUM_PINGS { + let rem = rem.clone(); + let done_tx = done_tx.clone(); + + tokio::spawn(async move { + let (tx1, rx1) = oneshot::channel(); + let (tx2, rx2) = oneshot::channel(); + + tokio::spawn(async move { + rx1.await.unwrap(); + tx2.send(()).unwrap(); + }); + + tx1.send(()).unwrap(); + rx2.await.unwrap(); + + if 1 == rem.fetch_sub(1, Relaxed) { + done_tx.send(()).unwrap(); + } + }); + } + }); + + done_rx.recv().unwrap(); + }); + }); +} + +fn chained_spawn(b: &mut Bencher) { + const ITER: usize = 1_000; + + let mut rt = rt(); + + fn iter(done_tx: mpsc::SyncSender<()>, n: usize) { + if n == 0 { + done_tx.send(()).unwrap(); + } else { + tokio::spawn(async move { + iter(done_tx, n - 1); + }); + } + } + + let (done_tx, done_rx) = mpsc::sync_channel(1000); + + b.iter(move || { + let done_tx = done_tx.clone(); + + rt.block_on(async { + tokio::spawn(async move { + iter(done_tx, ITER); + }); + + done_rx.recv().unwrap(); + }); + }); +} + +fn rt() -> Runtime { + runtime::Builder::new() + .threaded_scheduler() + .core_threads(4) + .enable_all() + .build() + .unwrap() +} + +benchmark_group!(scheduler, spawn_many, ping_pong, yield_many, chained_spawn,); + +benchmark_main!(scheduler); diff --git a/ci/azure-check-features.yml b/ci/azure-check-features.yml index a80af0d39..f5985843e 100644 --- a/ci/azure-check-features.yml +++ b/ci/azure-check-features.yml @@ -6,7 +6,7 @@ jobs: Linux: vmImage: ubuntu-16.04 MacOS: - vmImage: macOS-10.13 + vmImage: macos-latest Windows: vmImage: vs2017-win2016 pool: diff --git a/ci/azure-install-rust.yml b/ci/azure-install-rust.yml index 2b4feb2a9..4cf5ca3f9 100644 --- a/ci/azure-install-rust.yml +++ b/ci/azure-install-rust.yml @@ -2,6 +2,13 @@ steps: # Linux and macOS. - script: | set -e + + if [ "$RUSTUP_TOOLCHAIN" == "nightly" ]; then + echo "++ getting latest miri version" + export RUSTUP_TOOLCHAIN="nightly-$(curl -s https://rust-lang.github.io/rustup-components-history/x86_64-unknown-linux-gnu/miri)" + echo "$RUSTUP_TOOLCHAIN" + fi + curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal --default-toolchain none export PATH=$PATH:$HOME/.cargo/bin rustup toolchain install $RUSTUP_TOOLCHAIN diff --git a/ci/azure-loom.yml b/ci/azure-loom.yml index fdfb86705..001aedec2 100644 --- a/ci/azure-loom.yml +++ b/ci/azure-loom.yml @@ -1,6 +1,18 @@ jobs: - job: ${{ parameters.name }} displayName: Loom tests + strategy: + matrix: + rest: + scope: --skip loom_pool + pool_group_a: + scope: loom_pool::group_a + pool_group_b: + scope: loom_pool::group_b + pool_group_c: + scope: loom_pool::group_c + pool_group_d: + scope: loom_pool::group_d pool: vmImage: ubuntu-16.04 @@ -9,10 +21,9 @@ jobs: parameters: rust_version: ${{ parameters.rust }} - - ${{ each crate in parameters.crates }}: - - script: RUSTFLAGS="--cfg loom" cargo test --lib --release --features "full" -- --test-threads=1 --nocapture - env: - LOOM_MAX_PREEMPTIONS: 1 - CI: 'True' - displayName: test ${{ crate }} - workingDirectory: $(Build.SourcesDirectory)/${{ crate }} + - script: RUSTFLAGS="--cfg loom" cargo test --lib --release --features "full" -- --nocapture $(scope) + env: + LOOM_MAX_PREEMPTIONS: 2 + CI: 'True' + displayName: $(scope) + workingDirectory: $(Build.SourcesDirectory)/tokio diff --git a/ci/azure-miri.yml b/ci/azure-miri.yml new file mode 100644 index 000000000..fb886edc7 --- /dev/null +++ b/ci/azure-miri.yml @@ -0,0 +1,23 @@ +jobs: +- job: ${{ parameters.name }} + displayName: Miri + pool: + vmImage: ubuntu-16.04 + + steps: + - template: azure-install-rust.yml + parameters: + rust_version: nightly + + - script: | + rustup component add miri + cargo miri setup + rm -rf $(Build.SourcesDirectory)/tokio/tests + displayName: Install miri + + # TODO: enable all tests once they pass + - script: cargo miri test --features rt-core,rt-threaded,rt-util,sync -- -- task + env: + CI: 'True' + displayName: cargo miri test + workingDirectory: $(Build.SourcesDirectory)/tokio diff --git a/ci/azure-test-integration.yml b/ci/azure-test-integration.yml index fc45429aa..f498a6496 100644 --- a/ci/azure-test-integration.yml +++ b/ci/azure-test-integration.yml @@ -6,7 +6,7 @@ jobs: Linux: vmImage: ubuntu-16.04 MacOS: - vmImage: macOS-10.13 + vmImage: macos-latest Windows: vmImage: vs2017-win2016 pool: diff --git a/ci/azure-test-stable.yml b/ci/azure-test-stable.yml index a924c9289..bc93febd0 100644 --- a/ci/azure-test-stable.yml +++ b/ci/azure-test-stable.yml @@ -8,7 +8,7 @@ jobs: ${{ if parameters.cross }}: MacOS: - vmImage: macOS-10.13 + vmImage: macos-latest Windows: vmImage: vs2017-win2016 pool: diff --git a/tokio/src/lib.rs b/tokio/src/lib.rs index 843a1e64b..1498c2b5f 100644 --- a/tokio/src/lib.rs +++ b/tokio/src/lib.rs @@ -382,10 +382,6 @@ cfg_macros! { } } -// Tests -#[cfg(test)] -mod tests; - // TODO: rm #[cfg(feature = "io-util")] #[cfg(test)] diff --git a/tokio/src/loom/std/alloc.rs b/tokio/src/loom/std/alloc.rs deleted file mode 100644 index 25b199b1b..000000000 --- a/tokio/src/loom/std/alloc.rs +++ /dev/null @@ -1,18 +0,0 @@ -#[derive(Debug)] -pub(crate) struct Track { - value: T, -} - -impl Track { - pub(crate) fn new(value: T) -> Track { - Track { value } - } - - pub(crate) fn get_mut(&mut self) -> &mut T { - &mut self.value - } - - pub(crate) fn into_inner(self) -> T { - self.value - } -} diff --git a/tokio/src/loom/std/causal_cell.rs b/tokio/src/loom/std/causal_cell.rs index c4917e5f8..8300437a1 100644 --- a/tokio/src/loom/std/causal_cell.rs +++ b/tokio/src/loom/std/causal_cell.rs @@ -18,15 +18,6 @@ impl CausalCell { f(self.0.get()) } - pub(crate) fn with_unchecked(&self, f: F) -> R - where - F: FnOnce(*const T) -> R, - { - f(self.0.get()) - } - - pub(crate) fn check(&self) {} - pub(crate) fn with_deferred(&self, f: F) -> (R, CausalCheck) where F: FnOnce(*const T) -> R, diff --git a/tokio/src/loom/std/mod.rs b/tokio/src/loom/std/mod.rs index e4bae357b..a56d778ab 100644 --- a/tokio/src/loom/std/mod.rs +++ b/tokio/src/loom/std/mod.rs @@ -5,8 +5,6 @@ mod atomic_u64; mod atomic_usize; mod causal_cell; -pub(crate) mod alloc; - pub(crate) mod cell { pub(crate) use super::causal_cell::{CausalCell, CausalCheck}; } diff --git a/tokio/src/macros/assert.rs b/tokio/src/macros/assert.rs deleted file mode 100644 index 4b1cf272e..000000000 --- a/tokio/src/macros/assert.rs +++ /dev/null @@ -1,18 +0,0 @@ -/// Asserts option is some -macro_rules! assert_some { - ($e:expr) => {{ - match $e { - Some(v) => v, - _ => panic!("expected some, was none"), - } - }}; -} - -/// Asserts option is none -macro_rules! assert_none { - ($e:expr) => {{ - if let Some(v) = $e { - panic!("expected none, was {:?}", v); - } - }}; -} diff --git a/tokio/src/macros/mod.rs b/tokio/src/macros/mod.rs index a37b3e49d..2643c3601 100644 --- a/tokio/src/macros/mod.rs +++ b/tokio/src/macros/mod.rs @@ -1,9 +1,5 @@ #![cfg_attr(not(feature = "full"), allow(unused_macros))] -#[macro_use] -#[cfg(test)] -mod assert; - #[macro_use] mod cfg; @@ -19,6 +15,10 @@ mod ready; #[macro_use] mod thread_local; +#[macro_use] +#[cfg(feature = "rt-core")] +pub(crate) mod scoped_tls; + cfg_macros! { #[macro_use] mod select; diff --git a/tokio/src/macros/scoped_tls.rs b/tokio/src/macros/scoped_tls.rs new file mode 100644 index 000000000..666f382b2 --- /dev/null +++ b/tokio/src/macros/scoped_tls.rs @@ -0,0 +1,80 @@ +use crate::loom::thread::LocalKey; + +use std::cell::Cell; +use std::marker; + +/// Set a reference as a thread-local +#[macro_export] +macro_rules! scoped_thread_local { + ($(#[$attrs:meta])* $vis:vis static $name:ident: $ty:ty) => ( + $(#[$attrs])* + $vis static $name: $crate::macros::scoped_tls::ScopedKey<$ty> + = $crate::macros::scoped_tls::ScopedKey { + inner: { + thread_local!(static FOO: ::std::cell::Cell<*const ()> = { + std::cell::Cell::new(::std::ptr::null()) + }); + &FOO + }, + _marker: ::std::marker::PhantomData, + }; + ) +} + +/// Type representing a thread local storage key corresponding to a reference +/// to the type parameter `T`. +pub(crate) struct ScopedKey { + #[doc(hidden)] + pub(crate) inner: &'static LocalKey>, + #[doc(hidden)] + pub(crate) _marker: marker::PhantomData, +} + +unsafe impl Sync for ScopedKey {} + +impl ScopedKey { + /// Inserts a value into this scoped thread local storage slot for a + /// duration of a closure. + pub(crate) fn set(&'static self, t: &T, f: F) -> R + where + F: FnOnce() -> R, + { + struct Reset { + key: &'static LocalKey>, + val: *const (), + } + + impl Drop for Reset { + fn drop(&mut self) { + self.key.with(|c| c.set(self.val)); + } + } + + let prev = self.inner.with(|c| { + let prev = c.get(); + c.set(t as *const _ as *const ()); + prev + }); + + let _reset = Reset { + key: self.inner, + val: prev, + }; + + f() + } + + /// Gets a value out of this scoped variable. + pub(crate) fn with(&'static self, f: F) -> R + where + F: FnOnce(Option<&T>) -> R, + { + let val = self.inner.with(|c| c.get()); + + if val.is_null() { + f(None) + } else { + unsafe { f(Some(&*(val as *const T))) } + } + } +} diff --git a/tokio/src/park/thread.rs b/tokio/src/park/thread.rs index 12ef9717e..a8cdf1432 100644 --- a/tokio/src/park/thread.rs +++ b/tokio/src/park/thread.rs @@ -129,6 +129,10 @@ impl Inner { return; } + if dur == Duration::from_millis(0) { + return; + } + let m = self.mutex.lock().unwrap(); match self.state.compare_exchange(EMPTY, PARKED, SeqCst, SeqCst) { diff --git a/tokio/src/runtime/basic_scheduler.rs b/tokio/src/runtime/basic_scheduler.rs index f625920d7..a494f9e3d 100644 --- a/tokio/src/runtime/basic_scheduler.rs +++ b/tokio/src/runtime/basic_scheduler.rs @@ -1,46 +1,33 @@ use crate::park::{Park, Unpark}; -use crate::task::{self, queue::MpscQueues, JoinHandle, Schedule, ScheduleSendOnly, Task}; +use crate::runtime; +use crate::runtime::task::{self, JoinHandle, Schedule, Task}; +use crate::util::linked_list::LinkedList; +use crate::util::{waker_ref, Wake}; -use std::cell::Cell; +use std::cell::RefCell; +use std::collections::VecDeque; use std::fmt; use std::future::Future; -use std::mem::ManuallyDrop; -use std::ptr; -use std::sync::Arc; -use std::task::{RawWaker, RawWakerVTable, Waker}; +use std::sync::{Arc, Mutex}; +use std::task::Poll::Ready; use std::time::Duration; /// Executes tasks on the current thread -#[derive(Debug)] pub(crate) struct BasicScheduler

where P: Park, { - /// Scheduler component - scheduler: Arc, + /// Scheduler run queue + /// + /// When the scheduler is executed, the queue is removed from `self` and + /// moved into `Context`. + /// + /// This indirection is to allow `BasicScheduler` to be `Send`. + tasks: Option, - /// Local state - local: LocalState

, -} + /// Sendable task spawner + spawner: Spawner, -#[derive(Debug, Clone)] -pub(crate) struct Spawner { - scheduler: Arc, -} - -/// The scheduler component. -pub(super) struct SchedulerPriv { - queues: MpscQueues, - /// Unpark the blocked thread - unpark: Box, -} - -unsafe impl Send for SchedulerPriv {} -unsafe impl Sync for SchedulerPriv {} - -/// Local state -#[derive(Debug)] -struct LocalState

{ /// Current tick tick: u8, @@ -48,33 +35,76 @@ struct LocalState

{ park: P, } +#[derive(Clone)] +pub(crate) struct Spawner { + shared: Arc, +} + +struct Tasks { + /// Collection of all active tasks spawned onto this executor. + owned: LinkedList>>, + + /// Local run queue. + /// + /// Tasks notified from the current thread are pushed into this queue. + queue: VecDeque>>, +} + +/// Scheduler state shared between threads. +struct Shared { + /// Remote run queue + queue: Mutex>>>, + + /// Unpark the blocked thread + unpark: Box, +} + +/// Thread-local context +struct Context { + /// Shared scheduler state + shared: Arc, + + /// Local queue + tasks: RefCell, +} + +/// Initial queue capacity +const INITIAL_CAPACITY: usize = 64; + /// Max number of tasks to poll per tick. const MAX_TASKS_PER_TICK: usize = 61; -thread_local! { - static ACTIVE: Cell<*const SchedulerPriv> = Cell::new(ptr::null()) -} +/// How often ot check the remote queue first +const REMOTE_FIRST_INTERVAL: u8 = 31; + +// Tracks the current BasicScheduler +scoped_thread_local!(static CURRENT: Context); impl

BasicScheduler

where P: Park, { pub(crate) fn new(park: P) -> BasicScheduler

{ - let unpark = park.unpark(); + let unpark = Box::new(park.unpark()); BasicScheduler { - scheduler: Arc::new(SchedulerPriv { - queues: MpscQueues::new(), - unpark: Box::new(unpark), + tasks: Some(Tasks { + owned: LinkedList::new(), + queue: VecDeque::with_capacity(INITIAL_CAPACITY), }), - local: LocalState { tick: 0, park }, + spawner: Spawner { + shared: Arc::new(Shared { + queue: Mutex::new(VecDeque::with_capacity(INITIAL_CAPACITY)), + unpark: unpark as Box, + }), + }, + tick: 0, + park, } } - pub(crate) fn spawner(&self) -> Spawner { - Spawner { - scheduler: self.scheduler.clone(), - } + pub(crate) fn spawner(&self) -> &Spawner { + &self.spawner } /// Spawns a future onto the thread pool @@ -83,74 +113,146 @@ where F: Future + Send + 'static, F::Output: Send + 'static, { - let (task, handle) = task::joinable(future); - self.scheduler.schedule(task, true); - handle + self.spawner.spawn(future) } - pub(crate) fn block_on(&mut self, mut future: F) -> F::Output + pub(crate) fn block_on(&mut self, future: F) -> F::Output where F: Future, { - use crate::runtime; - use std::pin::Pin; - use std::task::Context; - use std::task::Poll::Ready; + enter(self, |scheduler, context| { + let _enter = runtime::enter(); + let waker = waker_ref(&scheduler.spawner.shared); + let mut cx = std::task::Context::from_waker(&waker); - let local = &mut self.local; - let scheduler = &*self.scheduler; + pin!(future); - struct Guard { - old: *const SchedulerPriv, - } + 'outer: loop { + if let Ready(v) = future.as_mut().poll(&mut cx) { + return v; + } - impl Drop for Guard { - fn drop(&mut self) { - ACTIVE.with(|cell| cell.set(self.old)); + for _ in 0..MAX_TASKS_PER_TICK { + // Get and increment the current tick + let tick = scheduler.tick; + scheduler.tick = scheduler.tick.wrapping_add(1); + + let next = if tick % REMOTE_FIRST_INTERVAL == 0 { + scheduler + .spawner + .pop() + .or_else(|| context.tasks.borrow_mut().queue.pop_front()) + } else { + context + .tasks + .borrow_mut() + .queue + .pop_front() + .or_else(|| scheduler.spawner.pop()) + }; + + match next { + Some(task) => task.run(), + None => { + // Park until the thread is signaled + scheduler.park.park().ok().expect("failed to park"); + + // Try polling the `block_on` future next + continue 'outer; + } + } + } + + // Yield to the park, this drives the timer and pulls any pending + // I/O events. + scheduler + .park + .park_timeout(Duration::from_millis(0)) + .ok() + .expect("failed to park"); } - } - - // Track the current scheduler - let _guard = ACTIVE.with(|cell| { - let guard = Guard { old: cell.get() }; - - cell.set(scheduler as *const SchedulerPriv); - - guard - }); - - let mut _enter = runtime::enter(); - - let raw_waker = RawWaker::new( - scheduler as *const SchedulerPriv as *const (), - &RawWakerVTable::new(sched_clone_waker, sched_noop, sched_wake_by_ref, sched_noop), - ); - - let waker = ManuallyDrop::new(unsafe { Waker::from_raw(raw_waker) }); - let mut cx = Context::from_waker(&waker); - - // `block_on` takes ownership of `f`. Once it is pinned here, the - // original `f` binding can no longer be accessed, making the - // pinning safe. - let mut future = unsafe { Pin::new_unchecked(&mut future) }; - - loop { - if let Ready(v) = future.as_mut().poll(&mut cx) { - return v; - } - - scheduler.tick(local); - - // Maintenance work - unsafe { - // safety: this function is safe to call only from the - // thread the basic scheduler is running on (which we are). - scheduler.queues.drain_pending_drop(); - } - } + }) } } +/// Enter the scheduler context. This sets the queue and other necessary +/// scheduler state in the thread-local +fn enter(scheduler: &mut BasicScheduler

, f: F) -> R +where + F: FnOnce(&mut BasicScheduler

, &Context) -> R, + P: Park, +{ + // Ensures the run queue is placed back in the `BasicScheduler` instance + // once `block_on` returns.` + struct Guard<'a, P: Park> { + context: Option, + scheduler: &'a mut BasicScheduler

, + } + + impl Drop for Guard<'_, P> { + fn drop(&mut self) { + let Context { tasks, .. } = self.context.take().expect("context missing"); + self.scheduler.tasks = Some(tasks.into_inner()); + } + } + + // Remove `tasks` from `self` and place it in a `Context`. + let tasks = scheduler.tasks.take().expect("invalid state"); + + let guard = Guard { + context: Some(Context { + shared: scheduler.spawner.shared.clone(), + tasks: RefCell::new(tasks), + }), + scheduler, + }; + + let context = guard.context.as_ref().unwrap(); + let scheduler = &mut *guard.scheduler; + + CURRENT.set(context, || f(scheduler, context)) +} + +impl

Drop for BasicScheduler

+where + P: Park, +{ + fn drop(&mut self) { + enter(self, |scheduler, context| { + // Loop required here to ensure borrow is dropped between iterations + #[allow(clippy::while_let_loop)] + loop { + let task = match context.tasks.borrow_mut().owned.pop_back() { + Some(task) => task, + None => break, + }; + + task.shutdown(); + } + + // Drain local queue + for task in context.tasks.borrow_mut().queue.drain(..) { + task.shutdown(); + } + + // Drain remote queue + for task in scheduler.spawner.shared.queue.lock().unwrap().drain(..) { + task.shutdown(); + } + + assert!(context.tasks.borrow().owned.is_empty()); + }); + } +} + +impl fmt::Debug for BasicScheduler

{ + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("BasicScheduler").finish() + } +} + +// ===== impl Spawner ===== + impl Spawner { /// Spawns a future onto the thread pool pub(crate) fn spawn(&self, future: F) -> JoinHandle @@ -159,177 +261,66 @@ impl Spawner { F::Output: Send + 'static, { let (task, handle) = task::joinable(future); - self.scheduler.schedule(task, true); + self.shared.schedule(task); handle } -} -// === impl SchedulerPriv === - -impl SchedulerPriv { - fn tick(&self, local: &mut LocalState) { - for _ in 0..MAX_TASKS_PER_TICK { - // Get the current tick - let tick = local.tick; - - // Increment the tick - local.tick = tick.wrapping_add(1); - let next = unsafe { - // safety: this function is safe to call only from the - // thread the basic scheduler is running on. The `LocalState` - // parameter to this method implies that we are on that thread. - self.queues.next_task(tick) - }; - - let task = match next { - Some(task) => task, - None => { - local.park.park().ok().expect("failed to park"); - return; - } - }; - - if let Some(task) = task.run(&mut || Some(self.into())) { - unsafe { - // safety: this function is safe to call only from the - // thread the basic scheduler is running on. The `LocalState` - // parameter to this method implies that we are on that thread. - self.queues.push_local(task); - } - } - } - - local - .park - .park_timeout(Duration::from_millis(0)) - .ok() - .expect("failed to park"); - } - - /// Schedule the provided task on the scheduler. - /// - /// If this scheduler is the `ACTIVE` scheduler, enqueue this task on the local queue, otherwise - /// the task is enqueued on the remote queue. - fn schedule(&self, task: Task, spawn: bool) { - let is_current = ACTIVE.with(|cell| cell.get() == self as *const SchedulerPriv); - - if is_current { - unsafe { - // safety: this function is safe to call only from the - // thread the basic scheduler is running on. If `is_current` is - // then we are on that thread. - self.queues.push_local(task) - }; - } else { - let mut lock = self.queues.remote(); - lock.schedule(task, spawn); - - // while locked, call unpark - self.unpark.unpark(); - - drop(lock); - } + fn pop(&self) -> Option>> { + self.shared.queue.lock().unwrap().pop_front() } } -impl Schedule for SchedulerPriv { - fn bind(&self, task: &Task) { - unsafe { - // safety: `Queues::add_task` is only safe to call from the thread - // that owns the queues (the thread the scheduler is running on). - // `Scheduler::bind` is called when polling a task that - // doesn't have a scheduler set. We will only poll new tasks from - // the thread that the scheduler is running on. Therefore, this is - // safe to call. - self.queues.add_task(task); - } - } - - fn release(&self, task: Task) { - self.queues.release_remote(task); - } - - fn release_local(&self, task: &Task) { - unsafe { - // safety: `Scheduler::release_local` is only called from the - // thread that the scheduler is running on. The `Schedule` trait's - // contract is that releasing a task from another thread should call - // `release` rather than `release_local`. - self.queues.release_local(task); - } - } - - fn schedule(&self, task: Task) { - SchedulerPriv::schedule(self, task, false); - } -} - -impl ScheduleSendOnly for SchedulerPriv {} - -impl

Drop for BasicScheduler

-where - P: Park, -{ - fn drop(&mut self) { - unsafe { - // safety: the `Drop` impl owns the scheduler's queues. these fields - // will only be accessed when running the scheduler, and it can no - // longer be run, since we are in the process of dropping it. - - // Shut down the task queues. - self.scheduler.queues.shutdown(); - } - - // Wait until all tasks have been released. - loop { - unsafe { - self.scheduler.queues.drain_pending_drop(); - self.scheduler.queues.drain_queues(); - - if !self.scheduler.queues.has_tasks_remaining() { - break; - } - - self.local.park.park().ok().expect("park failed"); - } - } - } -} - -impl fmt::Debug for SchedulerPriv { +impl fmt::Debug for Spawner { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Scheduler") - .field("queues", &self.queues) - .finish() + fmt.debug_struct("Spawner").finish() } } -unsafe fn sched_clone_waker(ptr: *const ()) -> RawWaker { - let s1 = ManuallyDrop::new(Arc::from_raw(ptr as *const SchedulerPriv)); +// ===== impl Shared ===== - #[allow(clippy::redundant_clone)] - let s2 = s1.clone(); +impl Schedule for Arc { + fn bind(task: Task) -> Arc { + CURRENT.with(|maybe_cx| { + let cx = maybe_cx.expect("scheduler context missing"); + cx.tasks.borrow_mut().owned.push_front(task); + cx.shared.clone() + }) + } - RawWaker::new( - &**s2 as *const SchedulerPriv as *const (), - &RawWakerVTable::new(sched_clone_waker, sched_wake, sched_wake_by_ref, sched_drop), - ) + fn release(&self, task: &Task) -> Option> { + use std::ptr::NonNull; + + CURRENT.with(|maybe_cx| { + let cx = maybe_cx.expect("scheduler context missing"); + + // safety: the task is inserted in the list in `bind`. + unsafe { + let ptr = NonNull::from(task.header()); + cx.tasks.borrow_mut().owned.remove(ptr) + } + }) + } + + fn schedule(&self, task: task::Notified) { + CURRENT.with(|maybe_cx| match maybe_cx { + Some(cx) if Arc::ptr_eq(self, &cx.shared) => { + cx.tasks.borrow_mut().queue.push_back(task); + } + _ => { + self.queue.lock().unwrap().push_back(task); + self.unpark.unpark(); + } + }); + } } -unsafe fn sched_wake(ptr: *const ()) { - let scheduler = Arc::from_raw(ptr as *const SchedulerPriv); - scheduler.unpark.unpark(); -} +impl Wake for Shared { + fn wake(self: Arc) { + Wake::wake_by_ref(&self) + } -unsafe fn sched_wake_by_ref(ptr: *const ()) { - let scheduler = ManuallyDrop::new(Arc::from_raw(ptr as *const SchedulerPriv)); - scheduler.unpark.unpark(); -} - -unsafe fn sched_drop(ptr: *const ()) { - let _ = Arc::from_raw(ptr as *const SchedulerPriv); -} - -unsafe fn sched_noop(_ptr: *const ()) { - unreachable!(); + /// Wake by reference + fn wake_by_ref(arc_self: &Arc) { + arc_self.unpark.unpark(); + } } diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 0b9d2209c..a3b208d17 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -5,8 +5,8 @@ use crate::loom::thread; use crate::runtime::blocking::schedule::NoopSchedule; use crate::runtime::blocking::shutdown; use crate::runtime::blocking::task::BlockingTask; +use crate::runtime::task::{self, JoinHandle}; use crate::runtime::{Builder, Callback, Handle}; -use crate::task::{self, JoinHandle}; use std::collections::VecDeque; use std::fmt; @@ -53,7 +53,7 @@ struct Shared { shutdown_tx: Option, } -type Task = task::Task; +type Task = task::Notified; const KEEP_ALIVE: Duration = Duration::from_secs(10); @@ -227,7 +227,7 @@ impl Inner { // BUSY while let Some(task) = shared.queue.pop_front() { drop(shared); - run_task(task); + task.run(); shared = self.shared.lock().unwrap(); } @@ -305,9 +305,3 @@ impl fmt::Debug for Spawner { fmt.debug_struct("blocking::Spawner").finish() } } - -fn run_task(f: Task) { - let scheduler: &'static NoopSchedule = &NoopSchedule; - let res = f.run(|| Some(scheduler.into())); - assert!(res.is_none()); -} diff --git a/tokio/src/runtime/blocking/schedule.rs b/tokio/src/runtime/blocking/schedule.rs index 5d2cd5f53..e10778d53 100644 --- a/tokio/src/runtime/blocking/schedule.rs +++ b/tokio/src/runtime/blocking/schedule.rs @@ -1,20 +1,24 @@ -use crate::task::{Schedule, ScheduleSendOnly, Task}; +use crate::runtime::task::{self, Task}; /// `task::Schedule` implementation that does nothing. This is unique to the /// blocking scheduler as tasks scheduled are not really futures but blocking /// operations. +/// +/// We avoid storing the task by forgetting it in `bind` and re-materializing it +/// in `release. pub(super) struct NoopSchedule; -impl Schedule for NoopSchedule { - fn bind(&self, _task: &Task) {} +impl task::Schedule for NoopSchedule { + fn bind(_task: Task) -> NoopSchedule { + // Do nothing w/ the task + NoopSchedule + } - fn release(&self, _task: Task) {} + fn release(&self, _task: &Task) -> Option> { + None + } - fn release_local(&self, _task: &Task) {} - - fn schedule(&self, _task: Task) { + fn schedule(&self, _task: task::Notified) { unreachable!(); } } - -impl ScheduleSendOnly for NoopSchedule {} diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index dce265c03..cfde99825 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -425,7 +425,7 @@ cfg_rt_core! { // the reactor to generate some new stimuli for the futures to continue // in their life. let scheduler = BasicScheduler::new(driver); - let spawner = Spawner::Basic(scheduler.spawner()); + let spawner = Spawner::Basic(scheduler.spawner().clone()); // Blocking pool let blocking_pool = blocking::create_blocking_pool(self, self.max_threads); @@ -470,7 +470,7 @@ cfg_rt_threaded! { let (io_driver, io_handle) = io::create_driver(self.enable_io)?; let (driver, time_handle) = time::create_driver(self.enable_time, io_driver, clock.clone()); - let (scheduler, workers) = ThreadPool::new(core_threads, Parker::new(driver)); + let (scheduler, launch) = ThreadPool::new(core_threads, Parker::new(driver)); let spawner = Spawner::ThreadPool(scheduler.spawner().clone()); // Create the blocking pool @@ -487,7 +487,7 @@ cfg_rt_threaded! { }; // Spawn the thread pool workers - workers.spawn(&handle); + handle.enter(|| launch.launch()); Ok(Runtime { kind: Kind::ThreadPool(scheduler), diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 6922ef59b..3aafc4060 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -187,11 +187,14 @@ #[cfg(test)] #[macro_use] mod tests; + pub(crate) mod context; cfg_rt_core! { mod basic_scheduler; use basic_scheduler::BasicScheduler; + + pub(crate) mod task; } mod blocking; @@ -215,7 +218,7 @@ mod io; cfg_rt_threaded! { mod park; - use park::{Parker, Unparker}; + use park::Parker; } mod shell; @@ -334,7 +337,7 @@ impl Runtime { /// [threaded scheduler]: index.html#threaded-scheduler /// [basic scheduler]: index.html#basic-scheduler /// [runtime builder]: crate::runtime::Builder - pub fn new() -> io::Result { + pub fn new() -> io::Result { #[cfg(feature = "rt-threaded")] let ret = Builder::new().threaded_scheduler().enable_all().build(); diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs new file mode 100644 index 000000000..43e6b4716 --- /dev/null +++ b/tokio/src/runtime/task/core.rs @@ -0,0 +1,280 @@ +use crate::loom::cell::CausalCell; +use crate::runtime::task::raw::{self, Vtable}; +use crate::runtime::task::state::State; +use crate::runtime::task::waker::waker_ref; +use crate::runtime::task::{Notified, Schedule, Task}; +use crate::util::linked_list; + +use std::cell::UnsafeCell; +use std::future::Future; +use std::pin::Pin; +use std::ptr::NonNull; +use std::task::{Context, Poll, Waker}; + +/// The task cell. Contains the components of the task. +/// +/// It is critical for `Header` to be the first field as the task structure will +/// be referenced by both *mut Cell and *mut Header. +#[repr(C)] +pub(super) struct Cell { + /// Hot task state data + pub(super) header: Header, + + /// Either the future or output, depending on the execution stage. + pub(super) core: Core, + + /// Cold data + pub(super) trailer: Trailer, +} + +/// The core of the task. +/// +/// Holds the future or output, depending on the stage of execution. +pub(super) struct Core { + /// Scheduler used to drive this future + pub(super) scheduler: CausalCell>, + + /// Either the future or the output + pub(super) stage: CausalCell>, +} + +/// Crate public as this is also needed by the pool. +#[repr(C)] +pub(crate) struct Header { + /// Task state + pub(super) state: State, + + pub(crate) owned: UnsafeCell>, + + /// Pointer to next task, used with the injection queue + pub(crate) queue_next: UnsafeCell>>, + + /// Pointer to the next task in the transfer stack + pub(super) stack_next: UnsafeCell>>, + + /// Table of function pointers for executing actions on the task. + pub(super) vtable: &'static Vtable, +} + +unsafe impl Send for Header {} +unsafe impl Sync for Header {} + +/// Cold data is stored after the future. +pub(super) struct Trailer { + /// Consumer task waiting on completion of this task. + pub(super) waker: CausalCell>, +} + +/// Either the future or the output. +pub(super) enum Stage { + Running(T), + Finished(super::Result), + Consumed, +} + +impl Cell { + /// Allocates a new task cell, containing the header, trailer, and core + /// structures. + pub(super) fn new(future: T, state: State) -> Box> { + Box::new(Cell { + header: Header { + state, + owned: UnsafeCell::new(linked_list::Pointers::new()), + queue_next: UnsafeCell::new(None), + stack_next: UnsafeCell::new(None), + vtable: raw::vtable::(), + }, + core: Core { + scheduler: CausalCell::new(None), + stage: CausalCell::new(Stage::Running(future)), + }, + trailer: Trailer { + waker: CausalCell::new(None), + }, + }) + } +} + +impl Core { + /// If needed, bind a scheduler to the task. + /// + /// This only happens on the first poll. + pub(super) fn bind_scheduler(&self, task: Task) { + use std::mem::ManuallyDrop; + + // TODO: it would be nice to not have to wrap with a ManuallyDrop + let task = ManuallyDrop::new(task); + + // This function may be called concurrently, but the __first__ time it + // is called, the caller has unique access to this field. All subsequent + // concurrent calls will be via the `Waker`, which will "happens after" + // the first poll. + // + // In other words, it is always safe to read the field and it is safe to + // write to the field when it is `None`. + if self.is_bound() { + return; + } + + // Bind the task to the scheduler + let scheduler = S::bind(ManuallyDrop::into_inner(task)); + + // Safety: As `scheduler` is not set, this is the first poll + self.scheduler.with_mut(|ptr| unsafe { + *ptr = Some(scheduler); + }); + } + + /// Returns true if the task is bound to a scheduler. + pub(super) fn is_bound(&self) -> bool { + // Safety: never called concurrently w/ a mutation. + self.scheduler.with(|ptr| unsafe { (*ptr).is_some() }) + } + + /// Poll the future + /// + /// # Safety + /// + /// The caller must ensure it is safe to mutate the `state` field. This + /// requires ensuring mutal exclusion between any concurrent thread that + /// might modify the future or output field. + /// + /// The mutual exclusion is implemented by `Harness` and the `Lifecycle` + /// component of the task state. + /// + /// `self` must also be pinned. This is handled by storing the task on the + /// heap. + pub(super) fn poll(&self, header: &Header) -> Poll { + let res = { + self.stage.with_mut(|ptr| { + // Safety: The caller ensures mutual exclusion to the field. + let future = match unsafe { &mut *ptr } { + Stage::Running(future) => future, + _ => unreachable!("unexpected stage"), + }; + + // Safety: The caller ensures the future is pinned. + let future = unsafe { Pin::new_unchecked(future) }; + + // The waker passed into the `poll` function does not require a ref + // count increment. + let waker_ref = waker_ref::(header); + let mut cx = Context::from_waker(&*waker_ref); + + future.poll(&mut cx) + }) + }; + + if res.is_ready() { + self.drop_future_or_output(); + } + + res + } + + /// Drop the future + /// + /// # Safety + /// + /// The caller must ensure it is safe to mutate the `stage` field. + pub(super) fn drop_future_or_output(&self) { + self.stage.with_mut(|ptr| { + // Safety: The caller ensures mutal exclusion to the field. + unsafe { *ptr = Stage::Consumed }; + }); + } + + /// Store the task output + /// + /// # Safety + /// + /// The caller must ensure it is safe to mutate the `stage` field. + pub(super) fn store_output(&self, output: super::Result) { + self.stage.with_mut(|ptr| { + // Safety: the caller ensures mutual exclusion to the field. + unsafe { *ptr = Stage::Finished(output) }; + }); + } + + /// Take the task output + /// + /// # Safety + /// + /// The caller must ensure it is safe to mutate the `stage` field. + pub(super) fn take_output(&self) -> super::Result { + use std::mem; + + self.stage.with_mut(|ptr| { + // Safety:: the caller ensures mutal exclusion to the field. + match mem::replace(unsafe { &mut *ptr }, Stage::Consumed) { + Stage::Finished(output) => output, + _ => panic!("unexpected task state"), + } + }) + } + + /// Schedule the future for execution + pub(super) fn schedule(&self, task: Notified) { + self.scheduler.with(|ptr| { + // Safety: Can only be called after initial `poll`, which is the + // only time the field is mutated. + match unsafe { &*ptr } { + Some(scheduler) => scheduler.schedule(task), + None => panic!("no scheduler set"), + } + }); + } + + /// Schedule the future for execution in the near future, yielding the + /// thread to other tasks. + pub(super) fn yield_now(&self, task: Notified) { + self.scheduler.with(|ptr| { + // Safety: Can only be called after initial `poll`, which is the + // only time the field is mutated. + match unsafe { &*ptr } { + Some(scheduler) => scheduler.yield_now(task), + None => panic!("no scheduler set"), + } + }); + } + + /// Release the task + /// + /// If the `Scheduler` implementation is able to, it returns the `Task` + /// handle immediately. The caller of this function will batch a ref-dec + /// with a state change. + pub(super) fn release(&self, task: Task) -> Option> { + use std::mem::ManuallyDrop; + + let task = ManuallyDrop::new(task); + + self.scheduler.with(|ptr| { + // Safety: Can only be called after initial `poll`, which is the + // only time the field is mutated. + match unsafe { &*ptr } { + Some(scheduler) => scheduler.release(&*task), + // Task was never polled + None => None, + } + }) + } +} + +cfg_rt_threaded! { + impl Header { + pub(crate) fn shutdown(&self) { + use crate::runtime::task::RawTask; + + let task = unsafe { RawTask::from_raw(self.into()) }; + task.shutdown(); + } + } +} + +#[test] +#[cfg(not(loom))] +fn header_lte_cache_line() { + use std::mem::size_of; + + assert!(size_of::

() <= 8 * size_of::<*const ()>()); +} diff --git a/tokio/src/task/error.rs b/tokio/src/runtime/task/error.rs similarity index 100% rename from tokio/src/task/error.rs rename to tokio/src/runtime/task/error.rs diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs new file mode 100644 index 000000000..f9cf5e75c --- /dev/null +++ b/tokio/src/runtime/task/harness.rs @@ -0,0 +1,369 @@ +use crate::runtime::task::core::{Cell, Core, Header, Trailer}; +use crate::runtime::task::state::Snapshot; +use crate::runtime::task::{JoinError, Notified, Schedule, Task}; + +use std::future::Future; +use std::mem; +use std::panic; +use std::ptr::NonNull; +use std::task::{Poll, Waker}; + +/// Typed raw task handle +pub(super) struct Harness { + cell: NonNull>, +} + +impl Harness +where + T: Future, + S: 'static, +{ + pub(super) unsafe fn from_raw(ptr: NonNull
) -> Harness { + Harness { + cell: ptr.cast::>(), + } + } + + fn header(&self) -> &Header { + unsafe { &self.cell.as_ref().header } + } + + fn trailer(&self) -> &Trailer { + unsafe { &self.cell.as_ref().trailer } + } + + fn core(&self) -> &Core { + unsafe { &self.cell.as_ref().core } + } +} + +impl Harness +where + T: Future, + S: Schedule, +{ + /// Polls the inner future. + /// + /// All necessary state checks and transitions are performed. + /// + /// Panics raised while polling the future are handled. + pub(super) fn poll(self) { + // If this is the first time the task is polled, the task will be bound + // to the scheduler, in which case the task ref count must be + // incremented. + let ref_inc = !self.core().is_bound(); + + // Transition the task to the running state. + // + // A failure to transition here indicates the task has been cancelled + // while in the run queue pending execution. + let snapshot = match self.header().state.transition_to_running(ref_inc) { + Ok(snapshot) => snapshot, + Err(_) => { + // The task was shutdown while in the run queue. At this point, + // we just hold a ref counted reference. Drop it here. + self.drop_reference(); + return; + } + }; + + // Ensure the task is bound to a scheduler instance. If this is the + // first time polling the task, a scheduler instance is pulled from the + // local context and assigned to the task. + // + // The scheduler maintains ownership of the task and responds to `wake` + // calls. + // + // The task reference count has been incremented. + self.core().bind_scheduler(self.to_task()); + + // The transition to `Running` done above ensures that a lock on the + // future has been obtained. This also ensures the `*mut T` pointer + // contains the future (as opposed to the output) and is initialized. + + let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { + struct Guard<'a, T: Future, S: Schedule> { + core: &'a Core, + polled: bool, + } + + impl Drop for Guard<'_, T, S> { + fn drop(&mut self) { + if !self.polled { + self.core.drop_future_or_output(); + } + } + } + + let mut guard = Guard { + core: self.core(), + polled: false, + }; + + // If the task is cancelled, avoid polling it, instead signalling it + // is complete. + if snapshot.is_cancelled() { + Poll::Ready(Err(JoinError::cancelled2())) + } else { + let res = guard.core.poll(self.header()); + + // prevent the guard from dropping the future + guard.polled = true; + + res.map(Ok) + } + })); + + match res { + Ok(Poll::Ready(out)) => { + self.complete(out, snapshot.is_join_interested()); + } + Ok(Poll::Pending) => { + match self.header().state.transition_to_idle() { + Ok(snapshot) => { + if snapshot.is_notified() { + // Signal yield + self.core().yield_now(Notified(self.to_task())); + } + } + Err(_) => self.cancel_task(), + } + } + Err(err) => { + self.complete(Err(JoinError::panic2(err)), snapshot.is_join_interested()); + } + } + } + + pub(super) fn dealloc(self) { + // Release the join waker, if there is one. + self.trailer().waker.with_mut(|_| ()); + + // Check causality + self.core().stage.with_mut(|_| {}); + self.core().scheduler.with_mut(|_| {}); + + unsafe { + drop(Box::from_raw(self.cell.as_ptr())); + } + } + + // ===== join handle ===== + + /// Read the task output into `dst`. + pub(super) fn try_read_output(self, dst: &mut Poll>, waker: &Waker) { + // Load a snapshot of the current task state + let snapshot = self.header().state.load(); + + debug_assert!(snapshot.is_join_interested()); + + if !snapshot.is_complete() { + // The waker must be stored in the task struct. + let res = if snapshot.has_join_waker() { + // There already is a waker stored in the struct. If it matches + // the provided waker, then there is no further work to do. + // Otherwise, the waker must be swapped. + let will_wake = unsafe { + // Safety: when `JOIN_INTEREST` is set, only `JOIN_HANDLE` + // may mutate the `waker` field. + self.trailer() + .waker + .with(|ptr| (*ptr).as_ref().unwrap().will_wake(waker)) + }; + + if will_wake { + // The task is not complete **and** the waker is up to date, + // there is nothing further that needs to be done. + return; + } + + // Unset the `JOIN_WAKER` to gain mutable access to the `waker` + // field then update the field with the new join worker. + // + // This requires two atomic operations, unsetting the bit and + // then resetting it. If the task transitions to complete + // concurrently to either one of those operations, then setting + // the join waker fails and we proceed to reading the task + // output. + self.header() + .state + .unset_waker() + .and_then(|snapshot| self.set_join_waker(waker.clone(), snapshot)) + } else { + self.set_join_waker(waker.clone(), snapshot) + }; + + match res { + Ok(_) => return, + Err(snapshot) => { + assert!(snapshot.is_complete()); + } + } + } + + *dst = Poll::Ready(self.core().take_output()); + } + + fn set_join_waker(&self, waker: Waker, snapshot: Snapshot) -> Result { + assert!(snapshot.is_join_interested()); + assert!(!snapshot.has_join_waker()); + + // Safety: Only the `JoinHandle` may set the `waker` field. When + // `JOIN_INTEREST` is **not** set, nothing else will touch the field. + unsafe { + self.trailer().waker.with_mut(|ptr| { + *ptr = Some(waker); + }); + } + + // Update the `JoinWaker` state accordingly + let res = self.header().state.set_join_waker(); + + // If the state could not be updated, then clear the join waker + if res.is_err() { + unsafe { + self.trailer().waker.with_mut(|ptr| { + *ptr = None; + }); + } + } + + res + } + + pub(super) fn drop_join_handle_slow(self) { + // Try to unset `JOIN_INTEREST`. This must be done as a first step in + // case the task concurrently completed. + if self.header().state.unset_join_interested().is_err() { + // It is our responsibility to drop the output. This is critical as + // the task output may not be `Send` and as such must remain with + // the scheduler or `JoinHandle`. i.e. if the output remains in the + // task structure until the task is deallocated, it may be dropped + // by a Waker on any arbitrary thread. + self.core().drop_future_or_output(); + } + + // Drop the `JoinHandle` reference, possibly deallocating the task + self.drop_reference(); + } + + // ===== waker behavior ===== + + pub(super) fn wake_by_val(self) { + self.wake_by_ref(); + self.drop_reference(); + } + + pub(super) fn wake_by_ref(&self) { + if self.header().state.transition_to_notified() { + self.core().schedule(Notified(self.to_task())); + } + } + + pub(super) fn drop_reference(self) { + if self.header().state.ref_dec() { + self.dealloc(); + } + } + + /// Forcibly shutdown the task + /// + /// Attempt to transition to `Running` in order to forcibly shutdown the + /// task. If the task is currently running or in a state of completion, then + /// there is nothing further to do. When the task completes running, it will + /// notice the `CANCELLED` bit and finalize the task. + pub(super) fn shutdown(self) { + if !self.header().state.transition_to_shutdown() { + // The task is concurrently running. No further work needed. + return; + } + + // By transitioning the lifcycle to `Running`, we have permission to + // drop the future. + self.cancel_task(); + } + + // ====== internal ====== + + fn cancel_task(self) { + // Drop the future from a panic guard. + let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { + self.core().drop_future_or_output(); + })); + + if let Err(err) = res { + // Dropping the future panicked, complete the join + // handle with the panic to avoid dropping the panic + // on the ground. + self.complete(Err(JoinError::panic2(err)), true); + } else { + self.complete(Err(JoinError::cancelled2()), true); + } + } + + fn complete(mut self, output: super::Result, is_join_interested: bool) { + if is_join_interested { + // Store the output. The future has already been dropped + // + // Safety: Mutual exclusion is obtained by having transitioned the task + // state -> Running + self.core().store_output(output); + + // Transition to `Complete`, notifying the `JoinHandle` if necessary. + self.transition_to_complete(); + } + + // The task has completed execution and will no longer be scheduled. + // + // Attempts to batch a ref-dec with the state transition below. + let ref_dec = if self.core().is_bound() { + if let Some(task) = self.core().release(self.to_task()) { + mem::forget(task); + true + } else { + false + } + } else { + false + }; + + // This might deallocate + let snapshot = self + .header() + .state + .transition_to_terminal(!is_join_interested, ref_dec); + + if snapshot.ref_count() == 0 { + self.dealloc() + } + } + + /// Transitions the task's lifecycle to `Complete`. Notifies the + /// `JoinHandle` if it still has interest in the completion. + fn transition_to_complete(&mut self) { + // Transition the task's lifecycle to `Complete` and get a snapshot of + // the task's sate. + let snapshot = self.header().state.transition_to_complete(); + + if !snapshot.is_join_interested() { + // The `JoinHandle` is not interested in the output of this task. It + // is our responsibility to drop the output. + self.core().drop_future_or_output(); + } else if snapshot.has_join_waker() { + // Notify the join handle. The previous transition obtains the + // lock on the waker cell. + self.wake_join(); + } + } + + fn wake_join(&self) { + self.trailer().waker.with(|ptr| match unsafe { &*ptr } { + Some(waker) => waker.wake_by_ref(), + None => panic!("waker missing"), + }); + } + + fn to_task(&self) -> Task { + unsafe { Task::from_raw(self.header().into()) } + } +} diff --git a/tokio/src/task/join.rs b/tokio/src/runtime/task/join.rs similarity index 75% rename from tokio/src/task/join.rs rename to tokio/src/runtime/task/join.rs index 8a8f25714..ed893a35c 100644 --- a/tokio/src/task/join.rs +++ b/tokio/src/runtime/task/join.rs @@ -1,5 +1,4 @@ -use crate::loom::alloc::Track; -use crate::task::RawTask; +use crate::runtime::task::RawTask; use std::fmt; use std::future::Future; @@ -99,46 +98,39 @@ impl Unpin for JoinHandle {} impl Future for JoinHandle { type Output = super::Result; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - use std::mem::MaybeUninit; + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut ret = Poll::Pending; - // Raw should always be set - let raw = self.raw.as_ref().unwrap(); - - // Load the current task state - let mut state = raw.header().state.load(); - - debug_assert!(state.is_join_interested()); - - if state.is_active() { - state = if state.has_join_waker() { - raw.swap_join_waker(cx.waker(), state) - } else { - raw.store_join_waker(cx.waker()) - }; - - if state.is_active() { - return Poll::Pending; - } - } - - let mut out = MaybeUninit::>::uninit(); + // Raw should always be set. If it is not, this is due to polling after + // completion + let raw = self + .raw + .as_ref() + .expect("polling after `JoinHandle` already completed"); + // Try to read the task output. If the task is not yet complete, the + // waker is stored and is notified once the task does complete. + // + // The function must go via the vtable, which requires erasing generic + // types. To do this, the function "return" is placed on the stack + // **before** calling the function and is passed into the function using + // `*mut ()`. + // + // Safety: + // + // The type of `T` must match the task's output type. unsafe { - // This could result in the task being freed. - raw.read_output(out.as_mut_ptr() as *mut (), state); - - self.raw = None; - - Poll::Ready(out.assume_init().into_inner()) + raw.try_read_output(&mut ret as *mut _ as *mut (), cx.waker()); } + + ret } } impl Drop for JoinHandle { fn drop(&mut self) { if let Some(raw) = self.raw.take() { - if raw.header().state.drop_join_handle_fast() { + if raw.header().state.drop_join_handle_fast().is_ok() { return; } diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs new file mode 100644 index 000000000..1ea60a9b6 --- /dev/null +++ b/tokio/src/runtime/task/mod.rs @@ -0,0 +1,219 @@ +mod core; +use self::core::Cell; +pub(crate) use self::core::Header; + +mod error; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::error::JoinError; + +mod harness; +use self::harness::Harness; + +mod join; +#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 +pub use self::join::JoinHandle; + +mod raw; +use self::raw::RawTask; + +mod state; +use self::state::State; + +mod waker; + +cfg_rt_threaded! { + mod stack; + pub(crate) use self::stack::TransferStack; +} + +use crate::util::linked_list; + +use std::future::Future; +use std::marker::PhantomData; +use std::ptr::NonNull; +use std::{fmt, mem}; + +/// An owned handle to the task, tracked by ref count +#[repr(transparent)] +pub(crate) struct Task { + raw: RawTask, + _p: PhantomData, +} + +unsafe impl Send for Task {} +unsafe impl Sync for Task {} + +/// A task was notified +#[repr(transparent)] +pub(crate) struct Notified(Task); + +unsafe impl Send for Notified {} +unsafe impl Sync for Notified {} + +/// Task result sent back +pub(crate) type Result = std::result::Result; + +pub(crate) trait Schedule: Sync + Sized + 'static { + /// Bind a task to the executor. + /// + /// Guaranteed to be called from the thread that called `poll` on the task. + /// The returned `Schedule` instance is associated with the task and is used + /// as `&self` in the other methods on this trait. + fn bind(task: Task) -> Self; + + /// The task has completed work and is ready to be released. The scheduler + /// is free to drop it whenever. + /// + /// If the scheduler can immediately release the task, it should return + /// it as part of the function. This enables the task module to batch + /// the ref-dec with other options. + fn release(&self, task: &Task) -> Option>; + + /// Schedule the task + fn schedule(&self, task: Notified); + + /// Schedule the task to run in the near future, yielding the thread to + /// other tasks. + fn yield_now(&self, task: Notified) { + self.schedule(task); + } +} + +/// Create a new task with an associated join handle +pub(crate) fn joinable(task: T) -> (Notified, JoinHandle) +where + T: Future + Send + 'static, + S: Schedule, +{ + let raw = RawTask::new::<_, S>(task); + + let task = Task { + raw, + _p: PhantomData, + }; + + let join = JoinHandle::new(raw); + + (Notified(task), join) +} + +cfg_rt_util! { + /// Create a new `!Send` task with an associated join handle + pub(crate) unsafe fn joinable_local(task: T) -> (Notified, JoinHandle) + where + T: Future + 'static, + S: Schedule, + { + let raw = RawTask::new::<_, S>(task); + + let task = Task { + raw, + _p: PhantomData, + }; + + let join = JoinHandle::new(raw); + + (Notified(task), join) + } +} + +impl Task { + pub(crate) unsafe fn from_raw(ptr: NonNull
) -> Task { + Task { + raw: RawTask::from_raw(ptr), + _p: PhantomData, + } + } + + pub(crate) fn header(&self) -> &Header { + self.raw.header() + } +} + +cfg_rt_threaded! { + impl Notified { + pub(crate) unsafe fn from_raw(ptr: NonNull
) -> Notified { + Notified(Task::from_raw(ptr)) + } + + pub(crate) fn header(&self) -> &Header { + self.0.header() + } + } + + impl Task { + pub(crate) fn into_raw(self) -> NonNull
{ + let ret = self.header().into(); + mem::forget(self); + ret + } + } + + impl Notified { + pub(crate) fn into_raw(self) -> NonNull
{ + self.0.into_raw() + } + } +} + +impl Task { + /// Pre-emptively cancel the task as part of the shutdown process. + pub(crate) fn shutdown(&self) { + self.raw.shutdown(); + } +} + +impl Notified { + /// Run the task + pub(crate) fn run(self) { + self.0.raw.poll(); + mem::forget(self); + } + + /// Pre-emptively cancel the task as part of the shutdown process. + pub(crate) fn shutdown(self) { + self.0.shutdown(); + } +} + +impl Drop for Task { + fn drop(&mut self) { + // Decrement the ref count + if self.header().state.ref_dec() { + // Deallocate if this is the final ref count + self.raw.dealloc(); + } + } +} + +impl fmt::Debug for Task { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(fmt, "Task({:p})", self.header()) + } +} + +impl fmt::Debug for Notified { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(fmt, "task::Notified({:p})", self.0.header()) + } +} + +/// # Safety +/// +/// Tasks are pinned +unsafe impl linked_list::Link for Task { + type Handle = Task; + type Target = Header; + + fn as_raw(handle: &Task) -> NonNull
{ + handle.header().into() + } + + unsafe fn from_raw(ptr: NonNull
) -> Task { + Task::from_raw(ptr) + } + + unsafe fn pointers(target: NonNull
) -> NonNull> { + NonNull::from(&mut *target.as_ref().owned.get()) + } +} diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs new file mode 100644 index 000000000..cae56d037 --- /dev/null +++ b/tokio/src/runtime/task/raw.rs @@ -0,0 +1,131 @@ +use crate::runtime::task::{Cell, Harness, Header, Schedule, State}; + +use std::future::Future; +use std::ptr::NonNull; +use std::task::{Poll, Waker}; + +/// Raw task handle +pub(super) struct RawTask { + ptr: NonNull
, +} + +pub(super) struct Vtable { + /// Poll the future + pub(super) poll: unsafe fn(NonNull
), + + /// Deallocate the memory + pub(super) dealloc: unsafe fn(NonNull
), + + /// Read the task output, if complete + pub(super) try_read_output: unsafe fn(NonNull
, *mut (), &Waker), + + /// The join handle has been dropped + pub(super) drop_join_handle_slow: unsafe fn(NonNull
), + + /// Scheduler is being shutdown + pub(super) shutdown: unsafe fn(NonNull
), +} + +/// Get the vtable for the requested `T` and `S` generics. +pub(super) fn vtable() -> &'static Vtable { + &Vtable { + poll: poll::, + dealloc: dealloc::, + try_read_output: try_read_output::, + drop_join_handle_slow: drop_join_handle_slow::, + shutdown: shutdown::, + } +} + +impl RawTask { + pub(super) fn new(task: T) -> RawTask + where + T: Future, + S: Schedule, + { + let ptr = Box::into_raw(Cell::<_, S>::new(task, State::new())); + let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) }; + + RawTask { ptr } + } + + pub(super) unsafe fn from_raw(ptr: NonNull
) -> RawTask { + RawTask { ptr } + } + + /// Returns a reference to the task's meta structure. + /// + /// Safe as `Header` is `Sync`. + pub(super) fn header(&self) -> &Header { + unsafe { self.ptr.as_ref() } + } + + /// Safety: mutual exclusion is required to call this function. + pub(super) fn poll(self) { + let vtable = self.header().vtable; + unsafe { (vtable.poll)(self.ptr) } + } + + pub(super) fn dealloc(self) { + let vtable = self.header().vtable; + unsafe { + (vtable.dealloc)(self.ptr); + } + } + + /// Safety: `dst` must be a `*mut Poll>` where `T` + /// is the future stored by the task. + pub(super) unsafe fn try_read_output(self, dst: *mut (), waker: &Waker) { + let vtable = self.header().vtable; + (vtable.try_read_output)(self.ptr, dst, waker); + } + + pub(super) fn drop_join_handle_slow(self) { + let vtable = self.header().vtable; + unsafe { (vtable.drop_join_handle_slow)(self.ptr) } + } + + pub(super) fn shutdown(self) { + let vtable = self.header().vtable; + unsafe { (vtable.shutdown)(self.ptr) } + } +} + +impl Clone for RawTask { + fn clone(&self) -> Self { + RawTask { ptr: self.ptr } + } +} + +impl Copy for RawTask {} + +unsafe fn poll(ptr: NonNull
) { + let harness = Harness::::from_raw(ptr); + harness.poll(); +} + +unsafe fn dealloc(ptr: NonNull
) { + let harness = Harness::::from_raw(ptr); + harness.dealloc(); +} + +unsafe fn try_read_output( + ptr: NonNull
, + dst: *mut (), + waker: &Waker, +) { + let out = &mut *(dst as *mut Poll>); + + let harness = Harness::::from_raw(ptr); + harness.try_read_output(out, waker); +} + +unsafe fn drop_join_handle_slow(ptr: NonNull
) { + let harness = Harness::::from_raw(ptr); + harness.drop_join_handle_slow() +} + +unsafe fn shutdown(ptr: NonNull
) { + let harness = Harness::::from_raw(ptr); + harness.shutdown() +} diff --git a/tokio/src/runtime/task/stack.rs b/tokio/src/runtime/task/stack.rs new file mode 100644 index 000000000..b2d50bafd --- /dev/null +++ b/tokio/src/runtime/task/stack.rs @@ -0,0 +1,81 @@ +use crate::loom::sync::atomic::AtomicPtr; +use crate::runtime::task::{Header, Task}; + +use std::marker::PhantomData; +use std::ptr::{self, NonNull}; +use std::sync::atomic::Ordering::{Acquire, Relaxed, Release}; + +/// Concurrent stack of tasks, used to pass ownership of a task from one worker +/// to another. +pub(crate) struct TransferStack { + head: AtomicPtr
, + _p: PhantomData, +} + +impl TransferStack { + pub(crate) fn new() -> TransferStack { + TransferStack { + head: AtomicPtr::new(ptr::null_mut()), + _p: PhantomData, + } + } + + pub(crate) fn push(&self, task: Task) { + let task = task.into_raw(); + + // We don't care about any memory associated w/ setting the `head` + // field, just the current value. + // + // The compare-exchange creates a release sequence. + let mut curr = self.head.load(Relaxed); + + loop { + unsafe { + *task.as_ref().stack_next.get() = NonNull::new(curr); + } + + let res = self + .head + .compare_exchange(curr, task.as_ptr() as *mut _, Release, Relaxed); + + match res { + Ok(_) => return, + Err(actual) => { + curr = actual; + } + } + } + } + + pub(crate) fn drain(&self) -> impl Iterator> { + struct Iter(Option>, PhantomData); + + impl Iterator for Iter { + type Item = Task; + + fn next(&mut self) -> Option> { + let task = self.0?; + + // Move the cursor forward + self.0 = unsafe { *task.as_ref().stack_next.get() }; + + // Return the task + unsafe { Some(Task::from_raw(task)) } + } + } + + impl Drop for Iter { + fn drop(&mut self) { + use std::process; + + if self.0.is_some() { + // we have bugs + process::abort(); + } + } + } + + let ptr = self.head.swap(ptr::null_mut(), Acquire); + Iter(NonNull::new(ptr), PhantomData) + } +} diff --git a/tokio/src/runtime/task/state.rs b/tokio/src/runtime/task/state.rs new file mode 100644 index 000000000..653c508da --- /dev/null +++ b/tokio/src/runtime/task/state.rs @@ -0,0 +1,447 @@ +use crate::loom::sync::atomic::AtomicUsize; + +use std::fmt; +use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; +use std::usize; + +pub(super) struct State { + val: AtomicUsize, +} + +/// Current state value +#[derive(Copy, Clone)] +pub(super) struct Snapshot(usize); + +type UpdateResult = Result; + +/// The task is currently being run. +const RUNNING: usize = 0b0001; + +/// The task is complete. +/// +/// Once this bit is set, it is never unset +const COMPLETE: usize = 0b0010; + +/// Extracts the task's lifecycle value from the state +const LIFECYCLE_MASK: usize = 0b11; + +/// Flag tracking if the task has been pushed into a run queue. +const NOTIFIED: usize = 0b100; + +/// The join handle is still around +const JOIN_INTEREST: usize = 0b1_000; + +/// A join handle waker has been set +const JOIN_WAKER: usize = 0b10_000; + +/// The task has been forcibly cancelled. +const CANCELLED: usize = 0b100_000; + +/// All bits +const STATE_MASK: usize = LIFECYCLE_MASK | NOTIFIED | JOIN_INTEREST | JOIN_WAKER | CANCELLED; + +/// Bits used by the ref count portion of the state. +const REF_COUNT_MASK: usize = !STATE_MASK; + +/// Number of positions to shift the ref count +const REF_COUNT_SHIFT: usize = REF_COUNT_MASK.count_zeros() as usize; + +/// One ref count +const REF_ONE: usize = 1 << REF_COUNT_SHIFT; + +/// State a task is initialized with +/// +/// A task is initialized with two references: one for the scheduler and one for +/// the `JoinHandle`. As the task starts with a `JoinHandle`, `JOIN_INTERST` is +/// set. A new task is immediately pushed into the run queue for execution and +/// starts with the `NOTIFIED` flag set. +const INITIAL_STATE: usize = (REF_ONE * 2) | JOIN_INTEREST | NOTIFIED; + +/// All transitions are performed via RMW operations. This establishes an +/// unambiguous modification order. +impl State { + /// Return a task's initial state + pub(super) fn new() -> State { + // A task is initialized with three references: one for the scheduler, + // one for the `JoinHandle`, one for the task handle made available in + // release. As the task starts with a `JoinHandle`, `JOIN_INTERST` is + // set. A new task is immediately pushed into the run queue for + // execution and starts with the `NOTIFIED` flag set. + State { + val: AtomicUsize::new(INITIAL_STATE), + } + } + + /// Loads the current state, establishes `Acquire` ordering. + pub(super) fn load(&self) -> Snapshot { + Snapshot(self.val.load(Acquire)) + } + + /// Attempt to transition the lifecycle to `Running`. + /// + /// If `ref_inc` is set, the reference count is also incremented. + /// + /// The `NOTIFIED` bit is always unset. + pub(super) fn transition_to_running(&self, ref_inc: bool) -> UpdateResult { + self.fetch_update(|curr| { + assert!(curr.is_notified()); + + let mut next = curr; + + if !next.is_idle() { + return None; + } + + if ref_inc { + next.ref_inc(); + } + + next.set_running(); + next.unset_notified(); + Some(next) + }) + } + + /// Transitions the task from `Running` -> `Idle`. + /// + /// Returns `Ok` if the transition to `Idle` is successful, `Err` otherwise. + /// In both cases, a snapshot of the state from **after** the transition is + /// returned. + /// + /// The transition to `Idle` fails if the task has been flagged to be + /// cancelled. + pub(super) fn transition_to_idle(&self) -> UpdateResult { + self.fetch_update(|curr| { + assert!(curr.is_running()); + + if curr.is_cancelled() { + return None; + } + + let mut next = curr; + next.unset_running(); + Some(next) + }) + } + + /// Transitions the task from `Running` -> `Complete`. + pub(super) fn transition_to_complete(&self) -> Snapshot { + const DELTA: usize = RUNNING | COMPLETE; + + let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel)); + assert!(prev.is_running()); + assert!(!prev.is_complete()); + + Snapshot(prev.0 ^ DELTA) + } + + /// Transition from `Complete` -> `Terminal`, decrementing the reference + /// count by 1. + /// + /// When `ref_dec` is set, an additional ref count decrement is performed. + /// This is used to batch atomic ops when possible. + pub(super) fn transition_to_terminal(&self, complete: bool, ref_dec: bool) -> Snapshot { + self.fetch_update(|mut snapshot| { + if complete { + snapshot.set_complete(); + } else { + assert!(snapshot.is_complete()); + } + + // Decrement the primary handle + snapshot.ref_dec(); + + if ref_dec { + // Decrement a second time + snapshot.ref_dec(); + } + + Some(snapshot) + }) + .unwrap() + } + + /// Transitions the state to `NOTIFIED`. + /// + /// Returns `true` if the task needs to be submitted to the pool for + /// execution + pub(super) fn transition_to_notified(&self) -> bool { + let prev = Snapshot(self.val.fetch_or(NOTIFIED, AcqRel)); + prev.will_need_queueing() + } + + /// Set the `CANCELLED` bit and attempt to transition to `Running`. + /// + /// Returns `true` if the transition to `Running` succeeded. + pub(super) fn transition_to_shutdown(&self) -> bool { + let mut prev = Snapshot(0); + + let _ = self.fetch_update(|mut snapshot| { + prev = snapshot; + + if snapshot.is_idle() { + snapshot.set_running(); + + if snapshot.is_notified() { + // If the task is idle and notified, this indicates the task is + // in the run queue and is considered owned by the scheduler. + // The shutdown operation claims ownership of the task, which + // means we need to assign an additional ref-count to the task + // in the queue. + snapshot.ref_inc(); + } + } + + snapshot.set_cancelled(); + Some(snapshot) + }); + + prev.is_idle() + } + + /// Optimistically tries to swap the state assuming the join handle is + /// __immediately__ dropped on spawn + pub(super) fn drop_join_handle_fast(&self) -> Result<(), ()> { + use std::sync::atomic::Ordering::Relaxed; + + // Relaxed is acceptable as if this function is called and succeeds, + // then nothing has been done w/ the join handle. + // + // The moment the join handle is used (polled), the `JOIN_WAKER` flag is + // set, at which point the CAS will fail. + // + // Given this, there is no risk if this operation is reordered. + self.val + .compare_exchange_weak( + INITIAL_STATE, + (INITIAL_STATE - REF_ONE) & !JOIN_INTEREST, + Release, + Relaxed, + ) + .map(|_| ()) + .map_err(|_| ()) + } + + /// Try to unset the JOIN_INTEREST flag. + /// + /// Returns `Ok` if the operation happens before the task transitions to a + /// completed state, `Err` otherwise. + pub(super) fn unset_join_interested(&self) -> UpdateResult { + self.fetch_update(|curr| { + assert!(curr.is_join_interested()); + + if curr.is_complete() { + return None; + } + + let mut next = curr; + next.unset_join_interested(); + + Some(next) + }) + } + + /// Set the `JOIN_WAKER` bit. + /// + /// Returns `Ok` if the bit is set, `Err` otherwise. This operation fails if + /// the task has completed. + pub(super) fn set_join_waker(&self) -> UpdateResult { + self.fetch_update(|curr| { + assert!(curr.is_join_interested()); + assert!(!curr.has_join_waker()); + + if curr.is_complete() { + return None; + } + + let mut next = curr; + next.set_join_waker(); + + Some(next) + }) + } + + /// Unsets the `JOIN_WAKER` bit. + /// + /// Returns `Ok` has been unset, `Err` otherwise. This operation fails if + /// the task has completed. + pub(super) fn unset_waker(&self) -> UpdateResult { + self.fetch_update(|curr| { + assert!(curr.is_join_interested()); + assert!(curr.has_join_waker()); + + if curr.is_complete() { + return None; + } + + let mut next = curr; + next.unset_join_waker(); + + Some(next) + }) + } + + pub(super) fn ref_inc(&self) { + use std::process; + use std::sync::atomic::Ordering::Relaxed; + + // Using a relaxed ordering is alright here, as knowledge of the + // original reference prevents other threads from erroneously deleting + // the object. + // + // As explained in the [Boost documentation][1], Increasing the + // reference counter can always be done with memory_order_relaxed: New + // references to an object can only be formed from an existing + // reference, and passing an existing reference from one thread to + // another must already provide any required synchronization. + // + // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) + let prev = self.val.fetch_add(REF_ONE, Relaxed); + + // If the reference count overflowed, abort. + if prev > isize::max_value() as usize { + process::abort(); + } + } + + /// Returns `true` if the task should be released. + pub(super) fn ref_dec(&self) -> bool { + use crate::loom::sync::atomic; + + let prev = Snapshot(self.val.fetch_sub(REF_ONE, Release)); + let is_final_ref = prev.ref_count() == 1; + + if is_final_ref { + atomic::fence(Acquire); + } + + is_final_ref + } + + fn fetch_update(&self, mut f: F) -> Result + where + F: FnMut(Snapshot) -> Option, + { + let mut curr = self.load(); + + loop { + let next = match f(curr) { + Some(next) => next, + None => return Err(curr), + }; + + let res = self.val.compare_exchange(curr.0, next.0, AcqRel, Acquire); + + match res { + Ok(_) => return Ok(next), + Err(actual) => curr = Snapshot(actual), + } + } + } +} + +// ===== impl Snapshot ===== + +impl Snapshot { + /// Returns `true` if the task is in an idle state. + pub(super) fn is_idle(self) -> bool { + self.0 & (RUNNING | COMPLETE) == 0 + } + + /// Returns `true` if the task has been flagged as notified. + pub(super) fn is_notified(self) -> bool { + self.0 & NOTIFIED == NOTIFIED + } + + fn unset_notified(&mut self) { + self.0 &= !NOTIFIED + } + + pub(super) fn is_running(self) -> bool { + self.0 & RUNNING == RUNNING + } + + fn set_running(&mut self) { + self.0 |= RUNNING; + } + + fn unset_running(&mut self) { + self.0 &= !RUNNING; + } + + pub(super) fn is_cancelled(self) -> bool { + self.0 & CANCELLED == CANCELLED + } + + fn set_cancelled(&mut self) { + self.0 |= CANCELLED; + } + + fn set_complete(&mut self) { + self.0 |= COMPLETE; + } + + /// Returns `true` if the task's future has completed execution. + pub(super) fn is_complete(self) -> bool { + self.0 & COMPLETE == COMPLETE + } + + pub(super) fn is_join_interested(self) -> bool { + self.0 & JOIN_INTEREST == JOIN_INTEREST + } + + fn unset_join_interested(&mut self) { + self.0 &= !JOIN_INTEREST + } + + pub(super) fn has_join_waker(self) -> bool { + self.0 & JOIN_WAKER == JOIN_WAKER + } + + fn set_join_waker(&mut self) { + self.0 |= JOIN_WAKER; + } + + fn unset_join_waker(&mut self) { + self.0 &= !JOIN_WAKER + } + + pub(super) fn ref_count(self) -> usize { + (self.0 & REF_COUNT_MASK) >> REF_COUNT_SHIFT + } + + fn ref_inc(&mut self) { + assert!(self.0 <= isize::max_value() as usize); + self.0 += REF_ONE; + } + + pub(super) fn ref_dec(&mut self) { + assert!(self.ref_count() > 0); + self.0 -= REF_ONE + } + + fn will_need_queueing(self) -> bool { + !self.is_notified() && self.is_idle() + } +} + +impl fmt::Debug for State { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + let snapshot = self.load(); + snapshot.fmt(fmt) + } +} + +impl fmt::Debug for Snapshot { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Snapshot") + .field("is_running", &self.is_running()) + .field("is_complete", &self.is_complete()) + .field("is_notified", &self.is_notified()) + .field("is_cancelled", &self.is_cancelled()) + .field("is_join_interested", &self.is_join_interested()) + .field("has_join_waker", &self.has_join_waker()) + .field("ref_count", &self.ref_count()) + .finish() + } +} diff --git a/tokio/src/task/waker.rs b/tokio/src/runtime/task/waker.rs similarity index 71% rename from tokio/src/task/waker.rs rename to tokio/src/runtime/task/waker.rs index 9892f1be8..5c2d478fb 100644 --- a/tokio/src/task/waker.rs +++ b/tokio/src/runtime/task/waker.rs @@ -1,10 +1,11 @@ -use crate::task::harness::Harness; -use crate::task::{Header, Schedule}; +use crate::runtime::task::harness::Harness; +use crate::runtime::task::{Header, Schedule}; use std::future::Future; use std::marker::PhantomData; use std::mem::ManuallyDrop; use std::ops; +use std::ptr::NonNull; use std::task::{RawWaker, RawWakerVTable, Waker}; pub(super) struct WakerRef<'a, S: 'static> { @@ -14,7 +15,7 @@ pub(super) struct WakerRef<'a, S: 'static> { /// Returns a `WakerRef` which avoids having to pre-emptively increase the /// refcount if there is no need to do so. -pub(super) fn waker_ref(meta: &Header) -> WakerRef<'_, S> +pub(super) fn waker_ref(header: &Header) -> WakerRef<'_, S> where T: Future, S: Schedule, @@ -27,7 +28,7 @@ where // point and not an *owned* waker, we must ensure that `drop` is never // called on this waker instance. This is done by wrapping it with // `ManuallyDrop` and then never calling drop. - let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker::(meta))) }; + let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker::(header))) }; WakerRef { waker, @@ -48,9 +49,9 @@ where T: Future, S: Schedule, { - let meta = ptr as *const Header; - (*meta).state.ref_inc(); - raw_waker::(meta) + let header = ptr as *const Header; + (*header).state.ref_inc(); + raw_waker::(header) } unsafe fn drop_waker(ptr: *const ()) @@ -58,8 +59,9 @@ where T: Future, S: Schedule, { - let harness = Harness::::from_raw(ptr as *mut _); - harness.drop_waker(); + let ptr = NonNull::new_unchecked(ptr as *mut Header); + let harness = Harness::::from_raw(ptr); + harness.drop_reference(); } unsafe fn wake_by_val(ptr: *const ()) @@ -67,7 +69,8 @@ where T: Future, S: Schedule, { - let harness = Harness::::from_raw(ptr as *mut _); + let ptr = NonNull::new_unchecked(ptr as *mut Header); + let harness = Harness::::from_raw(ptr); harness.wake_by_val(); } @@ -77,16 +80,17 @@ where T: Future, S: Schedule, { - let harness = Harness::::from_raw(ptr as *mut _); + let ptr = NonNull::new_unchecked(ptr as *mut Header); + let harness = Harness::::from_raw(ptr); harness.wake_by_ref(); } -fn raw_waker(meta: *const Header) -> RawWaker +fn raw_waker(header: *const Header) -> RawWaker where T: Future, S: Schedule, { - let ptr = meta as *const (); + let ptr = header as *const (); let vtable = &RawWakerVTable::new( clone_waker::, wake_by_val::, diff --git a/tokio/src/runtime/tests/loom_pool.rs b/tokio/src/runtime/tests/loom_pool.rs new file mode 100644 index 000000000..275e0ff54 --- /dev/null +++ b/tokio/src/runtime/tests/loom_pool.rs @@ -0,0 +1,381 @@ +/// Full runtime loom tests. These are heavy tests and take significant time to +/// run on CI. +/// +/// Use `LOOM_MAX_PREEMPTIONS=1` to do a "quick" run as a smoke test. +/// +/// In order to speed up the C +use crate::future::poll_fn; +use crate::runtime::tests::loom_oneshot as oneshot; +use crate::runtime::{self, Runtime}; +use crate::{spawn, task}; +use tokio_test::assert_ok; + +use loom::sync::atomic::{AtomicBool, AtomicUsize}; +use loom::sync::{Arc, Mutex}; + +use pin_project_lite::pin_project; +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::Ordering::{Relaxed, SeqCst}; +use std::task::{Context, Poll}; + +/// Tests are divided into groups to make the runs faster on CI. +mod group_a { + use super::*; + + #[test] + fn racy_shutdown() { + loom::model(|| { + let pool = mk_pool(1); + + // here's the case we want to exercise: + // + // a worker that still has tasks in its local queue gets sent to the blocking pool (due to + // block_in_place). the blocking pool is shut down, so drops the worker. the worker's + // shutdown method never gets run. + // + // we do this by spawning two tasks on one worker, the first of which does block_in_place, + // and then immediately drop the pool. + + pool.spawn(track(async { + crate::task::block_in_place(|| {}); + })); + pool.spawn(track(async {})); + drop(pool); + }); + } + + #[test] + fn pool_multi_spawn() { + loom::model(|| { + let pool = mk_pool(2); + let c1 = Arc::new(AtomicUsize::new(0)); + + let (tx, rx) = oneshot::channel(); + let tx1 = Arc::new(Mutex::new(Some(tx))); + + // Spawn a task + let c2 = c1.clone(); + let tx2 = tx1.clone(); + pool.spawn(track(async move { + spawn(track(async move { + if 1 == c1.fetch_add(1, Relaxed) { + tx1.lock().unwrap().take().unwrap().send(()); + } + })); + })); + + // Spawn a second task + pool.spawn(track(async move { + spawn(track(async move { + if 1 == c2.fetch_add(1, Relaxed) { + tx2.lock().unwrap().take().unwrap().send(()); + } + })); + })); + + rx.recv(); + }); + } + + fn only_blocking_inner(first_pending: bool) { + loom::model(move || { + let pool = mk_pool(1); + let (block_tx, block_rx) = oneshot::channel(); + + pool.spawn(track(async move { + crate::task::block_in_place(move || { + block_tx.send(()); + }); + if first_pending { + task::yield_now().await + } + })); + + block_rx.recv(); + drop(pool); + }); + } + + #[test] + fn only_blocking_without_pending() { + only_blocking_inner(false) + } + + #[test] + fn only_blocking_with_pending() { + only_blocking_inner(true) + } +} + +mod group_b { + use super::*; + + fn blocking_and_regular_inner(first_pending: bool) { + const NUM: usize = 3; + loom::model(move || { + let pool = mk_pool(1); + let cnt = Arc::new(AtomicUsize::new(0)); + + let (block_tx, block_rx) = oneshot::channel(); + let (done_tx, done_rx) = oneshot::channel(); + let done_tx = Arc::new(Mutex::new(Some(done_tx))); + + pool.spawn(track(async move { + crate::task::block_in_place(move || { + block_tx.send(()); + }); + if first_pending { + task::yield_now().await + } + })); + + for _ in 0..NUM { + let cnt = cnt.clone(); + let done_tx = done_tx.clone(); + + pool.spawn(track(async move { + if NUM == cnt.fetch_add(1, Relaxed) + 1 { + done_tx.lock().unwrap().take().unwrap().send(()); + } + })); + } + + done_rx.recv(); + block_rx.recv(); + + drop(pool); + }); + } + + #[test] + fn blocking_and_regular() { + blocking_and_regular_inner(false); + } + + #[test] + fn blocking_and_regular_with_pending() { + blocking_and_regular_inner(true); + } + + #[test] + fn pool_shutdown() { + loom::model(|| { + let pool = mk_pool(2); + + pool.spawn(track(async move { + gated2(true).await; + })); + + pool.spawn(track(async move { + gated2(false).await; + })); + + drop(pool); + }); + } + + #[test] + fn join_output() { + loom::model(|| { + let mut rt = mk_pool(1); + + rt.block_on(async { + let t = crate::spawn(track(async { "hello" })); + + let out = assert_ok!(t.await); + assert_eq!("hello", out.into_inner()); + }); + }); + } + + #[test] + fn poll_drop_handle_then_drop() { + loom::model(|| { + let mut rt = mk_pool(1); + + rt.block_on(async move { + let mut t = crate::spawn(track(async { "hello" })); + + poll_fn(|cx| { + let _ = Pin::new(&mut t).poll(cx); + Poll::Ready(()) + }) + .await; + }); + }) + } + + #[test] + fn complete_block_on_under_load() { + loom::model(|| { + let mut pool = mk_pool(1); + + pool.block_on(async { + // Trigger a re-schedule + crate::spawn(track(async { + for _ in 0..2 { + task::yield_now().await; + } + })); + + gated2(true).await + }); + }); + } +} + +mod group_c { + use super::*; + + #[test] + fn shutdown_with_notification() { + use crate::stream::StreamExt; + use crate::sync::{mpsc, oneshot}; + + loom::model(|| { + let rt = mk_pool(2); + let (done_tx, done_rx) = oneshot::channel::<()>(); + + rt.spawn(track(async move { + let (mut tx, mut rx) = mpsc::channel::<()>(10); + + crate::spawn(async move { + crate::task::spawn_blocking(move || { + let _ = tx.try_send(()); + }); + + let _ = done_rx.await; + }); + + while let Some(_) = rx.next().await {} + + let _ = done_tx.send(()); + })); + }); + } +} + +mod group_d { + use super::*; + + #[test] + fn pool_multi_notify() { + loom::model(|| { + let pool = mk_pool(2); + + let c1 = Arc::new(AtomicUsize::new(0)); + + let (done_tx, done_rx) = oneshot::channel(); + let done_tx1 = Arc::new(Mutex::new(Some(done_tx))); + + // Spawn a task + let c2 = c1.clone(); + let done_tx2 = done_tx1.clone(); + pool.spawn(track(async move { + gated().await; + gated().await; + + if 1 == c1.fetch_add(1, Relaxed) { + done_tx1.lock().unwrap().take().unwrap().send(()); + } + })); + + // Spawn a second task + pool.spawn(track(async move { + gated().await; + gated().await; + + if 1 == c2.fetch_add(1, Relaxed) { + done_tx2.lock().unwrap().take().unwrap().send(()); + } + })); + + done_rx.recv(); + }); + } +} + +fn mk_pool(num_threads: usize) -> Runtime { + runtime::Builder::new() + .threaded_scheduler() + .core_threads(num_threads) + .build() + .unwrap() +} + +fn gated() -> impl Future { + gated2(false) +} + +fn gated2(thread: bool) -> impl Future { + use loom::thread; + use std::sync::Arc; + + let gate = Arc::new(AtomicBool::new(false)); + let mut fired = false; + + poll_fn(move |cx| { + if !fired { + let gate = gate.clone(); + let waker = cx.waker().clone(); + + if thread { + thread::spawn(move || { + gate.store(true, SeqCst); + waker.wake_by_ref(); + }); + } else { + spawn(track(async move { + gate.store(true, SeqCst); + waker.wake_by_ref(); + })); + } + + fired = true; + + return Poll::Pending; + } + + if gate.load(SeqCst) { + Poll::Ready("hello world") + } else { + Poll::Pending + } + }) +} + +fn track(f: T) -> Track { + Track { + inner: f, + arc: Arc::new(()), + } +} + +pin_project! { + struct Track { + #[pin] + inner: T, + // Arc is used to hook into loom's leak tracking. + arc: Arc<()>, + } +} + +impl Track { + fn into_inner(self) -> T { + self.inner + } +} + +impl Future for Track { + type Output = Track; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let me = self.project(); + + Poll::Ready(Track { + inner: ready!(me.inner.poll(cx)), + arc: me.arc.clone(), + }) + } +} diff --git a/tokio/src/runtime/tests/mod.rs b/tokio/src/runtime/tests/mod.rs index da592f761..b932956ca 100644 --- a/tokio/src/runtime/tests/mod.rs +++ b/tokio/src/runtime/tests/mod.rs @@ -1,7 +1,8 @@ -//! Testing utilities +cfg_loom! { + mod loom_blocking; + mod loom_oneshot; + mod loom_pool; +} -#[cfg(loom)] -pub(crate) mod loom_oneshot; - -#[cfg(loom)] -pub(crate) mod loom_blocking; +#[cfg(miri)] +mod task; diff --git a/tokio/src/runtime/tests/task.rs b/tokio/src/runtime/tests/task.rs new file mode 100644 index 000000000..82315a04f --- /dev/null +++ b/tokio/src/runtime/tests/task.rs @@ -0,0 +1,159 @@ +use crate::runtime::task::{self, Schedule, Task}; +use crate::util::linked_list::LinkedList; +use crate::util::TryLock; + +use std::collections::VecDeque; +use std::sync::Arc; + +#[test] +fn create_drop() { + let _ = task::joinable::<_, Runtime>(async { unreachable!() }); +} + +#[test] +fn schedule() { + with(|rt| { + let (task, _) = task::joinable(async { + crate::task::yield_now().await; + }); + + rt.schedule(task); + + assert_eq!(2, rt.tick()); + }) +} + +#[test] +fn shutdown() { + with(|rt| { + let (task, _) = task::joinable(async { + loop { + crate::task::yield_now().await; + } + }); + + rt.schedule(task); + rt.tick_max(1); + + rt.shutdown(); + }) +} + +fn with(f: impl FnOnce(Runtime)) { + struct Reset; + + impl Drop for Reset { + fn drop(&mut self) { + let _rt = CURRENT.try_lock().unwrap().take(); + } + } + + let _reset = Reset; + + let rt = Runtime(Arc::new(Inner { + released: task::TransferStack::new(), + core: TryLock::new(Core { + queue: VecDeque::new(), + tasks: LinkedList::new(), + }), + })); + + *CURRENT.try_lock().unwrap() = Some(rt.clone()); + f(rt) +} + +#[derive(Clone)] +struct Runtime(Arc); + +struct Inner { + released: task::TransferStack, + core: TryLock, +} + +struct Core { + queue: VecDeque>, + tasks: LinkedList>, +} + +static CURRENT: TryLock> = TryLock::new(None); + +impl Runtime { + fn tick(&self) -> usize { + self.tick_max(usize::max_value()) + } + + fn tick_max(&self, max: usize) -> usize { + let mut n = 0; + + while !self.is_empty() && n < max { + let task = self.next_task(); + n += 1; + task.run(); + } + + self.0.maintenance(); + + n + } + + fn is_empty(&self) -> bool { + self.0.core.try_lock().unwrap().queue.is_empty() + } + + fn next_task(&self) -> task::Notified { + self.0.core.try_lock().unwrap().queue.pop_front().unwrap() + } + + fn shutdown(&self) { + let mut core = self.0.core.try_lock().unwrap(); + + for task in core.tasks.iter() { + task.shutdown(); + } + + while let Some(task) = core.queue.pop_back() { + task.shutdown(); + } + + drop(core); + + while !self.0.core.try_lock().unwrap().tasks.is_empty() { + self.0.maintenance(); + } + } +} + +impl Inner { + fn maintenance(&self) { + use std::mem::ManuallyDrop; + + for task in self.released.drain() { + let task = ManuallyDrop::new(task); + + // safety: see worker.rs + unsafe { + let ptr = task.header().into(); + self.core.try_lock().unwrap().tasks.remove(ptr); + } + } + } +} + +impl Schedule for Runtime { + fn bind(task: Task) -> Runtime { + let rt = CURRENT.try_lock().unwrap().as_ref().unwrap().clone(); + rt.0.core.try_lock().unwrap().tasks.push_front(task); + rt + } + + fn release(&self, task: &Task) -> Option> { + // safety: copying worker.rs + let task = unsafe { Task::from_raw(task.header().into()) }; + self.0.released.push(task); + None + } + + fn schedule(&self, task: task::Notified) { + self.0.core.try_lock().unwrap().queue.push_back(task); + } +} diff --git a/tokio/src/runtime/thread_pool/atomic_cell.rs b/tokio/src/runtime/thread_pool/atomic_cell.rs new file mode 100644 index 000000000..2bda0fc73 --- /dev/null +++ b/tokio/src/runtime/thread_pool/atomic_cell.rs @@ -0,0 +1,52 @@ +use crate::loom::sync::atomic::AtomicPtr; + +use std::ptr; +use std::sync::atomic::Ordering::AcqRel; + +pub(super) struct AtomicCell { + data: AtomicPtr, +} + +unsafe impl Send for AtomicCell {} +unsafe impl Sync for AtomicCell {} + +impl AtomicCell { + pub(super) fn new(data: Option>) -> AtomicCell { + AtomicCell { + data: AtomicPtr::new(to_raw(data)), + } + } + + pub(super) fn swap(&self, val: Option>) -> Option> { + let old = self.data.swap(to_raw(val), AcqRel); + from_raw(old) + } + + #[cfg(feature = "blocking")] + pub(super) fn set(&self, val: Box) { + let _ = self.swap(Some(val)); + } + + pub(super) fn take(&self) -> Option> { + self.swap(None) + } +} + +fn to_raw(data: Option>) -> *mut T { + data.map(Box::into_raw).unwrap_or(ptr::null_mut()) +} + +fn from_raw(val: *mut T) -> Option> { + if val.is_null() { + None + } else { + Some(unsafe { Box::from_raw(val) }) + } +} + +impl Drop for AtomicCell { + fn drop(&mut self) { + // Free any data still held by the cell + let _ = self.take(); + } +} diff --git a/tokio/src/runtime/thread_pool/current.rs b/tokio/src/runtime/thread_pool/current.rs deleted file mode 100644 index 60a207234..000000000 --- a/tokio/src/runtime/thread_pool/current.rs +++ /dev/null @@ -1,84 +0,0 @@ -use crate::loom::sync::Arc; -use crate::runtime::thread_pool::{slice, Owned}; - -use std::cell::Cell; -use std::ptr; - -/// Tracks the current worker -#[derive(Debug)] -pub(super) struct Current { - inner: Inner, -} - -#[derive(Debug, Copy, Clone)] -struct Inner { - // thread-local variables cannot track generics. However, the current worker - // is only checked when `P` is already known, so the type can be figured out - // on demand. - workers: *const (), - idx: usize, -} - -// Pointer to the current worker info -thread_local!(static CURRENT_WORKER: Cell = Cell::new(Inner::new())); - -pub(super) fn set(pool: &Arc, index: usize, f: F) -> R -where - F: FnOnce() -> R, -{ - CURRENT_WORKER.with(|cell| { - assert!(cell.get().workers.is_null()); - - struct Guard<'a>(&'a Cell); - - impl Drop for Guard<'_> { - fn drop(&mut self) { - self.0.set(Inner::new()); - } - } - - cell.set(Inner { - workers: pool.shared() as *const _ as *const (), - idx: index, - }); - - let _g = Guard(cell); - - f() - }) -} - -pub(super) fn clear() { - CURRENT_WORKER.with(|cell| cell.set(Inner::new())) -} - -pub(super) fn get(f: F) -> R -where - F: FnOnce(&Current) -> R, -{ - CURRENT_WORKER.with(|cell| { - let current = Current { inner: cell.get() }; - f(¤t) - }) -} - -impl Current { - pub(super) fn as_member<'a>(&self, set: &'a slice::Set) -> Option<&'a Owned> { - let inner = CURRENT_WORKER.with(|cell| cell.get()); - - if ptr::eq(inner.workers as *const _, set.shared().as_ptr()) { - Some(unsafe { &*set.owned()[inner.idx].get() }) - } else { - None - } - } -} - -impl Inner { - fn new() -> Inner { - Inner { - workers: ptr::null(), - idx: 0, - } - } -} diff --git a/tokio/src/runtime/thread_pool/mod.rs b/tokio/src/runtime/thread_pool/mod.rs index a52603560..8a74fe38b 100644 --- a/tokio/src/runtime/thread_pool/mod.rs +++ b/tokio/src/runtime/thread_pool/mod.rs @@ -1,45 +1,23 @@ //! Threadpool -mod current; +mod atomic_cell; +use atomic_cell::AtomicCell; mod idle; use self::idle::Idle; -mod owned; -use self::owned::Owned; - mod queue; -mod spawner; -pub(crate) use self::spawner::Spawner; - -mod slice; - -mod shared; -use self::shared::Shared; - mod worker; -use worker::Worker; +pub(crate) use worker::Launch; cfg_blocking! { pub(crate) use worker::block_in_place; } -/// Unit tests -#[cfg(test)] -mod tests; - -#[cfg(not(loom))] -const LOCAL_QUEUE_CAPACITY: usize = 256; - -// Shrink the size of the local queue when using loom. This shouldn't impact -// logic, but allows loom to test more edge cases in a reasonable a mount of -// time. -#[cfg(loom)] -const LOCAL_QUEUE_CAPACITY: usize = 2; - -use crate::runtime::{self, Parker}; -use crate::task::JoinHandle; +use crate::loom::sync::Arc; +use crate::runtime::task::{self, JoinHandle}; +use crate::runtime::Parker; use std::fmt; use std::future::Future; @@ -49,19 +27,32 @@ pub(crate) struct ThreadPool { spawner: Spawner, } -pub(crate) struct Workers { - workers: Vec, +/// Submit futures to the associated thread pool for execution. +/// +/// A `Spawner` instance is a handle to a single thread pool that allows the owner +/// of the handle to spawn futures onto the thread pool. +/// +/// The `Spawner` handle is *only* used for spawning new futures. It does not +/// impact the lifecycle of the thread pool in any way. The thread pool may +/// shutdown while there are outstanding `Spawner` instances. +/// +/// `Spawner` instances are obtained by calling [`ThreadPool::spawner`]. +/// +/// [`ThreadPool::spawner`]: struct.ThreadPool.html#method.spawner +#[derive(Clone)] +pub(crate) struct Spawner { + shared: Arc, } +// ===== impl ThreadPool ===== + impl ThreadPool { - pub(crate) fn new(pool_size: usize, parker: Parker) -> (ThreadPool, Workers) { - let (pool, workers) = worker::create_set(pool_size, parker); + pub(crate) fn new(size: usize, parker: Parker) -> (ThreadPool, Launch) { + let (shared, launch) = worker::create(size, parker); + let spawner = Spawner { shared }; + let thread_pool = ThreadPool { spawner }; - let spawner = Spawner::new(pool); - - let pool = ThreadPool { spawner }; - - (pool, Workers { workers }) + (thread_pool, launch) } /// Returns reference to `Spawner`. @@ -102,16 +93,27 @@ impl fmt::Debug for ThreadPool { impl Drop for ThreadPool { fn drop(&mut self) { - self.spawner.workers().close(); + self.spawner.shared.close(); } } -impl Workers { - pub(crate) fn spawn(self, rt: &runtime::Handle) { - rt.enter(|| { - for worker in self.workers { - runtime::spawn_blocking(move || worker.run()); - } - }); +// ==== impl Spawner ===== + +impl Spawner { + /// Spawns a future onto the thread pool + pub(crate) fn spawn(&self, future: F) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + let (task, handle) = task::joinable(future); + self.shared.schedule(task, false); + handle + } +} + +impl fmt::Debug for Spawner { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Spawner").finish() } } diff --git a/tokio/src/runtime/thread_pool/owned.rs b/tokio/src/runtime/thread_pool/owned.rs deleted file mode 100644 index b60eb7f3d..000000000 --- a/tokio/src/runtime/thread_pool/owned.rs +++ /dev/null @@ -1,83 +0,0 @@ -use crate::loom::sync::atomic::AtomicUsize; -use crate::runtime::thread_pool::{queue, Shared}; -use crate::task::{self, Task}; -use crate::util::FastRand; - -use std::cell::Cell; - -/// Per-worker data accessible only by the thread driving the worker. -#[derive(Debug)] -pub(super) struct Owned { - /// Worker generation. This guards concurrent access to the `Owned` struct. - /// When a worker starts running, it checks that the generation it has - /// assigned matches the current generation. When it does, the worker has - /// obtained unique access to the struct. When it fails, another thread has - /// gained unique access. - pub(super) generation: AtomicUsize, - - /// Worker tick number. Used to schedule bookkeeping tasks every so often. - pub(super) tick: Cell, - - /// Caches the pool run state. - pub(super) is_running: Cell, - - /// `true` if the worker is currently searching for more work. - pub(super) is_searching: Cell, - - /// `true` when worker notification should be delayed. - /// - /// This is used to batch notifications triggered by the parker. - pub(super) defer_notification: Cell, - - /// `true` if a task was submitted while `defer_notification` was set - pub(super) did_submit_task: Cell, - - /// Fast random number generator - pub(super) rand: FastRand, - - /// Work queue - pub(super) work_queue: queue::Worker, - - /// List of tasks owned by the worker - pub(super) owned_tasks: task::OwnedList, -} - -impl Owned { - pub(super) fn new(work_queue: queue::Worker, rand: FastRand) -> Owned { - Owned { - generation: AtomicUsize::new(0), - tick: Cell::new(1), - is_running: Cell::new(true), - is_searching: Cell::new(false), - defer_notification: Cell::new(false), - did_submit_task: Cell::new(false), - rand, - work_queue, - owned_tasks: task::OwnedList::new(), - } - } - - /// Returns `true` if a worker should be notified - pub(super) fn submit_local(&self, task: Task) -> bool { - let ret = self.work_queue.push(task); - - if self.defer_notification.get() { - self.did_submit_task.set(true); - false - } else { - ret - } - } - - pub(super) fn submit_local_yield(&self, task: Task) { - self.work_queue.push_yield(task); - } - - pub(super) fn bind_task(&mut self, task: &Task) { - self.owned_tasks.insert(task); - } - - pub(super) fn release_task(&mut self, task: &Task) { - self.owned_tasks.remove(task); - } -} diff --git a/tokio/src/runtime/thread_pool/queue.rs b/tokio/src/runtime/thread_pool/queue.rs new file mode 100644 index 000000000..66cec5040 --- /dev/null +++ b/tokio/src/runtime/thread_pool/queue.rs @@ -0,0 +1,568 @@ +//! Run-queue structures to support a work-stealing scheduler + +use crate::loom::cell::{CausalCell, CausalCheck}; +use crate::loom::sync::atomic::{self, AtomicU32, AtomicUsize}; +use crate::loom::sync::{Arc, Mutex}; +use crate::runtime::task; + +use std::marker::PhantomData; +use std::mem::MaybeUninit; +use std::ptr::{self, NonNull}; +use std::sync::atomic::Ordering::{Acquire, Release}; + +/// Producer handle. May only be used from a single thread. +pub(super) struct Local { + inner: Arc>, + + /// LIFO slot. Cannot be stolen. + next: Option>, +} + +/// Consumer handle. May be used from many threads. +pub(super) struct Steal(Arc>); + +/// Growable, MPMC queue used to inject new tasks into the scheduler and as an +/// overflow queue when the local, fixed-size, array queue overflows. +pub(super) struct Inject { + /// Pointers to the head and tail of the queue + pointers: Mutex, + + /// Number of pending tasks in the queue. This helps prevent unnecessary + /// locking in the hot path. + len: AtomicUsize, + + _p: PhantomData, +} + +pub(super) struct Inner { + /// Concurrently updated by many threads. + head: AtomicU32, + + /// Only updated by producer thread but read by many threads. + tail: AtomicU32, + + /// Elements + buffer: Box<[CausalCell>>]>, +} + +struct Pointers { + /// True if the queue is closed + is_closed: bool, + + /// Linked-list head + head: Option>, + + /// Linked-list tail + tail: Option>, +} + +unsafe impl Send for Inner {} +unsafe impl Sync for Inner {} +unsafe impl Send for Inject {} +unsafe impl Sync for Inject {} + +#[cfg(not(loom))] +const LOCAL_QUEUE_CAPACITY: usize = 256; + +// Shrink the size of the local queue when using loom. This shouldn't impact +// logic, but allows loom to test more edge cases in a reasonable a mount of +// time. +#[cfg(loom)] +const LOCAL_QUEUE_CAPACITY: usize = 2; + +const MASK: usize = LOCAL_QUEUE_CAPACITY - 1; + +/// Create a new local run-queue +pub(super) fn local() -> (Steal, Local) { + debug_assert!(LOCAL_QUEUE_CAPACITY >= 2 && LOCAL_QUEUE_CAPACITY.is_power_of_two()); + + let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY); + + for _ in 0..LOCAL_QUEUE_CAPACITY { + buffer.push(CausalCell::new(MaybeUninit::uninit())); + } + + let inner = Arc::new(Inner { + head: AtomicU32::new(0), + tail: AtomicU32::new(0), + buffer: buffer.into(), + }); + + let local = Local { + inner: inner.clone(), + next: None, + }; + + let remote = Steal(inner); + + (remote, local) +} + +impl Local { + /// Returns true if the queue has entries that can be stealed. + pub(super) fn is_stealable(&self) -> bool { + !self.inner.is_empty() + } + + /// Returns true if the queue has an unstealable entry. + pub(super) fn has_unstealable(&self) -> bool { + self.next.is_some() + } + + /// Push a task to the local queue. Returns `true` if a stealer should be + /// notified. + pub(super) fn push(&mut self, task: task::Notified, inject: &Inject) -> bool { + let prev = self.next.take(); + let ret = prev.is_some(); + + if let Some(prev) = prev { + self.push_back(prev, inject); + } + + self.next = Some(task); + + ret + } + + /// Pushes a task to the back of the local queue, skipping the LIFO slot. + pub(super) fn push_back(&mut self, mut task: task::Notified, inject: &Inject) { + loop { + let head = self.inner.head.load(Acquire); + + // safety: this is the **only** thread that updates this cell. + let tail = unsafe { self.inner.tail.unsync_load() }; + + if tail.wrapping_sub(head) < LOCAL_QUEUE_CAPACITY as u32 { + // Map the position to a slot index. + let idx = tail as usize & MASK; + + self.inner.buffer[idx].with_mut(|ptr| { + // Write the task to the slot + // + // Safety: There is only one producer and the above `if` + // condition ensures we don't touch a cell if there is a + // value, thus no consumer. + unsafe { + ptr::write((*ptr).as_mut_ptr(), task); + } + }); + + // Make the task available. Synchronizes with a load in + // `steal_into2`. + self.inner.tail.store(tail.wrapping_add(1), Release); + + return; + } + + // The local buffer is full. Push a batch of work to the inject + // queue. + match self.push_overflow(task, head, tail, inject) { + Ok(_) => return, + // Lost the race, try again + Err(v) => task = v, + } + + atomic::spin_loop_hint(); + } + } + + /// Moves a batch of tasks into the inject queue. + /// + /// This will temporarily make some of the tasks unavailable to stealers. + /// Once `push_overflow` is done, a notification is sent out, so if other + /// workers "missed" some of the tasks during a steal, they will get + /// another opportunity. + #[inline(never)] + fn push_overflow( + &mut self, + task: task::Notified, + head: u32, + tail: u32, + inject: &Inject, + ) -> Result<(), task::Notified> { + const BATCH_LEN: usize = LOCAL_QUEUE_CAPACITY / 2 + 1; + + let n = tail.wrapping_sub(head) / 2; + debug_assert_eq!(n as usize, LOCAL_QUEUE_CAPACITY / 2, "queue is not full"); + + // Claim a bunch of tasks + // + // We are claiming the tasks **before** reading them out of the buffer. + // This is safe because only the **current** thread is able to push new + // tasks. + // + // There isn't really any need for memory ordering... Relaxed would + // work. This is because all tasks are pushed into the queue from the + // current thread (or memory has been acquired if the local queue handle + // moved). + let actual = self.inner.head.compare_and_swap(head, head + n, Release); + if actual != head { + // We failed to claim the tasks, losing the race. Return out of + // this function and try the full `push` routine again. The queue + // may not be full anymore. + return Err(task); + } + + // link the tasks + for i in 0..n { + let j = i + 1; + + let i_idx = (i + head) as usize & MASK; + let j_idx = (j + head) as usize & MASK; + + // Get the next pointer + let next = if j == n { + // The last task in the local queue being moved + task.header().into() + } else { + // safety: The above CAS prevents a stealer from accessing these + // tasks and we are the only producer. + self.inner.buffer[j_idx].with(|ptr| unsafe { + let value = (*ptr).as_ptr(); + (*value).header().into() + }) + }; + + // safety: the above CAS prevents a stealer from accessing these + // tasks and we are the only producer. + self.inner.buffer[i_idx].with_mut(|ptr| unsafe { + let ptr = (*ptr).as_ptr(); + *(*ptr).header().queue_next.get() = Some(next); + }); + } + + // safety: the above CAS prevents a stealer from accessing these tasks + // and we are the only producer. + let head = self.inner.buffer[head as usize & MASK] + .with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); + + // Push the tasks onto the inject queue + inject.push_batch(head, task, BATCH_LEN); + + Ok(()) + } + + /// Pops a task from the local queue. + pub(super) fn pop(&mut self) -> Option> { + // If a task is available in the FIFO slot, return that. + if let Some(task) = self.next.take() { + return Some(task); + } + + loop { + let head = self.inner.head.load(Acquire); + + // safety: this is the **only** thread that updates this cell. + let tail = unsafe { self.inner.tail.unsync_load() }; + + if head == tail { + // queue is empty + return None; + } + + // Map the head position to a slot index. + let idx = head as usize & MASK; + + let task = self.inner.buffer[idx].with(|ptr| { + // Tentatively read the task at the head position. Note that we + // have not yet claimed the task. + // + // safety: reading this as uninitialized memory. + unsafe { ptr::read(ptr) } + }); + + // Attempt to claim the task read above. + let actual = self + .inner + .head + .compare_and_swap(head, head.wrapping_add(1), Release); + + if actual == head { + // safety: we claimed the task and the data we read is + // initialized memory. + return Some(unsafe { task.assume_init() }); + } + + atomic::spin_loop_hint(); + } + } +} + +impl Steal { + pub(super) fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// Steals half the tasks from self and place them into `dst`. + pub(super) fn steal_into(&self, dst: &mut Local) -> Option> { + // Safety: the caller is the only thread that mutates `dst.tail` and + // holds a mutable reference. + let dst_tail = unsafe { dst.inner.tail.unsync_load() }; + + // Steal the tasks into `dst`'s buffer. This does not yet expose the + // tasks in `dst`. + let mut n = self.steal_into2(dst, dst_tail); + + if n == 0 { + // No tasks were stolen + return None; + } + + // We are returning a task here + n -= 1; + + let ret_pos = dst_tail.wrapping_add(n); + let ret_idx = ret_pos as usize & MASK; + + // safety: the value was written as part of `steal_into2` and not + // exposed to stealers, so no other thread can access it. + let ret = dst.inner.buffer[ret_idx].with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); + + if n == 0 { + // The `dst` queue is empty, but a single task was stolen + return Some(ret); + } + + // Synchronize with stealers + let dst_head = dst.inner.head.load(Acquire); + + assert!(dst_tail.wrapping_sub(dst_head) + n <= LOCAL_QUEUE_CAPACITY as u32); + + // Make the stolen items available to consumers + dst.inner.tail.store(dst_tail.wrapping_add(n), Release); + + Some(ret) + } + + fn steal_into2(&self, dst: &mut Local, dst_tail: u32) -> u32 { + loop { + let src_head = self.0.head.load(Acquire); + let src_tail = self.0.tail.load(Acquire); + + // Number of available tasks to steal + let n = src_tail.wrapping_sub(src_head); + let n = n - n / 2; + + if n == 0 { + return 0; + } + + if n > LOCAL_QUEUE_CAPACITY as u32 / 2 { + atomic::spin_loop_hint(); + // inconsistent, try again + continue; + } + + // Track CausalCell causality checks. The check is deferred until + // the compare_and_swap claims ownership of the tasks. + let mut check = CausalCheck::default(); + + for i in 0..n { + // Compute the positions + let src_pos = src_head.wrapping_add(i); + let dst_pos = dst_tail.wrapping_add(i); + + // Map to slots + let src_idx = src_pos as usize & MASK; + let dst_idx = dst_pos as usize & MASK; + + // Read the task + // + // safety: this is being read as MaybeUninit -- potentially + // uninitialized memory (in the case a producer wraps). We don't + // assume it is initialized, but will just write the + // `MaybeUninit` in our slot below. + let (task, ch) = self.0.buffer[src_idx] + .with_deferred(|ptr| unsafe { ptr::read((*ptr).as_ptr()) }); + + check.join(ch); + + // Write the task to the new slot + // + // safety: `dst` queue is empty and we are the only producer to + // this queue. + dst.inner.buffer[dst_idx] + .with_mut(|ptr| unsafe { ptr::write((*ptr).as_mut_ptr(), task) }); + } + + // Claim all of those tasks! + let actual = self + .0 + .head + .compare_and_swap(src_head, src_head.wrapping_add(n), Release); + + if actual == src_head { + check.check(); + return n; + } + + atomic::spin_loop_hint(); + } + } +} + +impl Drop for Local { + fn drop(&mut self) { + if !std::thread::panicking() { + assert!(self.pop().is_none(), "queue not empty"); + } + } +} + +impl Inner { + fn is_empty(&self) -> bool { + let head = self.head.load(Acquire); + let tail = self.tail.load(Acquire); + + head == tail + } +} + +impl Inject { + pub(super) fn new() -> Inject { + Inject { + pointers: Mutex::new(Pointers { + is_closed: false, + head: None, + tail: None, + }), + len: AtomicUsize::new(0), + _p: PhantomData, + } + } + + pub(super) fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Close the injection queue, returns `true` if the queue is open when the + /// transition is made. + pub(super) fn close(&self) -> bool { + let mut p = self.pointers.lock().unwrap(); + + if p.is_closed { + return false; + } + + p.is_closed = true; + true + } + + pub(super) fn is_closed(&self) -> bool { + self.pointers.lock().unwrap().is_closed + } + + fn len(&self) -> usize { + self.len.load(Acquire) + } + + /// Pushes a value into the queue. + pub(super) fn push(&self, task: task::Notified) { + // Acquire queue lock + let mut p = self.pointers.lock().unwrap(); + + if p.is_closed { + // Drop the mutex to avoid a potential deadlock when + // re-entering. + drop(p); + drop(task); + return; + } + + // safety: only mutated with the lock held + let len = unsafe { self.len.unsync_load() }; + let task = task.into_raw(); + + // The next pointer should already be null + debug_assert!(get_next(task).is_none()); + + if let Some(tail) = p.tail { + set_next(tail, Some(task)); + } else { + p.head = Some(task); + } + + p.tail = Some(task); + + self.len.store(len + 1, Release); + } + + pub(super) fn push_batch( + &self, + batch_head: task::Notified, + batch_tail: task::Notified, + num: usize, + ) { + let batch_head = batch_head.into_raw(); + let batch_tail = batch_tail.into_raw(); + + debug_assert!(get_next(batch_tail).is_none()); + + let mut p = self.pointers.lock().unwrap(); + + if let Some(tail) = p.tail { + set_next(tail, Some(batch_head)); + } else { + p.head = Some(batch_head); + } + + p.tail = Some(batch_tail); + + // Increment the count. + // + // safety: All updates to the len atomic are guarded by the mutex. As + // such, a non-atomic load followed by a store is safe. + let len = unsafe { self.len.unsync_load() }; + + self.len.store(len + num, Release); + } + + pub(super) fn pop(&self) -> Option> { + // Fast path, if len == 0, then there are no values + if self.is_empty() { + return None; + } + + let mut p = self.pointers.lock().unwrap(); + + // It is possible to hit null here if another thread poped the last + // task between us checking `len` and acquiring the lock. + let task = p.head?; + + p.head = get_next(task); + + if p.head.is_none() { + p.tail = None; + } + + set_next(task, None); + + // Decrement the count. + // + // safety: All updates to the len atomic are guarded by the mutex. As + // such, a non-atomic load followed by a store is safe. + self.len + .store(unsafe { self.len.unsync_load() } - 1, Release); + + // safety: a `Notified` is pushed into the queue and now it is popped! + Some(unsafe { task::Notified::from_raw(task) }) + } +} + +impl Drop for Inject { + fn drop(&mut self) { + if !std::thread::panicking() { + assert!(self.pop().is_none(), "queue not empty"); + } + } +} + +fn get_next(header: NonNull) -> Option> { + unsafe { *header.as_ref().queue_next.get() } +} + +fn set_next(header: NonNull, val: Option>) { + unsafe { + *header.as_ref().queue_next.get() = val; + } +} diff --git a/tokio/src/runtime/thread_pool/queue/global.rs b/tokio/src/runtime/thread_pool/queue/global.rs deleted file mode 100644 index 7e16280aa..000000000 --- a/tokio/src/runtime/thread_pool/queue/global.rs +++ /dev/null @@ -1,209 +0,0 @@ -use crate::loom::sync::atomic::AtomicUsize; -use crate::loom::sync::Mutex; -use crate::task::{Header, Task}; - -use std::marker::PhantomData; -use std::ptr::{self, NonNull}; -use std::sync::atomic::Ordering::{Acquire, Release}; -use std::usize; - -pub(super) struct Queue { - /// Pointers to the head and tail of the queue - pointers: Mutex, - - /// Number of pending tasks in the queue. This helps prevent unnecessary - /// locking in the hot path. - /// - /// The LSB is a flag tracking whether or not the queue is open or not. - len: AtomicUsize, - - _p: PhantomData, -} - -struct Pointers { - head: *const Header, - tail: *const Header, -} - -const CLOSED: usize = 1; -const MAX_LEN: usize = usize::MAX >> 1; - -impl Queue { - pub(super) fn new() -> Queue { - Queue { - pointers: Mutex::new(Pointers { - head: ptr::null(), - tail: ptr::null(), - }), - len: AtomicUsize::new(0), - _p: PhantomData, - } - } - - pub(super) fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub(super) fn is_closed(&self) -> bool { - self.len.load(Acquire) & CLOSED == CLOSED - } - - /// Close the worker queue - pub(super) fn close(&self) -> bool { - // Acquire the lock - let p = self.pointers.lock().unwrap(); - - let len = unsafe { - // Set the queue as closed. Because all mutations are synchronized by - // the mutex, a read followed by a write is acceptable. - self.len.unsync_load() - }; - - let ret = len & CLOSED == 0; - - self.len.store(len | CLOSED, Release); - - drop(p); - - ret - } - - fn len(&self) -> usize { - self.len.load(Acquire) >> 1 - } - - pub(super) fn wait_for_unlocked(&self) { - // Acquire and release the lock immediately. This synchronizes the - // caller **after** all external waiters are done w/ the scheduler - // struct. - drop(self.pointers.lock().unwrap()); - } - - /// Pushes a value into the queue and call the closure **while still holding - /// the push lock** - pub(super) fn push(&self, task: Task, f: F) - where - F: FnOnce(Result<(), Task>), - { - unsafe { - // Acquire queue lock - let mut p = self.pointers.lock().unwrap(); - - // Check if the queue is closed. This must happen in the lock. - let len = self.len.unsync_load(); - if len & CLOSED == CLOSED { - drop(p); - f(Err(task)); - return; - } - - let task = task.into_raw(); - - // The next pointer should already be null - debug_assert!(get_next(task).is_null()); - - if let Some(tail) = NonNull::new(p.tail as *mut _) { - set_next(tail, task.as_ptr()); - } else { - p.head = task.as_ptr(); - } - - p.tail = task.as_ptr(); - - // Increment the count. - // - // All updates to the len atomic are guarded by the mutex. As such, - // a non-atomic load followed by a store is safe. - // - // We increment by 2 to avoid touching the shutdown flag - if (len >> 1) == MAX_LEN { - eprintln!("[ERROR] overflowed task counter. This is a bug and should be reported."); - std::process::abort(); - } - - self.len.store(len + 2, Release); - - f(Ok(())); - - drop(p); - } - } - - pub(super) fn push_batch(&self, batch_head: Task, batch_tail: Task, num: usize) { - unsafe { - let batch_head = batch_head.into_raw().as_ptr(); - let batch_tail = batch_tail.into_raw(); - - debug_assert!(get_next(batch_tail).is_null()); - - let mut p = self.pointers.lock().unwrap(); - - if let Some(tail) = NonNull::new(p.tail as *mut _) { - set_next(tail, batch_head); - } else { - p.head = batch_head; - } - - p.tail = batch_tail.as_ptr(); - - // Increment the count. - // - // All updates to the len atomic are guarded by the mutex. As such, - // a non-atomic load followed by a store is safe. - // - // Left shift by 1 to avoid touching the shutdown flag. - let len = self.len.unsync_load(); - - if (len >> 1) >= (MAX_LEN - num) { - std::process::abort(); - } - - self.len.store(len + (num << 1), Release); - - drop(p); - } - } - - pub(super) fn pop(&self) -> Option> { - // Fast path, if len == 0, then there are no values - if self.is_empty() { - return None; - } - - unsafe { - let mut p = self.pointers.lock().unwrap(); - - // It is possible to hit null here if another thread poped the last - // task between us checking `len` and acquiring the lock. - let task = NonNull::new(p.head as *mut _)?; - - p.head = get_next(task); - - if p.head.is_null() { - p.tail = ptr::null(); - } - - set_next(task, ptr::null()); - - // Decrement the count. - // - // All updates to the len atomic are guarded by the mutex. As such, - // a non-atomic load followed by a store is safe. - // - // Decrement by 2 to avoid touching the shutdown flag - self.len.store(self.len.unsync_load() - 2, Release); - - drop(p); - - Some(Task::from_raw(task)) - } - } -} - -unsafe fn get_next(meta: NonNull
) -> *const Header { - *meta.as_ref().queue_next.get() -} - -unsafe fn set_next(meta: NonNull
, val: *const Header) { - *meta.as_ref().queue_next.get() = val; -} diff --git a/tokio/src/runtime/thread_pool/queue/inject.rs b/tokio/src/runtime/thread_pool/queue/inject.rs deleted file mode 100644 index d83084643..000000000 --- a/tokio/src/runtime/thread_pool/queue/inject.rs +++ /dev/null @@ -1,41 +0,0 @@ -use crate::loom::sync::Arc; -use crate::runtime::thread_pool::queue::Cluster; -use crate::task::Task; - -pub(crate) struct Inject { - cluster: Arc>, -} - -impl Inject { - pub(super) fn new(cluster: Arc>) -> Inject { - Inject { cluster } - } - - /// Pushes a value onto the queue - pub(crate) fn push(&self, task: Task, f: F) - where - F: FnOnce(Result<(), Task>), - { - self.cluster.global.push(task, f) - } - - /// Checks if the queue has been closed - pub(crate) fn is_closed(&self) -> bool { - self.cluster.global.is_closed() - } - - /// Closes the queue - /// - /// Returns `true` if the channel was closed. `false` indicates the pool was - /// previously closed. - pub(crate) fn close(&self) -> bool { - self.cluster.global.close() - } - - /// Waits for all locks on the queue to drop. - /// - /// This is done by locking w/o doing anything. - pub(crate) fn wait_for_unlocked(&self) { - self.cluster.global.wait_for_unlocked(); - } -} diff --git a/tokio/src/runtime/thread_pool/queue/local.rs b/tokio/src/runtime/thread_pool/queue/local.rs deleted file mode 100644 index e913c4b0e..000000000 --- a/tokio/src/runtime/thread_pool/queue/local.rs +++ /dev/null @@ -1,298 +0,0 @@ -use crate::loom::cell::{CausalCell, CausalCheck}; -use crate::loom::sync::atomic::{self, AtomicU32}; -use crate::runtime::thread_pool::queue::global; -use crate::runtime::thread_pool::LOCAL_QUEUE_CAPACITY; -use crate::task::Task; - -use std::fmt; -use std::mem::MaybeUninit; -use std::ptr; -use std::sync::atomic::Ordering::{Acquire, Release}; - -pub(super) struct Queue { - /// Concurrently updated by many threads. - head: AtomicU32, - - /// Only updated by producer thread but read by many threads. - tail: AtomicU32, - - /// Elements - buffer: Box<[CausalCell>>]>, -} - -const MASK: usize = LOCAL_QUEUE_CAPACITY - 1; - -impl Queue { - pub(super) fn new() -> Queue { - debug_assert!(LOCAL_QUEUE_CAPACITY >= 2 && LOCAL_QUEUE_CAPACITY.is_power_of_two()); - - let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY); - - for _ in 0..LOCAL_QUEUE_CAPACITY { - buffer.push(CausalCell::new(MaybeUninit::uninit())); - } - - Queue { - head: AtomicU32::new(0), - tail: AtomicU32::new(0), - buffer: buffer.into(), - } - } -} - -impl Queue { - /// Pushes a task onto the local queue. - /// - /// This **must** be called by the producer thread. - pub(super) unsafe fn push(&self, mut task: Task, global: &global::Queue) { - loop { - let head = self.head.load(Acquire); - - // safety: this is the **only** thread that updates this cell. - let tail = self.tail.unsync_load(); - - if tail.wrapping_sub(head) < LOCAL_QUEUE_CAPACITY as u32 { - // Map the position to a slot index. - let idx = tail as usize & MASK; - - self.buffer[idx].with_mut(|ptr| { - // Write the task to the slot - ptr::write((*ptr).as_mut_ptr(), task); - }); - - // Make the task available - self.tail.store(tail.wrapping_add(1), Release); - - return; - } - - // The local buffer is full. Push a batch of work to the global - // queue. - match self.push_overflow(task, head, tail, global) { - Ok(_) => return, - // Lost the race, try again - Err(v) => task = v, - } - - atomic::spin_loop_hint(); - } - } - - /// Moves a batch of tasks into the global queue. - /// - /// This will temporarily make some of the tasks unavailable to stealers. - /// Once `push_overflow` is done, a notification is sent out, so if other - /// workers "missed" some of the tasks during a steal, they will get - /// another opportunity. - #[inline(never)] - unsafe fn push_overflow( - &self, - task: Task, - head: u32, - tail: u32, - global: &global::Queue, - ) -> Result<(), Task> { - const BATCH_LEN: usize = LOCAL_QUEUE_CAPACITY / 2 + 1; - - let n = tail.wrapping_sub(head) / 2; - assert_eq!(n as usize, LOCAL_QUEUE_CAPACITY / 2, "queue is not full"); - - // Claim a bunch of tasks - // - // We are claiming the tasks **before** reading them out of the buffer. - // This is safe because only the **current** thread is able to push new - // tasks. - // - // There isn't really any need for memory ordering... Relaxed would - // work. This is because all tasks are pushed into the queue from the - // current thread (or memory has been acquired if the local queue handle - // moved). - let actual = self.head.compare_and_swap(head, head + n, Release); - if actual != head { - // We failed to claim the tasks, losing the race. Return out of - // this function and try the full `push` routine again. The queue - // may not be full anymore. - return Err(task); - } - - // link the tasks - for i in 0..n { - let j = i + 1; - - let i_idx = (i + head) as usize & MASK; - let j_idx = (j + head) as usize & MASK; - - // Get the next pointer - let next = if j == n { - // The last task in the local queue being moved - task.header() as *const _ - } else { - self.buffer[j_idx].with(|ptr| { - let value = (*ptr).as_ptr(); - (*value).header() as *const _ - }) - }; - - self.buffer[i_idx].with_mut(|ptr| { - let ptr = (*ptr).as_ptr(); - debug_assert!((*(*ptr).header().queue_next.get()).is_null()); - *(*ptr).header().queue_next.get() = next; - }); - } - - let head = self.buffer[head as usize & MASK].with(|ptr| ptr::read((*ptr).as_ptr())); - - // Push the tasks onto the global queue - global.push_batch(head, task, BATCH_LEN); - - Ok(()) - } - - /// Pops a task from the local queue. - /// - /// This **must** be called by the producer thread - pub(super) unsafe fn pop(&self) -> Option> { - loop { - let head = self.head.load(Acquire); - - // safety: this is the **only** thread that updates this cell. - let tail = self.tail.unsync_load(); - - if head == tail { - // queue is empty - return None; - } - - // Map the head position to a slot index. - let idx = head as usize & MASK; - - let task = self.buffer[idx].with(|ptr| { - // Tentatively read the task at the head position. Note that we - // have not yet claimed the task. - // - ptr::read(ptr) - }); - - // Attempt to claim the task read above. - let actual = self - .head - .compare_and_swap(head, head.wrapping_add(1), Release); - - if actual == head { - return Some(task.assume_init()); - } - - atomic::spin_loop_hint(); - } - } - - pub(super) fn is_empty(&self) -> bool { - let head = self.head.load(Acquire); - let tail = self.tail.load(Acquire); - - head == tail - } - - /// Steals half the tasks from self and place them into `dst`. - pub(super) unsafe fn steal(&self, dst: &Queue) -> Option> { - let dst_tail = dst.tail.unsync_load(); - - // Steal the tasks into `dst`'s buffer. This does not yet expose the - // tasks in `dst`. - let mut n = self.steal2(dst, dst_tail); - - if n == 0 { - // No tasks were stolen - return None; - } - - // We are returning a task here - n -= 1; - - let ret_pos = dst_tail.wrapping_add(n); - let ret_idx = ret_pos as usize & MASK; - - let ret = dst.buffer[ret_idx].with(|ptr| ptr::read((*ptr).as_ptr())); - - if n == 0 { - // The `dst` queue is empty, but a single task was stolen - return Some(ret); - } - - // Synchronize with stealers - let dst_head = dst.head.load(Acquire); - - assert!(dst_tail.wrapping_sub(dst_head) + n <= LOCAL_QUEUE_CAPACITY as u32); - - // Make the stolen items available to consumers - dst.tail.store(dst_tail.wrapping_add(n), Release); - - Some(ret) - } - - unsafe fn steal2(&self, dst: &Queue, dst_tail: u32) -> u32 { - loop { - let src_head = self.head.load(Acquire); - let src_tail = self.tail.load(Acquire); - - // Number of available tasks to steal - let n = src_tail.wrapping_sub(src_head); - let n = n - n / 2; - - if n == 0 { - return 0; - } - - if n > LOCAL_QUEUE_CAPACITY as u32 / 2 { - atomic::spin_loop_hint(); - // inconsistent, try again - continue; - } - - // Track CausalCell causality checks. The check is deferred until - // the compare_and_swap claims ownership of the tasks. - let mut check = CausalCheck::default(); - - for i in 0..n { - // Compute the positions - let src_pos = src_head.wrapping_add(i); - let dst_pos = dst_tail.wrapping_add(i); - - // Map to slots - let src_idx = src_pos as usize & MASK; - let dst_idx = dst_pos as usize & MASK; - - // Read the task - let (task, ch) = - self.buffer[src_idx].with_deferred(|ptr| ptr::read((*ptr).as_ptr())); - - check.join(ch); - - // Write the task to the new slot - dst.buffer[dst_idx].with_mut(|ptr| ptr::write((*ptr).as_mut_ptr(), task)); - } - - // Claim all of those tasks! - let actual = self - .head - .compare_and_swap(src_head, src_head.wrapping_add(n), Release); - - if actual == src_head { - check.check(); - return n; - } - - atomic::spin_loop_hint(); - } - } -} - -impl fmt::Debug for Queue { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("local::Queue") - .field("head", &self.head) - .field("tail", &self.tail) - .field("buffer", &"[...]") - .finish() - } -} diff --git a/tokio/src/runtime/thread_pool/queue/mod.rs b/tokio/src/runtime/thread_pool/queue/mod.rs deleted file mode 100644 index 88633ee39..000000000 --- a/tokio/src/runtime/thread_pool/queue/mod.rs +++ /dev/null @@ -1,41 +0,0 @@ -//! The threadpool's task queue system. - -mod global; -mod inject; -mod local; -mod worker; - -pub(crate) use self::inject::Inject; -pub(crate) use self::worker::Worker; - -use crate::loom::sync::Arc; - -pub(crate) fn build(workers: usize) -> Vec> { - let local: Vec<_> = (0..workers).map(|_| local::Queue::new()).collect(); - - let cluster = Arc::new(Cluster { - local: local.into_boxed_slice(), - global: global::Queue::new(), - }); - - (0..workers) - .map(|index| Worker::new(cluster.clone(), index)) - .collect() -} - -struct Cluster { - /// per-worker local queues - local: Box<[local::Queue]>, - global: global::Queue, -} - -impl Drop for Cluster { - fn drop(&mut self) { - // Drain all the queues - for queue in &self.local[..] { - while let Some(_) = unsafe { queue.pop() } {} - } - - while let Some(_) = self.global.pop() {} - } -} diff --git a/tokio/src/runtime/thread_pool/queue/worker.rs b/tokio/src/runtime/thread_pool/queue/worker.rs deleted file mode 100644 index 6d3648967..000000000 --- a/tokio/src/runtime/thread_pool/queue/worker.rs +++ /dev/null @@ -1,127 +0,0 @@ -use crate::loom::sync::Arc; -use crate::runtime::thread_pool::queue::{local, Cluster, Inject}; -use crate::task::Task; - -use std::cell::Cell; -use std::fmt; - -pub(crate) struct Worker { - cluster: Arc>, - index: u16, - /// Task to pop next - next: Cell>>, -} - -impl Worker { - pub(super) fn new(cluster: Arc>, index: usize) -> Worker { - Worker { - cluster, - index: index as u16, - next: Cell::new(None), - } - } - - pub(crate) fn injector(&self) -> Inject { - Inject::new(self.cluster.clone()) - } - - /// Returns `true` if the queue is closed - pub(crate) fn is_closed(&self) -> bool { - self.cluster.global.is_closed() - } - - /// Pushes to the local queue. - /// - /// If the local queue is full, the task is pushed onto the global queue. - /// - /// # Return - /// - /// Returns `true` if the pushed task can be stolen by another worker. - pub(crate) fn push(&self, task: Task) -> bool { - let prev = self.next.take(); - let ret = prev.is_some(); - - if let Some(prev) = prev { - // safety: we guarantee that only one thread pushes to this local - // queue at a time. - unsafe { - self.local().push(prev, &self.cluster.global); - } - } - - self.next.set(Some(task)); - - ret - } - - pub(crate) fn push_yield(&self, task: Task) { - unsafe { self.local().push(task, &self.cluster.global) } - } - - /// Pops a task checking the local queue first. - pub(crate) fn pop_local_first(&self) -> Option> { - self.local_pop().or_else(|| self.cluster.global.pop()) - } - - /// Pops a task checking the global queue first. - pub(crate) fn pop_global_first(&self) -> Option> { - self.cluster.global.pop().or_else(|| self.local_pop()) - } - - /// Steals from other local queues. - /// - /// `start` specifies the queue from which to start stealing. - pub(crate) fn steal(&self, start: usize) -> Option> { - let num_queues = self.cluster.local.len(); - - for i in 0..num_queues { - let i = (start + i) % num_queues; - - if i == self.index as usize { - continue; - } - - // safety: we own the dst queue - let ret = unsafe { self.cluster.local[i].steal(self.local()) }; - - if ret.is_some() { - return ret; - } - } - - None - } - - /// An approximation of whether or not the queue is empty. - pub(crate) fn is_empty(&self) -> bool { - for local_queue in &self.cluster.local[..] { - if !local_queue.is_empty() { - return false; - } - } - - self.cluster.global.is_empty() - } - - fn local_pop(&self) -> Option> { - if let Some(task) = self.next.take() { - return Some(task); - } - // safety: we guarantee that only one thread pushes to this local queue - // at a time. - unsafe { self.local().pop() } - } - - fn local(&self) -> &local::Queue { - &self.cluster.local[self.index as usize] - } -} - -impl fmt::Debug for Worker { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("queue::Worker") - .field("cluster", &"...") - .field("index", &self.index) - .finish() - } -} diff --git a/tokio/src/runtime/thread_pool/shared.rs b/tokio/src/runtime/thread_pool/shared.rs deleted file mode 100644 index f61bafaba..000000000 --- a/tokio/src/runtime/thread_pool/shared.rs +++ /dev/null @@ -1,94 +0,0 @@ -use crate::park::Unpark; -use crate::runtime::thread_pool::slice; -use crate::runtime::Unparker; -use crate::task::{self, Schedule, ScheduleSendOnly, Task}; - -use std::ptr; - -/// Per-worker data accessible from any thread. -/// -/// Accessed by: -/// -/// - other workers -/// - tasks -/// -pub(crate) struct Shared { - /// Thread unparker - unpark: Unparker, - - /// Tasks pending drop. Any worker pushes tasks, only the "owning" worker - /// pops. - pub(super) pending_drop: task::TransferStack, - - /// Untracked pointer to the pool. - /// - /// The slice::Set itself is tracked by an `Arc`, but this pointer is not - /// included in the ref count. - slices: *const slice::Set, -} - -unsafe impl Send for Shared {} -unsafe impl Sync for Shared {} - -impl Shared { - pub(super) fn new(unpark: Unparker) -> Shared { - Shared { - unpark, - pending_drop: task::TransferStack::new(), - slices: ptr::null(), - } - } - - pub(crate) fn schedule(&self, task: Task) { - self.slices().schedule(task); - } - - pub(super) fn unpark(&self) { - self.unpark.unpark(); - } - - fn slices(&self) -> &slice::Set { - unsafe { &*self.slices } - } - - pub(super) fn set_slices_ptr(&mut self, slices: *const slice::Set) { - self.slices = slices; - } -} - -impl Schedule for Shared { - fn bind(&self, task: &Task) { - // Get access to the Owned component. This function can only be called - // when on the worker. - unsafe { - let index = self.slices().index_of(self); - let owned = &mut *self.slices().owned()[index].get(); - - owned.bind_task(task); - } - } - - fn release(&self, task: Task) { - // This stores the task with the owning worker. The worker is not - // notified. Instead, the worker will clean up the tasks "eventually". - // - self.pending_drop.push(task); - } - - fn release_local(&self, task: &Task) { - // Get access to the Owned component. This function can only be called - // when on the worker. - unsafe { - let index = self.slices().index_of(self); - let owned = &mut *self.slices().owned()[index].get(); - - owned.release_task(task); - } - } - - fn schedule(&self, task: Task) { - Self::schedule(self, task); - } -} - -impl ScheduleSendOnly for Shared {} diff --git a/tokio/src/runtime/thread_pool/shutdown.rs b/tokio/src/runtime/thread_pool/shutdown.rs deleted file mode 100644 index 414c1c84a..000000000 --- a/tokio/src/runtime/thread_pool/shutdown.rs +++ /dev/null @@ -1,44 +0,0 @@ -//! A shutdown channel. -//! -//! Each worker holds the `Sender` half. When all the `Sender` halves are -//! dropped, the `Receiver` receives a notification. - -use crate::loom::sync::Arc; -use crate::sync::oneshot; - -#[derive(Debug, Clone)] -pub(super) struct Sender { - tx: Arc>, -} - -#[derive(Debug)] -pub(super) struct Receiver { - rx: oneshot::Receiver<()>, -} - -pub(super) fn channel() -> (Sender, Receiver) { - let (tx, rx) = oneshot::channel(); - let tx = Sender { tx: Arc::new(tx) }; - let rx = Receiver { rx }; - - (tx, rx) -} - -impl Receiver { - /// Blocks the current thread until all `Sender` handles drop. - pub(crate) fn wait(&mut self) { - use crate::runtime::enter::{enter, try_enter}; - - let mut e = if std::thread::panicking() { - match try_enter() { - Some(enter) => enter, - _ => return, - } - } else { - enter() - }; - - // The oneshot completes with an Err - let _ = e.block_on(&mut self.rx); - } -} diff --git a/tokio/src/runtime/thread_pool/slice.rs b/tokio/src/runtime/thread_pool/slice.rs deleted file mode 100644 index 9a5fd334f..000000000 --- a/tokio/src/runtime/thread_pool/slice.rs +++ /dev/null @@ -1,172 +0,0 @@ -//! The scheduler is divided into multiple slices. Each slice is fairly -//! isolated, having its own queue. A worker is dedicated to processing a single -//! slice. - -use crate::loom::rand::seed; -use crate::park::Park; -use crate::runtime::thread_pool::{current, queue, Idle, Owned, Shared}; -use crate::runtime::Parker; -use crate::task::{self, JoinHandle, Task}; -use crate::util::{CachePadded, FastRand}; - -use std::cell::UnsafeCell; -use std::future::Future; - -pub(super) struct Set { - /// Data accessible from all workers. - shared: Box<[Shared]>, - - /// Data owned by the worker. - owned: Box<[UnsafeCell>]>, - - /// Submit work to the pool while *not* currently on a worker thread. - inject: queue::Inject, - - /// Coordinates idle workers - idle: Idle, -} - -unsafe impl Send for Set {} -unsafe impl Sync for Set {} - -impl Set { - /// Creates a new worker set using the provided queues. - pub(crate) fn new(parkers: &[Parker]) -> Self { - assert!(!parkers.is_empty()); - - let queues = queue::build(parkers.len()); - let inject = queues[0].injector(); - - let mut shared = Vec::with_capacity(queues.len()); - let mut owned = Vec::with_capacity(queues.len()); - - for (i, queue) in queues.into_iter().enumerate() { - let rand = FastRand::new(seed()); - - shared.push(Shared::new(parkers[i].unpark())); - owned.push(UnsafeCell::new(CachePadded::new(Owned::new(queue, rand)))); - } - - Set { - shared: shared.into_boxed_slice(), - owned: owned.into_boxed_slice(), - inject, - idle: Idle::new(parkers.len()), - } - } - - pub(crate) fn spawn_typed(&self, future: F) -> JoinHandle - where - F: Future + Send + 'static, - F::Output: Send + 'static, - { - let (task, handle) = task::joinable(future); - self.schedule(task); - handle - } - - fn inject_task(&self, task: Task) { - self.inject.push(task, |res| { - if let Err(task) = res { - task.shutdown(); - - // There may be a worker, in the process of being shutdown, that is - // waiting for this task to be released, so we notify all workers - // just in case. - // - // Over aggressive, but the runtime is in the process of shutting - // down, so efficiency is not critical. - self.notify_all(); - } else { - self.notify_work(); - } - }); - } - - pub(super) fn notify_work(&self) { - if let Some(index) = self.idle.worker_to_notify() { - self.shared[index].unpark(); - } - } - - pub(super) fn notify_all(&self) { - for shared in &self.shared[..] { - shared.unpark(); - } - } - - pub(crate) fn schedule(&self, task: Task) { - current::get(|current_worker| match current_worker.as_member(self) { - Some(worker) => { - if worker.submit_local(task) { - self.notify_work(); - } - } - None => { - self.inject_task(task); - } - }) - } - - pub(crate) fn set_ptr(&mut self) { - let ptr = self as *const _; - for shared in &mut self.shared[..] { - shared.set_slices_ptr(ptr); - } - } - - /// Signals the pool is closed - /// - /// Returns `true` if the transition to closed is successful. `false` - /// indicates the pool was already closed. - pub(crate) fn close(&self) -> bool { - if self.inject.close() { - self.notify_all(); - true - } else { - false - } - } - - pub(crate) fn is_closed(&self) -> bool { - self.inject.is_closed() - } - - pub(crate) fn len(&self) -> usize { - self.shared.len() - } - - pub(super) fn index_of(&self, shared: &Shared) -> usize { - use std::mem; - - let size = mem::size_of::(); - - ((shared as *const _ as usize) - (&self.shared[0] as *const _ as usize)) / size - } - - pub(super) fn shared(&self) -> &[Shared] { - &self.shared - } - - pub(super) fn owned(&self) -> &[UnsafeCell>] { - &self.owned - } - - pub(super) fn idle(&self) -> &Idle { - &self.idle - } - - /// Waits for all locks on the injection queue to drop. - /// - /// This is done by locking w/o doing anything. - pub(super) fn wait_for_unlocked(&self) { - self.inject.wait_for_unlocked(); - } -} - -impl Drop for Set { - fn drop(&mut self) { - // Before proceeding, wait for all concurrent wakers to exit - self.wait_for_unlocked(); - } -} diff --git a/tokio/src/runtime/thread_pool/spawner.rs b/tokio/src/runtime/thread_pool/spawner.rs deleted file mode 100644 index 56931c9ba..000000000 --- a/tokio/src/runtime/thread_pool/spawner.rs +++ /dev/null @@ -1,49 +0,0 @@ -use crate::loom::sync::Arc; -use crate::runtime::thread_pool::slice; -use crate::task::JoinHandle; - -use std::fmt; -use std::future::Future; - -/// Submit futures to the associated thread pool for execution. -/// -/// A `Spawner` instance is a handle to a single thread pool, allowing the owner -/// of the handle to spawn futures onto the thread pool. -/// -/// The `Spawner` handle is *only* used for spawning new futures. It does not -/// impact the lifecycle of the thread pool in any way. The thread pool may -/// shutdown while there are outstanding `Spawner` instances. -/// -/// `Spawner` instances are obtained by calling [`ThreadPool::spawner`]. -/// -/// [`ThreadPool::spawner`]: struct.ThreadPool.html#method.spawner -#[derive(Clone)] -pub(crate) struct Spawner { - workers: Arc, -} - -impl Spawner { - pub(super) fn new(workers: Arc) -> Spawner { - Spawner { workers } - } - - /// Spawns a future onto the thread pool - pub(crate) fn spawn(&self, future: F) -> JoinHandle - where - F: Future + Send + 'static, - F::Output: Send + 'static, - { - self.workers.spawn_typed(future) - } - - /// Reference to the worker set. Used by `ThreadPool` to initiate shutdown. - pub(super) fn workers(&self) -> &slice::Set { - &*self.workers - } -} - -impl fmt::Debug for Spawner { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Spawner").finish() - } -} diff --git a/tokio/src/runtime/thread_pool/tests/loom_pool.rs b/tokio/src/runtime/thread_pool/tests/loom_pool.rs deleted file mode 100644 index 98765ac02..000000000 --- a/tokio/src/runtime/thread_pool/tests/loom_pool.rs +++ /dev/null @@ -1,308 +0,0 @@ -use crate::runtime::tests::loom_oneshot as oneshot; -use crate::runtime::{self, Runtime}; -use crate::spawn; - -use loom::sync::atomic::{AtomicBool, AtomicUsize}; -use loom::sync::{Arc, Mutex}; - -use std::future::Future; -use std::sync::atomic::Ordering::{Acquire, Relaxed, Release}; - -#[test] -fn racy_shutdown() { - loom::model(|| { - let pool = mk_pool(1); - - // here's the case we want to exercise: - // - // a worker that still has tasks in its local queue gets sent to the blocking pool (due to - // block_in_place). the blocking pool is shut down, so drops the worker. the worker's - // shutdown method never gets run. - // - // we do this by spawning two tasks on one worker, the first of which does block_in_place, - // and then immediately drop the pool. - - pool.spawn(async { - crate::task::block_in_place(|| {}); - }); - pool.spawn(async {}); - drop(pool); - }); -} - -#[test] -fn pool_multi_spawn() { - loom::model(|| { - let pool = mk_pool(2); - let c1 = Arc::new(AtomicUsize::new(0)); - - let (tx, rx) = oneshot::channel(); - let tx1 = Arc::new(Mutex::new(Some(tx))); - - // Spawn a task - let c2 = c1.clone(); - let tx2 = tx1.clone(); - pool.spawn(async move { - spawn(async move { - if 1 == c1.fetch_add(1, Relaxed) { - tx1.lock().unwrap().take().unwrap().send(()); - } - }); - }); - - // Spawn a second task - pool.spawn(async move { - spawn(async move { - if 1 == c2.fetch_add(1, Relaxed) { - tx2.lock().unwrap().take().unwrap().send(()); - } - }); - }); - - rx.recv(); - }); -} - -fn only_blocking_inner(first_pending: bool) { - loom::model(move || { - let pool = mk_pool(1); - let (block_tx, block_rx) = oneshot::channel(); - - pool.spawn(async move { - crate::task::block_in_place(move || { - block_tx.send(()); - }); - if first_pending { - yield_once().await - } - }); - - block_rx.recv(); - drop(pool); - }); -} - -#[test] -fn only_blocking() { - only_blocking_inner(false) -} - -#[test] -fn only_blocking_with_pending() { - only_blocking_inner(true) -} - -fn blocking_and_regular_inner(first_pending: bool) { - const NUM: usize = 3; - loom::model(move || { - let pool = mk_pool(1); - let cnt = Arc::new(AtomicUsize::new(0)); - - let (block_tx, block_rx) = oneshot::channel(); - let (done_tx, done_rx) = oneshot::channel(); - let done_tx = Arc::new(Mutex::new(Some(done_tx))); - - pool.spawn(async move { - crate::task::block_in_place(move || { - block_tx.send(()); - }); - if first_pending { - yield_once().await - } - }); - - for _ in 0..NUM { - let cnt = cnt.clone(); - let done_tx = done_tx.clone(); - - pool.spawn(async move { - if NUM == cnt.fetch_add(1, Relaxed) + 1 { - done_tx.lock().unwrap().take().unwrap().send(()); - } - }); - } - - done_rx.recv(); - block_rx.recv(); - - drop(pool); - }); -} - -#[test] -fn blocking_and_regular() { - blocking_and_regular_inner(false); -} - -#[test] -fn blocking_and_regular_with_pending() { - blocking_and_regular_inner(true); -} - -#[test] -fn pool_multi_notify() { - loom::model(|| { - let pool = mk_pool(2); - - let c1 = Arc::new(AtomicUsize::new(0)); - - let (done_tx, done_rx) = oneshot::channel(); - let done_tx1 = Arc::new(Mutex::new(Some(done_tx))); - - // Spawn a task - let c2 = c1.clone(); - let done_tx2 = done_tx1.clone(); - pool.spawn(async move { - gated().await; - gated().await; - - if 1 == c1.fetch_add(1, Relaxed) { - done_tx1.lock().unwrap().take().unwrap().send(()); - } - }); - - // Spawn a second task - pool.spawn(async move { - gated().await; - gated().await; - - if 1 == c2.fetch_add(1, Relaxed) { - done_tx2.lock().unwrap().take().unwrap().send(()); - } - }); - - done_rx.recv(); - }); -} - -#[test] -fn pool_shutdown() { - loom::model(|| { - let pool = mk_pool(2); - - pool.spawn(async move { - gated2(true).await; - }); - - pool.spawn(async move { - gated2(false).await; - }); - - drop(pool); - }); -} - -#[test] -fn complete_block_on_under_load() { - use futures::FutureExt; - - loom::model(|| { - let mut pool = mk_pool(2); - - pool.block_on({ - futures::future::lazy(|_| ()).then(|_| { - // Spin hard - crate::spawn(async { - for _ in 0..2 { - yield_once().await; - } - }); - - gated2(true) - }) - }); - }); -} - -#[test] -fn shutdown_with_notification() { - use crate::stream::StreamExt; - use crate::sync::{mpsc, oneshot}; - - loom::model(|| { - let rt = mk_pool(2); - let (done_tx, done_rx) = oneshot::channel::<()>(); - - rt.spawn(async move { - let (mut tx, mut rx) = mpsc::channel::<()>(10); - - crate::spawn(async move { - crate::task::spawn_blocking(move || { - let _ = tx.try_send(()); - }); - - let _ = done_rx.await; - }); - - while let Some(_) = rx.next().await {} - - let _ = done_tx.send(()); - }); - }); -} - -fn mk_pool(num_threads: usize) -> Runtime { - runtime::Builder::new() - .threaded_scheduler() - .core_threads(num_threads) - .build() - .unwrap() -} - -use futures::future::poll_fn; -use std::task::Poll; -async fn yield_once() { - let mut yielded = false; - poll_fn(|cx| { - if yielded { - Poll::Ready(()) - } else { - loom::thread::yield_now(); - yielded = true; - cx.waker().wake_by_ref(); - Poll::Pending - } - }) - .await -} - -fn gated() -> impl Future { - gated2(false) -} - -fn gated2(thread: bool) -> impl Future { - use loom::thread; - use std::sync::Arc; - - let gate = Arc::new(AtomicBool::new(false)); - let mut fired = false; - - poll_fn(move |cx| { - if !fired { - let gate = gate.clone(); - let waker = cx.waker().clone(); - - if thread { - thread::spawn(move || { - gate.store(true, Release); - waker.wake_by_ref(); - }); - } else { - spawn(async move { - gate.store(true, Release); - waker.wake_by_ref(); - }); - } - - fired = true; - - return Poll::Pending; - } - - if gate.load(Acquire) { - Poll::Ready("hello world") - } else { - Poll::Pending - } - }) -} diff --git a/tokio/src/runtime/thread_pool/tests/loom_queue.rs b/tokio/src/runtime/thread_pool/tests/loom_queue.rs deleted file mode 100644 index a4e106201..000000000 --- a/tokio/src/runtime/thread_pool/tests/loom_queue.rs +++ /dev/null @@ -1,69 +0,0 @@ -use crate::runtime::thread_pool::queue; -use crate::task::{self, Task}; -use crate::tests::mock_schedule::{Noop, NOOP_SCHEDULE}; - -use loom::thread; - -use std::cell::Cell; -use std::rc::Rc; - -#[test] -fn multi_worker() { - const THREADS: usize = 2; - const PER_THREAD: usize = 7; - - fn work(_i: usize, q: queue::Worker, rem: Rc>) { - let mut rem_local = PER_THREAD; - - while rem.get() != 0 { - for _ in 0..3 { - if rem_local > 0 { - q.push(val(0)); - rem_local -= 1; - } - } - - // Try to work - while let Some(task) = q.pop_local_first() { - assert!(task.run(&mut || Some(From::from(&NOOP_SCHEDULE))).is_none()); - let r = rem.get(); - assert!(r > 0); - rem.set(r - 1); - } - - // Try to steal - if let Some(task) = q.steal(0) { - assert!(task.run(&mut || Some(From::from(&NOOP_SCHEDULE))).is_none()); - let r = rem.get(); - assert!(r > 0); - rem.set(r - 1); - } - - thread::yield_now(); - } - } - - loom::model(|| { - let rem = Rc::new(Cell::new(THREADS * PER_THREAD)); - - let mut qs = queue::build(THREADS); - let q1 = qs.remove(0); - - for i in 1..THREADS { - let q = qs.remove(0); - let rem = rem.clone(); - thread::spawn(move || { - work(i, q, rem); - }); - } - - work(0, q1, rem); - - // th.join().unwrap(); - }); -} - -fn val(num: u32) -> Task { - let (task, _) = task::joinable(async move { num }); - task -} diff --git a/tokio/src/runtime/thread_pool/tests/mod.rs b/tokio/src/runtime/thread_pool/tests/mod.rs deleted file mode 100644 index 6638c5587..000000000 --- a/tokio/src/runtime/thread_pool/tests/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -#[cfg(loom)] -mod loom_pool; - -#[cfg(loom)] -mod loom_queue; - -#[cfg(not(loom))] -mod queue; diff --git a/tokio/src/runtime/thread_pool/tests/queue.rs b/tokio/src/runtime/thread_pool/tests/queue.rs deleted file mode 100644 index 83a602758..000000000 --- a/tokio/src/runtime/thread_pool/tests/queue.rs +++ /dev/null @@ -1,277 +0,0 @@ -use crate::runtime::thread_pool::{queue, LOCAL_QUEUE_CAPACITY}; -use crate::task::{self, Task}; -use crate::tests::mock_schedule::{Noop, NOOP_SCHEDULE}; - -macro_rules! assert_pop { - ($q:expr, $expect:expr) => { - assert_eq!( - match $q.pop_local_first() { - Some(v) => num(v), - None => panic!("queue empty"), - }, - $expect - ) - }; -} - -macro_rules! assert_pop_global { - ($q:expr, $expect:expr) => { - assert_eq!( - match $q.pop_global_first() { - Some(v) => num(v), - None => panic!("queue empty"), - }, - $expect - ) - }; -} - -macro_rules! assert_steal { - ($q:expr, $n:expr, $expect:expr) => { - assert_eq!( - match $q.steal($n) { - Some(v) => num(v), - None => panic!("queue empty"), - }, - $expect - ) - }; -} - -macro_rules! assert_empty { - ($q:expr) => {{ - let q: &mut queue::Worker = &mut $q; - if let Some(v) = q.pop_local_first() { - panic!("expected emtpy queue; got {}", num(v)); - } - }}; -} - -#[test] -fn single_worker_push_pop() { - let mut q = queue::build(1).remove(0); - - // Queue is empty - assert_empty!(q); - - // Push a value - q.push(val(0)); - - // Pop the value - assert_pop!(q, 0); - - // Push two values - q.push(val(1)); - q.push(val(2)); - q.push(val(3)); - - // Pop the value - assert_pop!(q, 3); - assert_pop!(q, 1); - assert_pop!(q, 2); - assert_empty!(q); -} - -#[test] -fn multi_worker_push_pop() { - let (mut q1, mut q2) = queues_2(); - - // Queue is empty - assert_empty!(q1); - assert_empty!(q2); - - // Push a value - q1.push(val(0)); - - // Not available on other queue - assert_empty!(q2); - assert_pop!(q1, 0); - - q2.push(val(1)); - assert_pop!(q2, 1); - assert_empty!(q1); -} - -#[test] -fn multi_worker_inject_pop() { - let (mut q1, mut q2) = queues_2(); - let i = q1.injector(); - - // Push a value - i.push(val(0), is_ok); - assert_pop!(q1, 0); - assert_empty!(q2); - - // Push another value - i.push(val(1), is_ok); - assert_pop!(q2, 1); - assert_empty!(q1); - - i.push(val(2), is_ok); - i.push(val(3), is_ok); - i.push(val(4), is_ok); - assert_pop!(q2, 2); - assert_pop!(q1, 3); - assert_pop!(q1, 4); -} - -#[test] -fn overflow_local_queue() { - let (mut q1, mut q2) = queues_2(); - - for i in 0..LOCAL_QUEUE_CAPACITY { - q1.push(val(i as u32)); - } - - assert_empty!(q2); - - // Fill `next` slot - q1.push(val(999)); - - // overflow - q1.push(val(1000)); - - assert_pop!(q2, 0); - assert_pop!(q1, 1000); - - // Half the values were moved to the global queue - for i in 128..LOCAL_QUEUE_CAPACITY { - assert_pop!(q1, i as u32); - } - - for i in 1..128 { - assert_pop!(q2, i); - } - - assert_pop!(q2, 999); - assert_empty!(q2); - - assert_empty!(q1); -} - -#[test] -fn polling_global_first() { - let (q, _) = queues_2(); - let i = q.injector(); - - i.push(val(1000), is_ok); - i.push(val(1001), is_ok); - - for n in 0..5 { - q.push(val(n)); - } - - assert_pop_global!(q, 1000); - assert_pop!(q, 4); - assert_pop_global!(q, 1001); - assert_pop_global!(q, 0); - assert_pop!(q, 1); - assert_pop_global!(q, 2); - assert_pop_global!(q, 3); - - assert!(q.pop_global_first().is_none()); -} - -#[test] -fn steal() { - let mut qs = queue::build(3); - let (mut q1, mut q2, mut q3) = (qs.remove(0), qs.remove(0), qs.remove(0)); - - assert!(q1.steal(0).is_none()); - assert!(q2.steal(0).is_none()); - assert!(q3.steal(0).is_none()); - - // Steal one value, but not the first one - q1.push(val(0)); - q1.push(val(999)); - assert_steal!(q2, 0, 0); - assert!(q2.steal(0).is_none()); - assert_pop!(q1, 999); - - // Steals half the queue - for i in 0..4 { - q1.push(val(i)); - } - - q1.push(val(999)); - - assert_steal!(q2, 0, 1); - assert_pop!(q2, 0); - assert_empty!(q2); - assert_pop!(q1, 999); - assert_pop!(q1, 2); - assert_pop!(q1, 3); - assert_empty!(q1); - - // Searches multiple queues - q3.push(val(0)); - q3.push(val(999)); - assert_steal!(q2, 0, 0); - assert_pop!(q3, 999); - assert_empty!(q3); - - // Steals from one queue at a time - q1.push(val(0)); - q1.push(val(998)); - q2.push(val(1)); - q2.push(val(999)); - - assert_steal!(q3, 0, 0); - assert_pop!(q2, 999); - assert_pop!(q2, 1); - assert_empty!(q2); - - assert_pop!(q1, 998); - assert_empty!(q1); -} - -fn queues_2() -> (queue::Worker, queue::Worker) { - let mut qs = queue::build(2); - (qs.remove(0), qs.remove(0)) -} - -// pretty big hack to track tasks -use std::cell::RefCell; -use std::collections::HashMap; -thread_local! { - static TASKS: RefCell>> = RefCell::new(HashMap::new()) -} - -fn val(num: u32) -> Task { - let (task, join) = task::joinable(async move { num }); - let prev = TASKS.with(|t| t.borrow_mut().insert(num, join)); - assert!(prev.is_none()); - task -} - -fn num(task: Task) -> u32 { - use futures::task::noop_waker_ref; - use std::future::Future; - use std::pin::Pin; - use std::task::Context; - use std::task::Poll::*; - - assert!(task.run(&mut || Some(From::from(&NOOP_SCHEDULE))).is_none()); - - // Find the task that completed - TASKS.with(|c| { - let mut map = c.borrow_mut(); - let mut num = None; - - for (_, join) in map.iter_mut() { - let mut cx = Context::from_waker(noop_waker_ref()); - if let Ready(n) = Pin::new(join).poll(&mut cx) { - num = Some(n.unwrap()); - break; - } - } - - let num = num.expect("no task completed"); - map.remove(&num); - num - }) -} - -fn is_ok(r: Result) { - assert!(r.is_ok()) -} diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs index a4e9d83be..cc79530a9 100644 --- a/tokio/src/runtime/thread_pool/worker.rs +++ b/tokio/src/runtime/thread_pool/worker.rs @@ -1,19 +1,166 @@ -use crate::loom::cell::CausalCell; -use crate::loom::sync::Arc; -use crate::park::Park; -use crate::runtime; -use crate::runtime::park::Parker; -use crate::runtime::thread_pool::{current, slice, Owned, Shared}; -use crate::task::Task; +//! A scheduler is initialized with a fixed number of workers. Each worker is +//! driven by a thread. Each worker has a "core" which contains data such as the +//! run queue and other state. When `block_in_place` is called, the worker's +//! "core" is handed off to a new thread allowing the scheduler to continue to +//! make progress while the originating thread blocks. -use std::cell::Cell; -use std::marker::PhantomData; -use std::sync::atomic::Ordering::Relaxed; +use crate::loom::rand::seed; +use crate::loom::sync::{Arc, Mutex}; +use crate::park::{Park, Unpark}; +use crate::runtime; +use crate::runtime::park::{Parker, Unparker}; +use crate::runtime::task; +use crate::runtime::thread_pool::{queue, AtomicCell, Idle}; +use crate::util::linked_list::LinkedList; +use crate::util::FastRand; + +use std::cell::RefCell; use std::time::Duration; -thread_local! { - /// Used to handle block_in_place - static ON_BLOCK: Cell> = Cell::new(None) +/// A scheduler worker +pub(super) struct Worker { + /// Reference to shared state + shared: Arc, + + /// Index holding this worker's remote state + index: usize, + + /// Used to hand-off a worker's core to another thread. + core: AtomicCell, +} + +/// Core data +struct Core { + /// Used to schedule bookkeeping tasks every so often. + tick: u8, + + /// The worker-local run queue. + run_queue: queue::Local>, + + /// True if the worker is currently searching for more work. Searching + /// involves attempting to steal from other workers. + is_searching: bool, + + /// True if the scheduler is being shutdown + is_shutdown: bool, + + /// Tasks owned by the core + tasks: LinkedList, + + /// Parker + /// + /// Stored in an `Option` as the parker is added / removed to make the + /// borrow checker happy. + park: Option, + + /// Fast random number generator. + rand: FastRand, +} + +/// State shared across all workers +pub(super) struct Shared { + /// Per-worker remote state. All other workers have access to this and is + /// how they communicate between each other. + remotes: Box<[Remote]>, + + /// Submit work to the scheduler while **not** currently on a worker thread. + inject: queue::Inject>, + + /// Coordinates idle workers + idle: Idle, + + /// Workers have have observed the shutdown signal + /// + /// The core is **not** placed back in the worker to avoid it from being + /// stolen by a thread that was spawned as part of `block_in_place`. + shutdown_workers: Mutex, Arc)>>, +} + +/// Used to communicate with a worker from other threads. +struct Remote { + /// Steal tasks from this worker. + steal: queue::Steal>, + + /// Transfers tasks to be released. Any worker pushes tasks, only the owning + /// worker pops. + pending_drop: task::TransferStack>, + + /// Unparks the associated worker thread + unpark: Unparker, +} + +/// Thread-local context +struct Context { + /// Worker + worker: Arc, + + /// Core data + core: RefCell>>, +} + +/// Starts the workers +pub(crate) struct Launch(Vec>); + +/// Running a task may consume the core. If the core is still available when +/// running the task completes, it is returned. Otherwise, the worker will need +/// to stop processing. +type RunResult = Result, ()>; + +/// A task handle +type Task = task::Task>; + +/// A notified task handle +type Notified = task::Notified>; + +// Tracks thread-local state +scoped_thread_local!(static CURRENT: Context); + +pub(super) fn create(size: usize, park: Parker) -> (Arc, Launch) { + let mut cores = vec![]; + let mut remotes = vec![]; + + // Create the local queues + for _ in 0..size { + let (steal, run_queue) = queue::local(); + + let park = park.clone(); + let unpark = park.unpark(); + + cores.push(Box::new(Core { + tick: 0, + run_queue, + is_searching: false, + is_shutdown: false, + tasks: LinkedList::new(), + park: Some(park), + rand: FastRand::new(seed()), + })); + + remotes.push(Remote { + steal, + pending_drop: task::TransferStack::new(), + unpark, + }); + } + + let shared = Arc::new(Shared { + remotes: remotes.into_boxed_slice(), + inject: queue::Inject::new(), + idle: Idle::new(size), + shutdown_workers: Mutex::new(vec![]), + }); + + let mut launch = Launch(vec![]); + + for (index, core) in cores.drain(..).enumerate() { + launch.0.push(Arc::new(Worker { + shared: shared.clone(), + index, + core: AtomicCell::new(Some(core)), + })); + } + + (shared, launch) } cfg_blocking! { @@ -21,597 +168,543 @@ cfg_blocking! { where F: FnOnce() -> R, { - // Make the current worker give away its Worker to another thread so that we can safely block - // this one without preventing progress on other futures the worker owns. - ON_BLOCK.with(|ob| { - let allow_blocking = ob - .get() - // `block_in_place` can only be called from a spawned task when - // working with the threaded scheduler. - .expect("can call blocking only when running in a spawned task"); + // Try to steal the worker core back + struct Reset; - // This is safe, because ON_BLOCK was set from an &mut dyn FnMut in the worker that wraps - // the worker's operation, and is unset just prior to when the FnMut is dropped. - let allow_blocking = unsafe { &*allow_blocking }; + impl Drop for Reset { + fn drop(&mut self) { + CURRENT.with(|maybe_cx| { + if let Some(cx) = maybe_cx { + let core = cx.worker.core.take(); + *cx.core.borrow_mut() = core; + } + }); + } + } - allow_blocking(); - f() - }) + CURRENT.with(|maybe_cx| { + let cx = maybe_cx.expect("can call blocking only when running in a spawned task"); + + // Get the worker core. If none is set, then blocking is fine! + let core = match cx.core.borrow_mut().take() { + Some(core) => core, + None => return, + }; + + // The parker should be set here + assert!(core.park.is_some()); + + // In order to block, the core must be sent to another thread for + // execution. + // + // First, move the core back into the worker's shared core slot. + cx.worker.core.set(core); + + // Next, clone the worker handle and send it to a new thread for + // processing. + // + // Once the blocking task is done executing, we will attempt to + // steal the core back. + let worker = cx.worker.clone(); + runtime::spawn_blocking(move || run(worker)); + }); + + let _reset = Reset; + + f() } } -pub(crate) struct Worker { - /// Parks the thread. Requires the calling worker to have obtained unique - /// access via the generation synchronization action. - inner: Arc, - - /// Scheduler slices - slices: Arc, - - /// Slice assigned to this worker - index: usize, - - /// Worker generation. This is used to synchronize access to the internal - /// data. - generation: usize, - - /// To indicate that the Worker has been given away and should no longer be used - gone: Cell, -} - -/// Internal worker state. This may be referenced from multiple threads, but the -/// generation guard protects unsafe access -struct Inner { - /// Used to park the thread - park: CausalCell, -} - -unsafe impl Send for Worker {} - -/// Used to ensure the invariants are respected -struct GenerationGuard<'a> { - /// Worker reference - worker: &'a Worker, - - /// Prevent `Sync` access - _p: PhantomData>, -} - -struct WorkerGone; - -// TODO: Move into slices -pub(super) fn create_set(pool_size: usize, parker: Parker) -> (Arc, Vec) { - // Create the parks... - let parkers: Vec<_> = (0..pool_size).map(|_| parker.clone()).collect(); - - let mut slices = Arc::new(slice::Set::new(&parkers)); - - // Establish the circular link between the individual worker state - // structure and the container. - Arc::get_mut(&mut slices).unwrap().set_ptr(); - - // This will contain each worker. - let workers = parkers - .into_iter() - .enumerate() - .map(|(index, parker)| Worker::new(slices.clone(), index, parker)) - .collect(); - - (slices, workers) -} - /// After how many ticks is the global queue polled. This helps to ensure /// fairness. /// /// The number is fairly arbitrary. I believe this value was copied from golang. -const GLOBAL_POLL_INTERVAL: u16 = 61; +const GLOBAL_POLL_INTERVAL: u8 = 61; -impl Worker { - // Safe as aquiring a lock is required before doing anything potentially - // dangerous. - pub(super) fn new(slices: Arc, index: usize, park: Parker) -> Self { - Worker { - inner: Arc::new(Inner { - park: CausalCell::new(park), - }), - slices, - index, - generation: 0, - gone: Cell::new(false), - } - } - - pub(super) fn run(self) { - // First, acquire a lock on the worker. - let guard = match self.acquire_lock() { - Some(guard) => guard, - None => return, - }; - - // Track the current worker - current::set(&self.slices, self.index, || { - // Enter a runtime context - let _enter = crate::runtime::enter(); - - ON_BLOCK.with(|ob| { - // Ensure that the ON_BLOCK 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(ob); - - let allow_blocking: &dyn Fn() = &|| self.block_in_place(); - - ob.set(Some(unsafe { - // NOTE: We cannot use a safe cast to raw pointer here, since we are - // _also_ erasing the lifetime of these pointers. That is safe here, - // because we know that ob will set back to None before allow_blocking - // is dropped. - #[allow(clippy::useless_transmute)] - std::mem::transmute::<_, *const dyn Fn()>(allow_blocking) - })); - - let _ = guard.run(); - - // Ensure that we reset ob before allow_blocking is dropped. - drop(_reset); - }); - }); - - if self.gone.get() { - // Synchronize with the pool for load(Acquire) in is_closed to get - // up-to-date value. - self.slices.wait_for_unlocked(); - - if self.slices.is_closed() { - // If the pool is shutting down, some other thread may be - // waiting to clean up after the task that we were holding on - // to. If we completed that task, we did nothing (because - // task.run() returned None), and so crucially we did not wait - // up any such thread. - // - // So, we have to do that here. - self.slices.notify_all(); - } - } - } - - /// Acquires the lock - fn acquire_lock(&self) -> Option> { - // Safety: Only getting `&self` access to access atomic field - let owned = unsafe { &*self.slices.owned()[self.index].get() }; - - // The lock is only to establish mutual exclusion. Other synchronization - // handles memory orderings - let prev = owned.generation.compare_and_swap( - self.generation, - self.generation.wrapping_add(1), - Relaxed, - ); - - if prev == self.generation { - Some(GenerationGuard { - worker: self, - _p: PhantomData, - }) - } else { - None - } - } - - /// Enters an in-place blocking section - fn block_in_place(&self) { - // If our Worker has already been given away, then blocking is fine! - if self.gone.get() { - return; - } - - // If this method is called, we need to move the entire worker onto a - // separate (blocking) thread before returning. Once we return, the - // caller is going to execute some blocking code which would otherwise - // block our reactor from making progress. Since we are _in the middle_ - // of running a task, this isn't trivial, as the Worker is "active". - // We do have the luxury of knowing that we are on the worker thread, - // so we can assert exclusive access to any Worker-specific state. - // - // More specifically, the caller is _currently_ "stuck" in - // Entry::run_task at: - // - // if let Some(task) = task.run(self.shared().into()) { - // - // And _we_ get to decide when it continues (specifically, by choosing - // when we return from the second callback (i.e., after the FnOnce - // passed to blocking has returned). - // - // Here's what we'll have to do: - // - // - Reconstruct our `Worker` struct - // - Spawn the reconstructed `Worker` on another blocking thread - // - Clear any state indicating what worker we are on, since at this - // point we are effectively no longer "on" that worker. - // - Allow the caller of `blocking` to continue. - // - // Once the caller completes the blocking operations, we need to ensure - // that async code can continue running in that context. Luckily, since - // `Arc` has a fallback for when - // current::get() is None, we can just let the task - // run until it yields, and then put it back into - // the pool. - - let worker = Worker { - inner: self.inner.clone(), - slices: self.slices.clone(), - index: self.index, - generation: self.generation + 1, - gone: Cell::new(false), - }; - - // Give away the worker - // - // Returns `Err` if the spawn failed due to the runtime shutting down - let res = runtime::try_spawn_blocking(move || worker.run()); - - // If the worker hand-off was successful, clear the local state. - // Otherwise, the runtime is in the process of shutting down, so we will - // just block on the worker. - if res.is_ok() { - // make sure no subsequent code thinks that it is on a worker - current::clear(); - - // Track that the worker is gone - self.gone.set(true); +impl Launch { + pub(crate) fn launch(mut self) { + for worker in self.0.drain(..) { + runtime::spawn_blocking(move || run(worker)); } } } -impl GenerationGuard<'_> { - fn run(self) -> Result<(), WorkerGone> { - let mut me = self; +fn run(worker: Arc) { + // Acquire a core. If this fails, then another thread is running this + // worker and there is nothing further to do. + let core = match worker.core.take() { + Some(core) => core, + None => return, + }; - while me.is_running() { - me = me.process_available_work()?; + // Set the worker context. + let cx = Context { + worker, + core: RefCell::new(None), + }; - if me.is_running() { - me.park(); - } - } + let _enter = crate::runtime::enter(); - me.shutdown(); - Ok(()) - } + CURRENT.set(&cx, || { + // This should always be an error. It only returns a `Result` to support + // using `?` to short circuit. + assert!(cx.run(core).is_err()); + }); +} - fn is_running(&self) -> bool { - self.owned().is_running.get() - } +impl Context { + fn run(&self, mut core: Box) -> RunResult { + while !core.is_shutdown { + // Increment the tick + core.tick(); - /// Returns `true` if the worker needs to park - fn process_available_work(self) -> Result { - let mut me = self; + // Run maintenance, if needed + core = self.maintenance(core); - loop { - // Local queue loop - loop { - let task = match me.find_local_work() { - Some(task) => task, - None => { - if !me.is_running() { - // The scheduler is in the process of shutting down. - return Ok(me); - } - - // Break out of the local task loop and try to steal - break; - } - }; - - me = me.run_task(task)?; + // First, check work available to the current worker. + if let Some(task) = core.next_task(&self.worker) { + core = self.run_task(task, core)?; + continue; } - // No more **local** work to process, try transitioning to searching - // in order to attempt to steal work from other workers. - // - // On `false`, the worker has entered the parked state - if me.transition_to_searching() { - // Try to steal tasks from other workers - if let Some(task) = me.steal_work() { - me = me.run_task(task)?; - } else { - // No work to steal, perform some routine work - me.drain_tasks_pending_drop(); - return Ok(me); - } + // There is no more **local** work to process, try to steal work + // from other workers. + if let Some(task) = core.steal_work(&self.worker) { + core = self.run_task(task, core)?; } else { - return Ok(me); + // Wait for work + core = self.park(core); } + } - // Start checking the local queue again + // Signal shutdown + self.worker.shared.shutdown(core, self.worker.clone()); + Err(()) + } + + fn run_task(&self, task: Notified, mut core: Box) -> RunResult { + // Make sure thew orker is not in the **searching** state. This enables + // another idle worker to try to steal work. + core.transition_from_searching(&self.worker); + + // Make the core available to the runtime context + *self.core.borrow_mut() = Some(core); + + // Run the task + task.run(); + + // Try to take the core back + match self.core.borrow_mut().take() { + Some(core) => Ok(core), + None => Err(()), } } - /// Finds local work - fn find_local_work(&mut self) -> Option> { - let tick = self.tick_fetch_inc(); + fn maintenance(&self, mut core: Box) -> Box { + if core.tick % GLOBAL_POLL_INTERVAL == 0 { + // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... + // to run without actually putting the thread to sleep. + core = self.park_timeout(core, Some(Duration::from_millis(0))); - if tick % GLOBAL_POLL_INTERVAL == 0 { - // Sleep light... - self.park_light(); + // Run regularly scheduled maintenance + core.maintenance(&self.worker); + } - // Perform regularly scheduled maintenance work. - self.maintenance(); + core + } - if !self.is_running() { - return None; + fn park(&self, mut core: Box) -> Box { + core.transition_to_parked(&self.worker); + + while !core.is_shutdown { + core = self.park_timeout(core, None); + + // Run regularly scheduled maintenance + core.maintenance(&self.worker); + + if core.transition_from_parked(&self.worker) { + return core; } + } - // Check the global queue - self.owned().work_queue.pop_global_first() + core + } + + fn park_timeout(&self, mut core: Box, duration: Option) -> Box { + // Take the parker out of core + let mut park = core.park.take().expect("park missing"); + + // Store `core` in context + *self.core.borrow_mut() = Some(core); + + // Park thread + if let Some(timeout) = duration { + park.park_timeout(timeout).expect("park failed"); } else { - self.owned().work_queue.pop_local_first() + park.park().expect("park failed"); + } + + // Remove `core` from context + core = self.core.borrow_mut().take().expect("core missing"); + + // Place `park` back in `core` + core.park = Some(park); + + // If there are tasks available to steal, notify a worker + if core.run_queue.is_stealable() { + self.worker.shared.notify_parked(); + } + + core + } +} + +impl Core { + /// Increment the tick + fn tick(&mut self) { + self.tick = self.tick.wrapping_add(1); + } + + /// Return the next notified task available to this worker. + fn next_task(&mut self, worker: &Worker) -> Option { + if self.tick % GLOBAL_POLL_INTERVAL == 0 { + worker.inject().pop().or_else(|| self.run_queue.pop()) + } else { + self.run_queue.pop().or_else(|| worker.inject().pop()) } } - fn steal_work(&mut self) -> Option> { - let num_slices = self.worker.slices.len(); - let start = self.owned().rand.fastrand_n(num_slices as u32); + fn steal_work(&mut self, worker: &Worker) -> Option { + if !self.transition_to_searching(worker) { + return None; + } - self.owned() - .work_queue - .steal(start as usize) - // Fallback on checking the local queue, which will also check the - // injector. - .or_else(|| self.owned().work_queue.pop_global_first()) + let num = worker.shared.remotes.len(); + // Start from a random worker + let start = self.rand.fastrand_n(num as u32) as usize; + + for i in 0..num { + let i = (start + i) % num; + + // Don't steal from ourself! We know we don't have work. + if i == worker.index { + continue; + } + + let target = &worker.shared.remotes[i]; + if let Some(task) = target.steal.steal_into(&mut self.run_queue) { + return Some(task); + } + } + + // Fallback on checking the global queue + worker.shared.inject.pop() + } + + fn transition_to_searching(&mut self, worker: &Worker) -> bool { + if !self.is_searching { + self.is_searching = worker.shared.idle.transition_worker_to_searching(); + } + + self.is_searching + } + + fn transition_from_searching(&mut self, worker: &Worker) { + if !self.is_searching { + return; + } + + self.is_searching = false; + worker.shared.transition_worker_from_searching(); + } + + /// Prepare the worker state for parking + fn transition_to_parked(&mut self, worker: &Worker) { + // When the final worker transitions **out** of searching to parked, it + // must check all the queues one last time in case work materialized + // between the last work scan and transitioning out of searching. + let is_last_searcher = worker + .shared + .idle + .transition_worker_to_parked(worker.index, self.is_searching); + + // The worker is no longer searching. Setting this is the local cache + // only. + self.is_searching = false; + + if is_last_searcher { + worker.shared.notify_if_work_pending(); + } + } + + /// Returns `true` if the transition happened. + fn transition_from_parked(&mut self, worker: &Worker) -> bool { + // If there is a non-stealable task, then we must unpark regardless of + // being notified + if self.run_queue.has_unstealable() { + worker.shared.idle.unpark_worker_by_id(worker.index); + self.is_searching = true; + return true; + } + + if worker.shared.idle.is_parked(worker.index) { + return false; + } + + // When unparked, the worker is in the searching state. + self.is_searching = true; + true } /// Runs maintenance work such as free pending tasks and check the pool's /// state. - fn maintenance(&mut self) { - // Free any completed tasks - self.drain_tasks_pending_drop(); + fn maintenance(&mut self, worker: &Worker) { + self.drain_pending_drop(worker); - // Update the pool state cache - let closed = self.owned().work_queue.is_closed(); - self.owned().is_running.set(!closed) - } - - fn transition_to_searching(&mut self) -> bool { - if self.is_searching() { - return true; - } - - let ret = self.slices().idle().transition_worker_to_searching(); - self.owned().is_searching.set(ret); - ret - } - - fn transition_from_searching(&mut self) { - self.owned().is_searching.set(false); - - if self.slices().idle().transition_worker_from_searching() { - // We are the final searching worker. Because work was found, we - // need to notify another worker. - self.slices().notify_work(); + if !self.is_shutdown { + // Check if the scheduler has been shutdown + self.is_shutdown = worker.inject().is_closed(); } } - /// Returns `true` if the worker must check for any work. - fn transition_to_parked(&mut self) -> bool { - let idx = self.index(); - let is_searching = self.is_searching(); - let ret = self - .slices() - .idle() - .transition_worker_to_parked(idx, is_searching); + // Shutdown the core + fn shutdown(&mut self, worker: &Worker) { + // Take the core + let mut park = self.park.take().expect("park missing"); - // The worker is no longer searching. Setting this is the local cache - // only. - self.owned().is_searching.set(false); - - // When tasks are submitted locally (from the parker), defer any - // notifications in hopes that the curent worker will grab those tasks. - self.owned().defer_notification.set(true); - - ret - } - - /// Returns `true` if the transition happened. - fn transition_from_parked(&mut self) -> bool { - if self.owned().did_submit_task.get() || !self.is_running() { - // Remove the worker from the sleep set. - self.slices().idle().unpark_worker_by_id(self.index()); - - self.owned().is_searching.set(true); - self.owned().defer_notification.set(false); - - true - } else { - let ret = !self.slices().idle().is_parked(self.index()); - - if ret { - self.owned().is_searching.set(true); - self.owned().defer_notification.set(false); - } - - ret - } - } - - /// Runs the task. During the task execution, it is possible for worker to - /// transition to a new thread. In this case, the caller loses the guard to - /// access the generation and must stop processing. - fn run_task(mut self, task: Task) -> Result { - if self.is_searching() { - self.transition_from_searching(); + // Signal to all tasks to shut down. + for header in self.tasks.iter() { + header.shutdown(); } - let gone = &self.worker.gone; - let executor = self.shared(); - - let task = task.run(&mut || { - if gone.get() { - None - } else { - Some(executor.into()) - } - }); - - if gone.get() { - // The Worker disappeared from under us. - // We need to return, because we no longer own all of our state! - // Make sure the task gets picked up again eventually. - if let Some(task) = task { - self.worker.slices.schedule(task); - } - - Err(WorkerGone) - } else { - if let Some(task) = task { - self.owned().submit_local_yield(task); - self.slices().notify_work(); - } - - Ok(self) - } - } - - fn final_work_sweep(&mut self) { - if !self.owned().work_queue.is_empty() { - self.slices().notify_work(); - } - } - - fn park(&mut self) { - if self.transition_to_parked() { - // We are the final searching worker, check if any work arrived - // before parking - self.final_work_sweep(); - } - - // The state has been transitioned to parked, we can now wait by - // calling the parker. This is done in a loop as spurious wakeups are - // permitted. loop { - self.park_mut().park().expect("park failed"); + self.drain_pending_drop(worker); - // We might have been woken to clean up a dropped task - self.maintenance(); - - if self.transition_from_parked() { - return; - } - } - } - - fn park_light(&mut self) { - // When tasks are submitted locally (from the parker), defer any - // notifications in hopes that the curent worker will grab those tasks. - self.owned().defer_notification.set(true); - - self.park_mut() - .park_timeout(Duration::from_millis(0)) - .expect("park failed"); - - self.owned().defer_notification.set(false); - - if self.owned().did_submit_task.get() { - self.slices().notify_work(); - self.owned().did_submit_task.set(false) - } - } - - fn drain_tasks_pending_drop(&mut self) { - for task in self.shared().pending_drop.drain() { - unsafe { - let owned = &mut *self.slices().owned()[self.index()].get(); - owned.release_task(&task); - } - drop(task); - } - } - - /// Shutdowns the worker. - /// - /// Once the shutdown flag has been observed, it is guaranteed that no - /// further tasks may be pushed into the global queue. - fn shutdown(&mut self) { - // Transition all tasks owned by the worker to canceled. - self.owned().owned_tasks.shutdown(); - - // Always notify the first time around. This flushes any released tasks - // that happened before the call to `Worker::shutdown` - let mut notify = true; - - // The worker can only shutdown once there are no further owned tasks. - loop { - // First, drain all tasks from both the local & global queue. - while let Some(task) = self.owned().work_queue.pop_local_first() { - notify = true; - task.shutdown(); - } - - if notify { - // If any tasks are shutdown, they may be pushed on another - // worker's `pending_drop` stack. However, we don't know which - // workers need to be notified, so we just notify all of them. - // Since this is a shutdown process, excessive notification is - // not a huge deal. - self.worker.slices.notify_all(); - notify = false; - } - - // Try draining more tasks - self.drain_tasks_pending_drop(); - - if self.owned().owned_tasks.is_empty() { + if self.tasks.is_empty() { break; } - // Wait until task that this worker owns are released. + // Wait until signalled + park.park().expect("park failed"); + } + + // Drain the queue + while let Some(_) = self.run_queue.pop() {} + } + + fn drain_pending_drop(&mut self, worker: &Worker) { + use std::mem::ManuallyDrop; + + for task in worker.remote().pending_drop.drain() { + let task = ManuallyDrop::new(task); + + // safety: tasks are only pushed into the `pending_drop` stacks that + // are associated with the list they are inserted into. When a task + // is pushed into `pending_drop`, the ref-inc is skipped, so we must + // not ref-dec here. // - // `transition_to_parked` is not called as we are not working - // anymore. When a task is released, the owning worker is unparked - // directly. - self.park_mut().park().expect("park failed"); + // See `bind` and `release` implementations. + unsafe { + self.tasks.remove(task.header().into()); + } + } + } +} + +impl Worker { + /// Returns a reference to the scheduler's injection queue + fn inject(&self) -> &queue::Inject> { + &self.shared.inject + } + + /// Return a reference to this worker's remote data + fn remote(&self) -> &Remote { + &self.shared.remotes[self.index] + } + + fn eq(&self, other: &Worker) -> bool { + self.shared.ptr_eq(&other.shared) && self.index == other.index + } +} + +impl task::Schedule for Arc { + fn bind(task: Task) -> Arc { + CURRENT.with(|maybe_cx| { + let cx = maybe_cx.expect("scheduler context missing"); + + // Track the task + cx.core + .borrow_mut() + .as_mut() + .expect("scheduler core missing") + .tasks + .push_front(task); + + // Return a clone of the worker + cx.worker.clone() + }) + } + + fn release(&self, task: &Task) -> Option { + use std::ptr::NonNull; + + CURRENT.with(|maybe_cx| { + let cx = maybe_cx.expect("scheduler context missing"); + + if self.eq(&cx.worker) { + let mut maybe_core = cx.core.borrow_mut(); + + if let Some(core) = &mut *maybe_core { + // Directly remove the task + // + // safety: the task is inserted in the list in `bind`. + unsafe { + let ptr = NonNull::from(task.header()); + return core.tasks.remove(ptr); + } + } + } + + // Track the task to be released by the worker that owns it + // + // Safety: We get a new handle without incrementing the ref-count. + // A ref-count is held by the "owned" linked list and it is only + // ever removed from that list as part of the release process: this + // method or popping the task from `pending_drop`. Thus, we can rely + // on the ref-count held by the linked-list to keep the memory + // alive. + // + // When the task is removed from the stack, it is forgotten instead + // of dropped. + let task = unsafe { Task::from_raw(task.header().into()) }; + + self.remote().pending_drop.push(task); + + if cx.core.borrow().is_some() { + return None; + } + + // The worker core has been handed off to another thread. In the + // event that the scheduler is currently shutting down, the thread + // that owns the task may be waiting on the release to complete + // shutdown. + if self.inject().is_closed() { + self.remote().unpark.unpark(); + } + + None + }) + } + + fn schedule(&self, task: Notified) { + self.shared.schedule(task, false); + } + + fn yield_now(&self, task: Notified) { + self.shared.schedule(task, true); + } +} + +impl Shared { + pub(super) fn schedule(&self, task: Notified, is_yield: bool) { + CURRENT.with(|maybe_cx| { + if let Some(cx) = maybe_cx { + // Make sure the task is part of the **current** scheduler. + if self.ptr_eq(&cx.worker.shared) { + // And the current thread still holds a core + if let Some(core) = cx.core.borrow_mut().as_mut() { + self.schedule_local(core, task, is_yield); + return; + } + } + } + + // Otherwise, use the inject queue + self.inject.push(task); + self.notify_parked(); + }); + } + + fn schedule_local(&self, core: &mut Core, task: Notified, is_yield: bool) { + // Spawning from the worker thread. If scheduling a "yield" then the + // task must always be pushed to the back of the queue, enabling other + // tasks to be executed. If **not** a yield, then there is more + // flexibility and the task may go to the front of the queue. + let should_notify = if is_yield { + core.run_queue.push_back(task, &self.inject); + true + } else { + core.run_queue.push(task, &self.inject) + }; + + // Only notify if not currently parked. If `park` is `None`, then the + // scheduling is from a resource driver. As notifications often come in + // batches, the notification is delayed until the park is complete. + if should_notify && core.park.is_some() { + self.notify_parked(); } } - /// Increments the tick, returning the value from before the increment. - fn tick_fetch_inc(&mut self) -> u16 { - let tick = self.owned().tick.get(); - self.owned().tick.set(tick.wrapping_add(1)); - tick + pub(super) fn close(&self) { + if self.inject.close() { + self.notify_all(); + } } - fn is_searching(&self) -> bool { - self.owned().is_searching.get() + fn notify_parked(&self) { + if let Some(index) = self.idle.worker_to_notify() { + self.remotes[index].unpark.unpark(); + } } - fn index(&self) -> usize { - self.worker.index + fn notify_all(&self) { + for remote in &self.remotes[..] { + remote.unpark.unpark(); + } } - fn slices(&self) -> &slice::Set { - &self.worker.slices + fn notify_if_work_pending(&self) { + for remote in &self.remotes[..] { + if !remote.steal.is_empty() { + self.notify_parked(); + return; + } + } + + if !self.inject.is_empty() { + self.notify_parked(); + } } - fn shared(&self) -> &Shared { - &self.slices().shared()[self.index()] + fn transition_worker_from_searching(&self) { + if self.idle.transition_worker_from_searching() { + // We are the final searching worker. Because work was found, we + // need to notify another worker. + self.notify_parked(); + } } - fn owned(&self) -> &Owned { - let index = self.index(); - // safety: we own the slot - unsafe { &*self.slices().owned()[index].get() } + /// Signals that a worker has observed the shutdown signal and has replaced + /// its core back into its handle. + /// + /// If all workers have reached this point, the final cleanup is performed. + fn shutdown(&self, core: Box, worker: Arc) { + let mut workers = self.shutdown_workers.lock().unwrap(); + workers.push((core, worker)); + + if workers.len() != self.remotes.len() { + return; + } + + for (mut core, worker) in workers.drain(..) { + core.shutdown(&worker); + } + + // Drain the injection queue + while let Some(_) = self.inject.pop() {} } - fn park_mut(&mut self) -> &mut Parker { - // Safety: `&mut self` on `GenerationGuard` implies it is safe to - // perform the action. - unsafe { self.worker.inner.park.with_mut(|ptr| &mut *ptr) } + fn ptr_eq(&self, other: &Shared) -> bool { + self as *const _ == other as *const _ } } diff --git a/tokio/src/sync/notify.rs b/tokio/src/sync/notify.rs index ef34ad2d4..56b1c9a85 100644 --- a/tokio/src/sync/notify.rs +++ b/tokio/src/sync/notify.rs @@ -306,10 +306,10 @@ fn notify_locked(waiters: &mut LinkedList, state: &AtomicU8, curr: u8) - // transition **out** of `WAITING`. // // Get a pending waiter - let waiter = waiters.pop_back().unwrap(); + let mut waiter = waiters.pop_back().unwrap(); // Safety: `waiters` lock is still held. - let waiter = unsafe { &mut *waiter }; + let waiter = unsafe { waiter.as_mut() }; assert!(!waiter.notified); @@ -423,7 +423,9 @@ impl Future for Notified<'_> { } // Insert the waiter into the linked list - waiters.push_front(waiter.get()); + // + // safety: pointers from `UnsafeCell` are never null. + waiters.push_front(unsafe { NonNull::new_unchecked(waiter.get()) }); *state = Waiting; } @@ -535,16 +537,15 @@ impl Drop for Notified<'_> { /// /// `Waiter` is forced to be !Unpin. unsafe impl linked_list::Link for Waiter { - type Handle = *mut Waiter; + type Handle = NonNull; type Target = Waiter; - fn to_raw(handle: *mut Waiter) -> NonNull { - debug_assert!(!handle.is_null()); - unsafe { NonNull::new_unchecked(handle) } + fn as_raw(handle: &NonNull) -> NonNull { + *handle } - unsafe fn from_raw(ptr: NonNull) -> *mut Waiter { - ptr.as_ptr() + unsafe fn from_raw(ptr: NonNull) -> NonNull { + ptr } unsafe fn pointers(mut target: NonNull) -> NonNull> { diff --git a/tokio/src/task/core.rs b/tokio/src/task/core.rs deleted file mode 100644 index b7c15a988..000000000 --- a/tokio/src/task/core.rs +++ /dev/null @@ -1,156 +0,0 @@ -use crate::loom::alloc::Track; -use crate::loom::cell::CausalCell; -use crate::task::raw::{self, Vtable}; -use crate::task::state::State; -use crate::task::waker::waker_ref; -use crate::task::Schedule; - -use std::cell::UnsafeCell; -use std::future::Future; -use std::mem::MaybeUninit; -use std::pin::Pin; -use std::ptr::{self, NonNull}; -use std::task::{Context, Poll, Waker}; - -/// The task cell. Contains the components of the task. -/// -/// It is critical for `Header` to be the first field as the task structure will -/// be referenced by both *mut Cell and *mut Header. -#[repr(C)] -pub(super) struct Cell { - /// Hot task state data - pub(super) header: Header, - - /// Either the future or output, depending on the execution stage. - pub(super) core: Core, - - /// Cold data - pub(super) trailer: Trailer, -} - -/// The core of the task. -/// -/// Holds the future or output, depending on the stage of execution. -pub(super) struct Core { - stage: Stage, -} - -/// Crate public as this is also needed by the pool. -#[repr(C)] -pub(crate) struct Header { - /// Task state - pub(super) state: State, - - /// Pointer to the executor owned by the task - pub(super) executor: CausalCell>>, - - /// Pointer to next task, used for misc task linked lists. - pub(crate) queue_next: UnsafeCell<*const Header>, - - /// Pointer to the next task in the ownership list. - pub(crate) owned_next: UnsafeCell>>, - - /// Pointer to the previous task in the ownership list. - pub(crate) owned_prev: UnsafeCell>>, - - /// Table of function pointers for executing actions on the task. - pub(super) vtable: &'static Vtable, - - /// Used by loom to track the causality of the future. Without loom, this is - /// unit. - pub(super) future_causality: CausalCell<()>, -} - -/// Cold data is stored after the future. -pub(super) struct Trailer { - /// Consumer task waiting on completion of this task. - pub(super) waker: CausalCell>>, -} - -/// Either the future or the output. -enum Stage { - Running(Track), - Finished(Track>), - Consumed, -} - -impl Cell { - /// Allocates a new task cell, containing the header, trailer, and core - /// structures. - pub(super) fn new(future: T, state: State) -> Box> - where - S: Schedule, - { - Box::new(Cell { - header: Header { - state, - executor: CausalCell::new(None), - queue_next: UnsafeCell::new(ptr::null()), - owned_next: UnsafeCell::new(None), - owned_prev: UnsafeCell::new(None), - vtable: raw::vtable::(), - future_causality: CausalCell::new(()), - }, - core: Core { - stage: Stage::Running(Track::new(future)), - }, - trailer: Trailer { - waker: CausalCell::new(MaybeUninit::new(None)), - }, - }) - } -} - -impl Core { - pub(super) fn transition_to_consumed(&mut self) { - self.stage = Stage::Consumed - } - - pub(super) fn poll(&mut self, header: &Header) -> Poll - where - S: Schedule, - { - let res = { - let future = match &mut self.stage { - Stage::Running(tracked) => tracked.get_mut(), - _ => unreachable!("unexpected stage"), - }; - - // The future is pinned within the task. The above state transition - // has ensured the safety of this action. - let future = unsafe { Pin::new_unchecked(future) }; - - // The waker passed into the `poll` function does not require a ref - // count increment. - let waker_ref = waker_ref::(header); - let mut cx = Context::from_waker(&*waker_ref); - - future.poll(&mut cx) - }; - - if res.is_ready() { - self.stage = Stage::Consumed; - } - - res - } - - pub(super) fn store_output(&mut self, output: super::Result) { - self.stage = Stage::Finished(Track::new(output)); - } - - pub(super) unsafe fn read_output(&mut self, dst: *mut Track>) { - use std::mem; - - dst.write(match mem::replace(&mut self.stage, Stage::Consumed) { - Stage::Finished(output) => output, - _ => unreachable!("unexpected state"), - }); - } -} - -impl Header { - pub(super) fn executor(&self) -> Option> { - unsafe { self.executor.with(|ptr| *ptr) } - } -} diff --git a/tokio/src/task/harness.rs b/tokio/src/task/harness.rs deleted file mode 100644 index 09fdbe4de..000000000 --- a/tokio/src/task/harness.rs +++ /dev/null @@ -1,558 +0,0 @@ -use crate::loom::alloc::Track; -use crate::task::core::{Cell, Core, Header, Trailer}; -use crate::task::state::Snapshot; -use crate::task::{JoinError, Schedule, Task}; - -use std::future::Future; -use std::marker::PhantomData; -use std::mem::{ManuallyDrop, MaybeUninit}; -use std::ptr::NonNull; -use std::task::{Poll, Waker}; - -/// Typed raw task handle -pub(super) struct Harness { - cell: NonNull>, - _p: PhantomData, -} - -impl Harness -where - T: Future, - S: 'static, -{ - pub(super) unsafe fn from_raw(ptr: *mut ()) -> Harness { - debug_assert!(!ptr.is_null()); - - Harness { - cell: NonNull::new_unchecked(ptr as *mut Cell), - _p: PhantomData, - } - } - - fn header(&self) -> &Header { - unsafe { &self.cell.as_ref().header } - } - - fn trailer(&self) -> &Trailer { - unsafe { &self.cell.as_ref().trailer } - } - - fn core(&mut self) -> &mut Core { - unsafe { &mut self.cell.as_mut().core } - } -} - -impl Harness -where - T: Future, - S: Schedule, -{ - /// Polls the inner future. - /// - /// All necessary state checks and transitions are performed. - /// - /// Panics raised while polling the future are handled. - /// - /// Returns `true` if the task needs to be scheduled again - /// - /// # Safety - /// - /// The pointer returned by the `executor` fn must be castable to `*mut S` - pub(super) unsafe fn poll(mut self, executor: &mut dyn FnMut() -> Option>) -> bool { - use std::panic; - - // Transition the task to the running state. - let res = self.header().state.transition_to_running(); - - if res.is_canceled() { - // The task was concurrently canceled. - self.do_cancel(res); - return false; - } - - let join_interest = res.is_join_interested(); - debug_assert!(join_interest || !res.has_join_waker()); - - // Get the cell components - let cell = &mut self.cell.as_mut(); - let header = &cell.header; - let core = &mut cell.core; - - // If the task's executor pointer is not yet set, then set it here. This - // is safe because a) this is the only time the value is set. b) at this - // point, there are no outstanding wakers which might access the - // field concurrently. - if header.executor().is_none() { - // We don't want the destructor to run because we don't really - // own the task here. - let task = ManuallyDrop::new(Task::from_raw(header.into())); - // Call the scheduler's bind callback - let executor = executor().expect("first poll must happen from an executor"); - executor.cast::().as_ref().bind(&task); - header.executor.with_mut(|ptr| *ptr = Some(executor.cast())); - } - - // The transition to `Running` done above ensures that a lock on the - // future has been obtained. This also ensures the `*mut T` pointer - // contains the future (as opposed to the output) and is initialized. - - let res = header.future_causality.with_mut(|_| { - panic::catch_unwind(panic::AssertUnwindSafe(|| { - struct Guard<'a, T: Future> { - core: &'a mut Core, - polled: bool, - } - - impl Drop for Guard<'_, T> { - fn drop(&mut self) { - if !self.polled { - self.core.transition_to_consumed(); - } - } - } - - let mut guard = Guard { - core, - polled: false, - }; - - let res = guard.core.poll::(header); - - // prevent the guard from dropping the future - guard.polled = true; - - res - })) - }); - - match res { - Ok(Poll::Ready(out)) => { - self.complete(executor, join_interest, Ok(out)); - false - } - Ok(Poll::Pending) => { - let res = self.header().state.transition_to_idle(); - - if res.is_canceled() { - self.do_cancel(res); - false - } else { - res.is_notified() - } - } - Err(err) => { - self.complete(executor, join_interest, Err(JoinError::panic2(err))); - false - } - } - } - - pub(super) unsafe fn drop_task(mut self) { - let might_drop_join_waker_on_release = self.might_drop_join_waker_on_release(); - - let join_waker = if might_drop_join_waker_on_release { - // Read the join waker cell just to have it - self.read_join_waker() - } else { - MaybeUninit::uninit() - }; - - // transition the task to released - let res = self.header().state.release_task(); - - assert!(res.is_terminal(), "state = {:?}", res); - - if might_drop_join_waker_on_release && !res.is_join_interested() { - debug_assert!(res.has_join_waker()); - - // Its our responsibility to drop the waker - let _ = join_waker.assume_init(); - } - - if res.is_final_ref() { - self.dealloc(); - } - } - - unsafe fn dealloc(self) { - // Check causality - self.header().executor.with_mut(|_| {}); - self.header().future_causality.with_mut(|_| {}); - self.trailer().waker.with_mut(|_| { - // we can't check the contents of this cell as it is considered - // "uninitialized" data at this point. - }); - - drop(Box::from_raw(self.cell.as_ptr())); - } - - // ===== join handle ===== - - pub(super) unsafe fn read_output( - mut self, - dst: *mut Track>, - state: Snapshot, - ) { - if state.is_canceled() { - dst.write(Track::new(Err(JoinError::cancelled2()))); - } else { - self.core().read_output(dst); - } - - // Before transitioning the state, the waker must be read. It is - // possible that, after the transition, we are responsible for dropping - // the waker but before the waker can be read from the struct, the - // struct is deallocated. - let waker = self.read_join_waker(); - - // The operation counts as dropping the join handle - let res = self.header().state.complete_join_handle(); - - if res.is_released() { - // We are responsible for freeing the waker handle - drop(waker.assume_init()); - } - - if res.is_final_ref() { - self.dealloc(); - } - } - - pub(super) fn store_join_waker(&self, waker: &Waker) -> Snapshot { - unsafe { - self.trailer().waker.with_mut(|ptr| { - (*ptr).as_mut_ptr().replace(Some(waker.clone())); - }); - } - - let res = self.header().state.store_join_waker(); - - if res.is_complete() || res.is_canceled() { - // Drop the waker here - self.trailer() - .waker - .with_mut(|ptr| unsafe { *(*ptr).as_mut_ptr() = None }); - } - - res - } - - pub(super) fn swap_join_waker(&self, waker: &Waker, prev: Snapshot) -> Snapshot { - unsafe { - let will_wake = self - .trailer() - .waker - .with(|ptr| (*(*ptr).as_ptr()).as_ref().unwrap().will_wake(waker)); - - if will_wake { - return prev; - } - - // Acquire the lock - let state = self.header().state.unset_waker(); - - if state.is_active() { - return self.store_join_waker(waker); - } - - state - } - } - - pub(super) fn drop_join_handle_slow(mut self) { - unsafe { - // Before transitioning the state, the waker must be read. It is - // possible that, after the transition, we are responsible for dropping - // the waker but before the waker can be read from the struct, the - // struct is deallocated. - let waker = self.read_join_waker(); - - // The operation counts as dropping the join handle - let res = match self.header().state.drop_join_handle_slow() { - Ok(res) => res, - Err(res) => { - // The task output must be read & dropped - debug_assert!(!(res.is_complete() && res.is_canceled())); - - if res.is_complete() { - self.core().transition_to_consumed(); - } - - self.header().state.complete_join_handle() - } - }; - - if !(res.is_complete() | res.is_canceled()) || res.is_released() { - // We are responsible for freeing the waker handle - drop(waker.assume_init()); - } - - if res.is_final_ref() { - self.dealloc(); - } - } - } - - // ===== waker behavior ===== - - pub(super) fn wake_by_val(self) { - self.wake_by_ref(); - self.drop_waker(); - } - - pub(super) fn wake_by_ref(&self) { - if self.header().state.transition_to_notified() { - unsafe { - let executor = match self.header().executor.with(|ptr| *ptr) { - Some(executor) => executor, - None => panic!("executor should be set"), - }; - - S::schedule(executor.cast().as_ref(), self.to_task()); - } - } - } - - pub(super) fn drop_waker(self) { - if self.header().state.ref_dec() { - unsafe { - self.dealloc(); - } - } - } - - /// Cancel the task. - /// - /// `from_queue` signals the caller is cancelling the task after popping it - /// from the queue. This indicates "polling" capability. - pub(super) fn cancel(self, from_queue: bool) { - let res = if from_queue { - self.header().state.transition_to_canceled_from_queue() - } else { - match self.header().state.transition_to_canceled_from_list() { - Some(res) => res, - None => return, - } - }; - - self.do_cancel(res); - } - - fn do_cancel(mut self, res: Snapshot) { - use std::panic; - - debug_assert!(!res.is_complete()); - - let cell = unsafe { &mut self.cell.as_mut() }; - let header = &cell.header; - let core = &mut cell.core; - - // Since we transitioned the task state to `canceled`, it won't ever be - // polled again. We are now responsible for all cleanup. - // - // We have to drop the future - // - header.future_causality.with_mut(|_| { - // Guard against potential panics in the drop handler - let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { - // Drop the future - core.transition_to_consumed(); - })); - }); - - // If there is a join waker, we must notify it so it can observe the - // task was canceled. - if res.is_join_interested() && res.has_join_waker() { - // Notify the join handle. The transition to cancelled obtained a - // lock on the waker cell. - unsafe { - self.wake_join(); - } - - // Also track that we might be responsible for releasing the waker. - self.set_might_drop_join_waker_on_release(); - } - - // The `RELEASED` flag is not set yet. - assert!(!res.is_final_ref()); - - // This **can** be null if the task is being cancelled before it was - // ever polled. - let bound_executor = unsafe { self.header().executor.with(|ptr| *ptr) }; - - unsafe { - let task = self.to_task(); - - if let Some(executor) = bound_executor { - executor.cast::().as_ref().release(task); - } else { - // Just drop the task. This will release / deallocate memory. - drop(task); - } - } - } - - // ====== internal ====== - - fn complete( - mut self, - executor: &mut dyn FnMut() -> Option>, - join_interest: bool, - output: super::Result, - ) { - if join_interest { - // Store the output. The future has already been dropped - self.core().store_output(output); - } - - let executor = executor(); - let bound_executor = unsafe { self.header().executor.with(|ptr| *ptr) }; - - // Handle releasing the task. First, check if the current - // executor is the one that is bound to the task: - if executor.is_some() && executor == bound_executor { - unsafe { - // perform a local release - let task = ManuallyDrop::new(self.to_task()); - executor - .as_ref() - .unwrap() - .cast::() - .as_ref() - .release_local(&task); - - if self.transition_to_released(join_interest).is_final_ref() { - self.dealloc(); - } - } - } else { - let res = self.transition_to_complete(join_interest); - assert!(!res.is_final_ref()); - - if res.has_join_waker() { - // The release step happens later once the task has migrated back to - // the worker that owns it. At that point, the releaser **may** also - // be responsible for dropping. This fact must be tracked until - // the release step happens. - self.set_might_drop_join_waker_on_release(); - } - - unsafe { - let task = self.to_task(); - - let executor = match bound_executor { - Some(executor) => executor, - None => panic!("executor should be set"), - }; - - executor.cast::().as_ref().release(task); - } - } - } - - /// Returns `true` if the task structure should be deallocated - fn transition_to_complete(&mut self, join_interest: bool) -> Snapshot { - let res = self.header().state.transition_to_complete(); - - self.notify_join_handle(join_interest, res); - - // Transition to complete last to ensure freeing does - // not happen until the above work is done. - res - } - - /// Returns `true` if the task structure should be deallocated - fn transition_to_released(&mut self, join_interest: bool) -> Snapshot { - if join_interest { - let res1 = self.transition_to_complete(join_interest); - - let join_waker = if res1.has_join_waker() { - // At this point, the join waker may not be changed. Once we perform - // `release_task` we may no longer read from the struct but we - // **may** be responsible for dropping the waker. We do an - // optimistic read here. - unsafe { self.read_join_waker() } - } else { - MaybeUninit::uninit() - }; - - let res2 = self.header().state.release_task(); - - if res1.has_join_waker() && !res2.is_join_interested() { - debug_assert!(res2.has_join_waker()); - - // Its our responsibility to drop the waker - unsafe { - drop(join_waker.assume_init()); - } - } - - res2 - } else { - self.header().state.transition_to_released() - } - } - - fn notify_join_handle(&mut self, join_interest: bool, res: Snapshot) { - if join_interest { - if !res.is_join_interested() { - debug_assert!(!res.has_join_waker()); - - // The join handle dropped interest before we could release - // the output. We are now responsible for releasing the - // output. - self.core().transition_to_consumed(); - } else if res.has_join_waker() { - if res.is_canceled() { - // The join handle will set the output to Cancelled without - // attempting to read the output. We must drop it here. - self.core().transition_to_consumed(); - } - - // Notify the join handle. The previous transition obtains the - // lock on the waker cell. - unsafe { - self.wake_join(); - } - } - } - } - - fn might_drop_join_waker_on_release(&self) -> bool { - unsafe { - let next = *self.header().queue_next.get() as usize; - next & 1 == 1 - } - } - - fn set_might_drop_join_waker_on_release(&self) { - unsafe { - debug_assert!( - (*self.header().queue_next.get()).is_null(), - "the task's queue_next field must be null when releasing" - ); - - *self.header().queue_next.get() = 1 as *const _; - } - } - - unsafe fn wake_join(&self) { - // LOOM: ensure we can make this call - self.trailer().waker.check(); - self.trailer().waker.with_unchecked(|ptr| { - (*(*ptr).as_ptr()) - .as_ref() - .expect("waker missing") - .wake_by_ref(); - }); - } - - unsafe fn read_join_waker(&mut self) -> MaybeUninit> { - self.trailer().waker.with(|ptr| ptr.read()) - } - - unsafe fn to_task(&self) -> Task { - let ptr = self.cell.as_ptr() as *mut Header; - Task::from_raw(NonNull::new_unchecked(ptr)) - } -} diff --git a/tokio/src/task/list.rs b/tokio/src/task/list.rs deleted file mode 100644 index 85ff3dc22..000000000 --- a/tokio/src/task/list.rs +++ /dev/null @@ -1,96 +0,0 @@ -use crate::task::{Header, Task}; - -use std::fmt; -use std::marker::PhantomData; -use std::ptr::NonNull; - -pub(crate) struct OwnedList { - head: Option>, - _p: PhantomData, -} - -impl OwnedList { - pub(crate) fn new() -> OwnedList { - OwnedList { - head: None, - _p: PhantomData, - } - } - - pub(crate) fn insert(&mut self, task: &Task) { - debug_assert!(!self.contains(task)); - - unsafe { - debug_assert!((*task.header().owned_next.get()).is_none()); - debug_assert!((*task.header().owned_prev.get()).is_none()); - - let ptr = Some(task.header().into()); - - if let Some(next) = self.head { - debug_assert!((*next.as_ref().owned_prev.get()).is_none()); - *next.as_ref().owned_prev.get() = ptr; - } - - *task.header().owned_next.get() = self.head; - self.head = ptr; - } - } - - pub(crate) fn remove(&mut self, task: &Task) { - debug_assert!(self.head.is_some()); - - unsafe { - if let Some(next) = *task.header().owned_next.get() { - *next.as_ref().owned_prev.get() = *task.header().owned_prev.get(); - } - - if let Some(prev) = *task.header().owned_prev.get() { - *prev.as_ref().owned_next.get() = *task.header().owned_next.get(); - } else { - debug_assert_eq!(self.head, Some(task.header().into())); - self.head = *task.header().owned_next.get(); - } - } - } - - pub(crate) fn is_empty(&self) -> bool { - self.head.is_none() - } - - /// Transition all tasks in the list to canceled as part of the shutdown - /// process. - pub(crate) fn shutdown(&self) { - let mut curr = self.head; - - while let Some(task) = curr { - unsafe { - let vtable = task.as_ref().vtable; - (vtable.cancel)(task.as_ptr() as *mut (), false); - curr = *task.as_ref().owned_next.get(); - } - } - } - - /// Only used by debug assertions - fn contains(&self, task: &Task) -> bool { - let mut curr = self.head; - - while let Some(p) = curr { - if p == task.header().into() { - return true; - } - - unsafe { - curr = *p.as_ref().owned_next.get(); - } - } - - false - } -} - -impl fmt::Debug for OwnedList { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("OwnedList").finish() - } -} diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index ed122f03f..a2d1ceb24 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -1,13 +1,15 @@ //! Runs `!Send` futures on the current thread. +use crate::runtime::task::{self, JoinHandle, Task}; use crate::sync::AtomicWaker; -use crate::task::{self, queue::MpscQueues, JoinHandle, Schedule, Task}; +use crate::util::linked_list::LinkedList; -use std::cell::Cell; +use std::cell::{Cell, RefCell}; +use std::collections::VecDeque; +use std::fmt; use std::future::Future; use std::pin::Pin; -use std::ptr::{self, NonNull}; -use std::rc::Rc; -use std::task::{Context, Poll}; +use std::sync::{Arc, Mutex}; +use std::task::Poll; use pin_project_lite::pin_project; @@ -106,35 +108,51 @@ cfg_rt_util! { /// [local task set]: struct.LocalSet.html /// [`Runtime::block_on`]: ../struct.Runtime.html#method.block_on /// [`task::spawn_local`]: fn.spawn.html - #[derive(Debug)] pub struct LocalSet { - scheduler: Rc, + /// Current scheduler tick + tick: Cell, + + /// State available from thread-local + context: Context, } } -#[derive(Debug)] -struct Scheduler { - tick: Cell, +/// State available from the thread-local +struct Context { + /// Owned task set and local run queue + tasks: RefCell, - queues: MpscQueues, + /// State shared between threads. + shared: Arc, +} - /// Used to notify the `LocalFuture` when a task in the local task set is - /// notified. +struct Tasks { + /// Collection of all active tasks spawned onto this executor. + owned: LinkedList>>, + + /// Local run queue sender and receiver. + queue: VecDeque>>, +} + +/// LocalSet state shared between threads. +struct Shared { + /// Remote run queue sender + queue: Mutex>>>, + + /// Wake the `LocalSet` task waker: AtomicWaker, } pin_project! { #[derive(Debug)] - struct LocalFuture { - scheduler: Rc, + struct RunUntil<'a, F> { + local_set: &'a LocalSet, #[pin] future: F, } } -thread_local! { - static CURRENT_TASK_SET: Cell>> = Cell::new(None); -} +scoped_thread_local!(static CURRENT: Context); cfg_rt_util! { /// Spawns a `!Send` future on the local task set. @@ -173,32 +191,43 @@ cfg_rt_util! { F: Future + 'static, F::Output: 'static, { - CURRENT_TASK_SET.with(|current| { - let current = current - .get() - .expect("`spawn_local` called from outside of a task::LocalSet!"); - let (task, handle) = task::joinable_local(future); - unsafe { - // safety: this function is unsafe to call outside of the local - // thread. Since the call above to get the current task set - // would not succeed if we were outside of a local set, this is - // safe. - current.as_ref().queues.push_local(task); - } + CURRENT.with(|maybe_cx| { + let cx = maybe_cx + .expect("`spawn_local` called from outside of a `task::LocalSet`"); + // Safety: Tasks are only polled and dropped from the thread that + // spawns them. + let (task, handle) = unsafe { task::joinable_local(future) }; + cx.tasks.borrow_mut().queue.push_back(task); handle }) } } +/// Initial queue capacity +const INITIAL_CAPACITY: usize = 64; + /// Max number of tasks to poll per tick. const MAX_TASKS_PER_TICK: usize = 61; +/// How often it check the remote queue first +const REMOTE_FIRST_INTERVAL: u8 = 31; + impl LocalSet { /// Returns a new local task set. - pub fn new() -> Self { - Self { - scheduler: Rc::new(Scheduler::new()), + pub fn new() -> LocalSet { + LocalSet { + tick: Cell::new(0), + context: Context { + tasks: RefCell::new(Tasks { + owned: LinkedList::new(), + queue: VecDeque::with_capacity(INITIAL_CAPACITY), + }), + shared: Arc::new(Shared { + queue: Mutex::new(VecDeque::with_capacity(INITIAL_CAPACITY)), + waker: AtomicWaker::new(), + }), + }, } } @@ -243,12 +272,8 @@ impl LocalSet { F: Future + 'static, F::Output: 'static, { - let (task, handle) = task::joinable_local(future); - unsafe { - // safety: since `LocalSet` is not Send or Sync, this is - // always being called from the local thread. - self.scheduler.queues.push_local(task); - } + let (task, handle) = unsafe { task::joinable_local(future) }; + self.context.tasks.borrow_mut().queue.push_back(task); handle } @@ -353,25 +378,83 @@ impl LocalSet { where F: Future, { - let scheduler = self.scheduler.clone(); - let future = LocalFuture { scheduler, future }; - future.await + let run_until = RunUntil { + future, + local_set: self, + }; + run_until.await + } + + /// Tick the scheduler, returning whether the local future needs to be + /// notified again. + fn tick(&self) -> bool { + for _ in 0..MAX_TASKS_PER_TICK { + match self.next_task() { + // Run the task + // + // Safety: As spawned tasks are `!Send`, `run_unchecked` must be + // used. We are responsible for maintaining the invariant that + // `run_unchecked` is only called on threads that spawned the + // task initially. Because `LocalSet` itself is `!Send`, and + // `spawn_local` spawns into the `LocalSet` on the current + // thread, the invariant is maintained. + Some(task) => task.run(), + // We have fully drained the queue of notified tasks, so the + // local future doesn't need to be notified again — it can wait + // until something else wakes a task in the local set. + None => return false, + } + } + + true + } + + fn next_task(&self) -> Option>> { + let tick = self.tick.get(); + self.tick.set(tick.wrapping_add(1)); + + if tick % REMOTE_FIRST_INTERVAL == 0 { + self.context + .shared + .queue + .lock() + .unwrap() + .pop_front() + .or_else(|| self.context.tasks.borrow_mut().queue.pop_front()) + } else { + self.context + .tasks + .borrow_mut() + .queue + .pop_front() + .or_else(|| self.context.shared.queue.lock().unwrap().pop_front()) + } + } + + fn with(&self, f: impl FnOnce() -> T) -> T { + CURRENT.set(&self.context, f) + } +} + +impl fmt::Debug for LocalSet { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("LocalSet").finish() } } impl Future for LocalSet { type Output = (); - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let scheduler = self.as_ref().scheduler.clone(); - scheduler.waker.register_by_ref(cx.waker()); + fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { + // Register the waker before starting to work + self.context.shared.waker.register_by_ref(cx.waker()); - if scheduler.with(|| scheduler.tick()) { + if self.with(|| self.tick()) { // If `tick` returns true, we need to notify the local future again: // there are still tasks remaining in the run queue. cx.waker().wake_by_ref(); Poll::Pending - } else if scheduler.is_empty() { + } else if self.context.tasks.borrow().owned.is_empty() { // If the scheduler has no remaining futures, we're done! Poll::Ready(()) } else { @@ -384,27 +467,59 @@ impl Future for LocalSet { } impl Default for LocalSet { - fn default() -> Self { - Self::new() + fn default() -> LocalSet { + LocalSet::new() + } +} + +impl Drop for LocalSet { + fn drop(&mut self) { + self.with(|| { + // Loop required here to ensure borrow is dropped between iterations + #[allow(clippy::while_let_loop)] + loop { + let task = match self.context.tasks.borrow_mut().owned.pop_back() { + Some(task) => task, + None => break, + }; + + // Safety: same as `run_unchecked`. + task.shutdown(); + } + + for task in self.context.tasks.borrow_mut().queue.drain(..) { + task.shutdown(); + } + + for task in self.context.shared.queue.lock().unwrap().drain(..) { + task.shutdown(); + } + + assert!(self.context.tasks.borrow().owned.is_empty()); + }); } } // === impl LocalFuture === -impl Future for LocalFuture { - type Output = F::Output; +impl Future for RunUntil<'_, T> { + type Output = T::Output; - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let this = self.project(); - let scheduler = this.scheduler; - let mut future = this.future; - scheduler.waker.register_by_ref(cx.waker()); - scheduler.with(|| { - if let Poll::Ready(output) = future.as_mut().poll(cx) { + fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll { + let me = self.project(); + + me.local_set.with(|| { + me.local_set + .context + .shared + .waker + .register_by_ref(cx.waker()); + + if let Poll::Ready(output) = me.future.poll(cx) { return Poll::Ready(output); } - if scheduler.tick() { + if me.local_set.tick() { // If `tick` returns `true`, we need to notify the local future again: // there are still tasks remaining in the run queue. cx.waker().wake_by_ref(); @@ -415,144 +530,50 @@ impl Future for LocalFuture { } } -// === impl Scheduler === - -impl Schedule for Scheduler { - fn bind(&self, task: &Task) { - assert!(self.is_current()); - unsafe { - self.queues.add_task(task); - } +impl Shared { + /// Schedule the provided task on the scheduler. + fn schedule(&self, task: task::Notified>) { + CURRENT.with(|maybe_cx| match maybe_cx { + Some(cx) if cx.shared.ptr_eq(self) => { + cx.tasks.borrow_mut().queue.push_back(task); + } + _ => { + self.queue.lock().unwrap().push_back(task); + self.waker.wake(); + } + }); } - fn release(&self, task: Task) { - // This will be called when dropping the local runtime. - self.queues.release_remote(task); - } - - fn release_local(&self, task: &Task) { - debug_assert!(self.is_current()); - unsafe { - self.queues.release_local(task); - } - } - - fn schedule(&self, task: Task) { - if self.is_current() { - unsafe { self.queues.push_local(task) }; - } else { - let mut lock = self.queues.remote(); - lock.schedule(task, false); - - self.waker.wake(); - - drop(lock); - } + fn ptr_eq(&self, other: &Shared) -> bool { + self as *const _ == other as *const _ } } -impl Scheduler { - fn new() -> Self { - Self { - tick: Cell::new(0), - queues: MpscQueues::new(), - waker: AtomicWaker::new(), - } - } - - fn with(&self, f: impl FnOnce() -> F) -> F { - struct Entered<'a> { - current: &'a Cell>>, - } - - impl<'a> Drop for Entered<'a> { - fn drop(&mut self) { - self.current.set(None); - } - } - - CURRENT_TASK_SET.with(|current| { - let prev = current.replace(Some(NonNull::from(self))); - assert!(prev.is_none(), "nested call to local::Scheduler::with"); - let _entered = Entered { current }; - f() +impl task::Schedule for Arc { + fn bind(task: Task) -> Arc { + CURRENT.with(|maybe_cx| { + let cx = maybe_cx.expect("scheduler context missing"); + cx.tasks.borrow_mut().owned.push_front(task); + cx.shared.clone() }) } - fn is_current(&self) -> bool { - CURRENT_TASK_SET - .try_with(|current| { - current - .get() - .iter() - .any(|current| ptr::eq(current.as_ptr(), self as *const _)) - }) - .unwrap_or(false) + fn release(&self, task: &Task) -> Option> { + use std::ptr::NonNull; + + CURRENT.with(|maybe_cx| { + let cx = maybe_cx.expect("scheduler context missing"); + + assert!(cx.shared.ptr_eq(self)); + + let ptr = NonNull::from(task.header()); + // safety: task must be contained by list. It is inserted into the + // list in `bind`. + unsafe { cx.tasks.borrow_mut().owned.remove(ptr) } + }) } - /// Tick the scheduler, returning whether the local future needs to be - /// notified again. - fn tick(&self) -> bool { - assert!(self.is_current()); - for _ in 0..MAX_TASKS_PER_TICK { - let tick = self.tick.get().wrapping_add(1); - self.tick.set(tick); - - let task = match unsafe { - // safety: we must be on the local thread to call this. The assertion - // the top of this method ensures that `tick` is only called locally. - self.queues.next_task(tick) - } { - Some(task) => task, - // We have fully drained the queue of notified tasks, so the - // local future doesn't need to be notified again — it can wait - // until something else wakes a task in the local set. - None => return false, - }; - - if let Some(task) = task.run(&mut || Some(self.into())) { - unsafe { - // safety: we must be on the local thread to call this. The - // the top of this method ensures that `tick` is only called locally. - self.queues.push_local(task); - } - } - } - - true - } - - fn is_empty(&self) -> bool { - unsafe { - // safety: this method may not be called from threads other than the - // thread that owns the `Queues`. since `Scheduler` is not `Send` or - // `Sync`, that shouldn't happen. - !self.queues.has_tasks_remaining() - } - } -} - -impl Drop for Scheduler { - fn drop(&mut self) { - unsafe { - // safety: these functions are unsafe to call outside of the local - // thread. Since the `Scheduler` type is not `Send` or `Sync`, we - // know it will be dropped only from the local thread. - self.queues.shutdown(); - - // Wait until all tasks have been released. - // XXX: this is a busy loop, but we don't really have any way to park - // the thread here? - loop { - self.queues.drain_pending_drop(); - self.queues.drain_queues(); - - if !self.queues.has_tasks_remaining() { - break; - } - - std::thread::yield_now(); - } - } + fn schedule(&self, task: task::Notified) { + Shared::schedule(self, task); } } diff --git a/tokio/src/task/mod.rs b/tokio/src/task/mod.rs index 073215e6e..5c89393a5 100644 --- a/tokio/src/task/mod.rs +++ b/tokio/src/task/mod.rs @@ -224,39 +224,11 @@ cfg_blocking! { } cfg_rt_core! { - mod core; - use self::core::Cell; - pub(crate) use self::core::Header; - - mod error; - pub use self::error::JoinError; - - mod harness; - use self::harness::Harness; - - mod join; - #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 - pub use self::join::JoinHandle; - - mod list; - pub(crate) use self::list::OwnedList; - - pub(crate) mod queue; - - mod raw; - use self::raw::RawTask; + pub use crate::runtime::task::{JoinError, JoinHandle}; mod spawn; pub use spawn::spawn; - mod stack; - pub(crate) use self::stack::TransferStack; - - mod state; - use self::state::{Snapshot, State}; - - mod waker; - mod yield_now; pub use yield_now::yield_now; } @@ -268,144 +240,3 @@ cfg_rt_util! { mod task_local; pub use task_local::LocalKey; } - -cfg_rt_core! { - /// Unit tests - #[cfg(test)] - mod tests; - - use std::future::Future; - use std::marker::PhantomData; - use std::ptr::NonNull; - use std::{fmt, mem}; - - /// An owned handle to the task, tracked by ref count - pub(crate) struct Task { - raw: RawTask, - _p: PhantomData, - } - - unsafe impl Send for Task {} - - /// Task result sent back - pub(crate) type Result = std::result::Result; - - pub(crate) trait Schedule: Sized + 'static { - /// Bind a task to the executor. - /// - /// Guaranteed to be called from the thread that called `poll` on the task. - fn bind(&self, task: &Task); - - /// The task has completed work and is ready to be released. The scheduler - /// is free to drop it whenever. - fn release(&self, task: Task); - - /// The has been completed by the executor it was bound to. - fn release_local(&self, task: &Task); - - /// Schedule the task - fn schedule(&self, task: Task); - } - - /// Marker trait indicating that a scheduler can only schedule tasks which - /// implement `Send`. - /// - /// Schedulers that implement this trait may not schedule `!Send` futures. If - /// trait is implemented, the corresponding `Task` type will implement `Send`. - pub(crate) trait ScheduleSendOnly: Schedule + Send + Sync {} - - /// Create a new task with an associated join handle - pub(crate) fn joinable(task: T) -> (Task, JoinHandle) - where - T: Future + Send + 'static, - S: ScheduleSendOnly, - { - let raw = RawTask::new_joinable::<_, S>(task); - - let task = Task { - raw, - _p: PhantomData, - }; - - let join = JoinHandle::new(raw); - - (task, join) - } - - cfg_rt_util! { - /// Create a new `!Send` task with an associated join handle - pub(crate) fn joinable_local(task: T) -> (Task, JoinHandle) - where - T: Future + 'static, - S: Schedule, - { - let raw = RawTask::new_joinable_local::<_, S>(task); - - let task = Task { - raw, - _p: PhantomData, - }; - - let join = JoinHandle::new(raw); - - (task, join) - } - } - - impl Task { - pub(crate) unsafe fn from_raw(ptr: NonNull
) -> Task { - Task { - raw: RawTask::from_raw(ptr), - _p: PhantomData, - } - } - - pub(crate) fn header(&self) -> &Header { - self.raw.header() - } - - pub(crate) fn into_raw(self) -> NonNull
{ - let raw = self.raw.into_raw(); - mem::forget(self); - raw - } - } - - impl Task { - /// Returns `self` when the task needs to be immediately re-scheduled - pub(crate) fn run(self, mut executor: F) -> Option - where - F: FnMut() -> Option>, - { - if unsafe { - self.raw - .poll(&mut || executor().map(|ptr| ptr.cast::<()>())) - } { - Some(self) - } else { - // Cleaning up the `Task` instance is done from within the poll - // function. - mem::forget(self); - None - } - } - - /// Pre-emptively cancel the task as part of the shutdown process. - pub(crate) fn shutdown(self) { - self.raw.cancel_from_queue(); - mem::forget(self); - } - } - - impl Drop for Task { - fn drop(&mut self) { - self.raw.drop_task(); - } - } - - impl fmt::Debug for Task { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Task").finish() - } - } -} diff --git a/tokio/src/task/queue.rs b/tokio/src/task/queue.rs deleted file mode 100644 index 5a2f5473f..000000000 --- a/tokio/src/task/queue.rs +++ /dev/null @@ -1,338 +0,0 @@ -use super::{OwnedList, Schedule, Task, TransferStack}; -use std::{ - cell::UnsafeCell, - collections::VecDeque, - fmt, - sync::{Mutex, MutexGuard}, -}; - -/// A set of multi-producer, single consumer task queues, suitable for use by a -/// single-threaded scheduler. -/// -/// This consists of a list of _all_ tasks bound to the scheduler, a run queue -/// of tasks notified from the thread the scheduler is running on (the "local -/// queue"), a run queue of tasks notified from another thread (the "remote -/// queue"), and a stack of tasks released from other threads which will -/// eventually need to be dropped by the scheduler on its own thread ("pending -/// drop"). -/// -/// Submitting tasks to or popping tasks from the local queue is unsafe, as it -/// must only be performed on the same thread as the scheduler. -pub(crate) struct MpscQueues { - /// List of all active tasks spawned onto this executor. - /// - /// # Safety - /// - /// Must only be accessed from the primary thread - owned_tasks: UnsafeCell>, - - /// Local run queue. - /// - /// Tasks notified from the current thread are pushed into this queue. - /// - /// # Safety - /// - /// References should not be handed out. Only call `push` / `pop` functions. - /// Only call from the owning thread. - local_queue: UnsafeCell>>, - - /// Remote run queue. - /// - /// Tasks notified from another thread are pushed into this queue. - remote_queue: Mutex>, - - /// Tasks pending drop - pending_drop: TransferStack, -} - -pub(crate) struct RemoteQueue { - /// FIFO list of tasks - queue: VecDeque>, - - /// `true` when a task can be pushed into the queue, `false` otherwise. - open: bool, -} - -// === impl Queues === - -impl MpscQueues -where - S: Schedule + 'static, -{ - pub(crate) const INITIAL_CAPACITY: usize = 64; - - /// How often to check the remote queue first - pub(crate) const CHECK_REMOTE_INTERVAL: u8 = 13; - - pub(crate) fn new() -> Self { - Self { - owned_tasks: UnsafeCell::new(OwnedList::new()), - local_queue: UnsafeCell::new(VecDeque::with_capacity(Self::INITIAL_CAPACITY)), - pending_drop: TransferStack::new(), - remote_queue: Mutex::new(RemoteQueue { - queue: VecDeque::with_capacity(Self::INITIAL_CAPACITY), - open: true, - }), - } - } - - /// Adds a new task to the scheduler. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn add_task(&self, task: &Task) { - (*self.owned_tasks.get()).insert(task); - } - - /// Pushes a task to the local queue. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn push_local(&self, task: Task) { - (*self.local_queue.get()).push_back(task); - } - - /// Removes a task from the local queue. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn release_local(&self, task: &Task) { - (*self.owned_tasks.get()).remove(task); - } - - /// Locks the remote queue, returning a `MutexGuard`. - /// - /// This can be used to push to the remote queue and perform other - /// operations while holding the lock. - /// - /// # Panics - /// - /// If the remote queue mutex is poisoned. - pub(crate) fn remote(&self) -> MutexGuard<'_, RemoteQueue> { - self.remote_queue - .lock() - .expect("failed to lock remote queue") - } - - /// Releases a task from outside of the thread that owns the scheduler. - /// - /// This simply pushes the task to the pending drop queue. - pub(crate) fn release_remote(&self, task: Task) { - self.pending_drop.push(task); - } - - /// Returns the next task from the remote *or* local queue. - /// - /// Typically, this checks the local queue before the remote queue, and only - /// checks the remote queue if the local queue is empty. However, to avoid - /// starving the remote queue, it is checked first every - /// `CHECK_REMOTE_INTERVAL` ticks. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn next_task(&self, tick: u8) -> Option> { - if 0 == tick % Self::CHECK_REMOTE_INTERVAL { - self.next_remote_task().or_else(|| self.next_local_task()) - } else { - self.next_local_task().or_else(|| self.next_remote_task()) - } - } - - /// Returns the next task from the local queue. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn next_local_task(&self) -> Option> { - (*self.local_queue.get()).pop_front() - } - - /// Returns the next task from the remote queue. - /// - /// # Panics - /// - /// If the mutex around the remote queue is poisoned _and_ the current - /// thread is not already panicking. This is safe to call in a `Drop` impl. - pub(crate) fn next_remote_task(&self) -> Option> { - // there is no semantic information in the `PoisonError`, and it - // doesn't implement `Debug`, but clippy thinks that it's bad to - // match all errors here... - #[allow(clippy::match_wild_err_arm)] - let mut lock = match self.remote_queue.lock() { - // If the lock is poisoned, but the thread is already panicking, - // avoid a double panic. This is necessary since `next_task` (which - // calls `next_remote_task`) can be called in the `Drop` impl. - Err(_) if std::thread::panicking() => return None, - Err(_) => panic!("mutex poisoned"), - Ok(lock) => lock, - }; - lock.queue.pop_front() - } - - /// Returns `true` if any owned tasks are still bound to this scheduler. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn has_tasks_remaining(&self) -> bool { - !(*self.owned_tasks.get()).is_empty() - } - - /// Drains any tasks that have previously been released from other threads. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn drain_pending_drop(&self) { - for task in self.pending_drop.drain() { - (*self.owned_tasks.get()).remove(&task); - drop(task); - } - } - - /// Shuts down the queues. - /// - /// This performs the following operations: - /// - /// 1. Close the remote queue (so that it will no longer accept new tasks). - /// 2. Drain the remote queue and shut down all tasks. - /// 3. Drain the local queue and shut down all tasks. - /// 4. Shut down the owned task list. - /// 5. Drain the list of tasks dropped externally and remove them from the - /// owned task list. - /// - /// This method should be called before dropping a `Queues`. It is provided - /// as a method rather than a `Drop` impl because types that own a `Queues` - /// wish to perform other work in their `Drop` implementations _after_ - /// shutting down the task queues. - /// - /// # Safety - /// - /// This method accesses the local task queue, and therefore *must* be - /// called only from the thread that owns the scheduler. - /// - /// # Panics - /// - /// If the mutex around the remote queue is poisoned _and_ the current - /// thread is not already panicking. This is safe to call in a `Drop` impl. - pub(crate) unsafe fn shutdown(&self) { - // Close and drain the remote queue. - self.close_remote(); - - // Drain the local queue. - self.close_local(); - - // Release owned tasks - self.shutdown_owned_tasks(); - - // Drain tasks pending drop. - self.drain_pending_drop(); - } - - /// Drains both the local and remote run queues, shutting down any tasks. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - pub(crate) unsafe fn drain_queues(&self) { - self.close_local(); - self.close_remote(); - } - - /// Shuts down the scheduler's owned task list. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - unsafe fn shutdown_owned_tasks(&self) { - (*self.owned_tasks.get()).shutdown(); - } - - /// Drains the remote queue, and shut down its tasks. - /// - /// This closes the remote queue. Any additional tasks added to it will be - /// shut down instead. - /// - /// # Panics - /// If the mutex around the remote queue is poisoned _and_ the current - /// thread is not already panicking. This is safe to call in a `Drop` impl. - fn close_remote(&self) { - loop { - #[allow(clippy::match_wild_err_arm)] - let mut lock = match self.remote_queue.lock() { - // If the lock is poisoned, but the thread is already panicking, - // avoid a double panic. This is necessary since this fn can be - // called in a drop impl. - Err(_) if std::thread::panicking() => return, - Err(_) => panic!("mutex poisoned"), - Ok(lock) => lock, - }; - lock.open = false; - - if let Some(task) = lock.queue.pop_front() { - // Release lock before dropping task, in case - // task tries to re-schedule in its Drop. - drop(lock); - task.shutdown(); - } else { - return; - } - } - } - - /// Drains the local queue, and shut down its tasks. - /// - /// # Safety - /// - /// This *must* be called only from the thread that owns the scheduler. - unsafe fn close_local(&self) { - while let Some(task) = self.next_local_task() { - task.shutdown(); - } - } -} - -impl fmt::Debug for MpscQueues { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("MpscQueues") - .field("owned_tasks", &self.owned_tasks) - .field("remote_queue", &self.remote_queue) - .field("local_queue", &self.local_queue) - .finish() - } -} - -// === impl RemoteQueue === - -impl RemoteQueue -where - S: Schedule, -{ - /// Schedule a remote task. - /// - /// If the queue is open to accept new tasks, the task is pushed to the back - /// of the queue. Otherwise, if the queue is closed (the scheduler is - /// shutting down), the new task will be shut down immediately. - /// - /// `spawn` should be set if the caller is spawning a new task. - pub(crate) fn schedule(&mut self, task: Task, spawn: bool) { - if !spawn || self.open { - self.queue.push_back(task); - } else { - task.shutdown(); - } - } -} - -impl fmt::Debug for RemoteQueue { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("RemoteQueue") - .field("queue", &self.queue) - .field("open", &self.open) - .finish() - } -} diff --git a/tokio/src/task/raw.rs b/tokio/src/task/raw.rs deleted file mode 100644 index 93bb79c82..000000000 --- a/tokio/src/task/raw.rs +++ /dev/null @@ -1,197 +0,0 @@ -use crate::loom::alloc::Track; -use crate::task::Cell; -use crate::task::Harness; -use crate::task::{Header, Schedule, ScheduleSendOnly}; -use crate::task::{Snapshot, State}; - -use std::future::Future; -use std::ptr::NonNull; -use std::task::Waker; - -/// Raw task handle -pub(super) struct RawTask { - ptr: NonNull
, -} - -pub(super) struct Vtable { - /// Poll the future - pub(super) poll: unsafe fn(*mut (), &mut dyn FnMut() -> Option>) -> bool, - - /// The task handle has been dropped and the join waker needs to be dropped - /// or the task struct needs to be deallocated - pub(super) drop_task: unsafe fn(*mut ()), - - /// Read the task output - pub(super) read_output: unsafe fn(*mut (), *mut (), Snapshot), - - /// Store the join handle's waker - /// - /// Returns a snapshot of the state **after** the transition - pub(super) store_join_waker: unsafe fn(*mut (), &Waker) -> Snapshot, - - /// Replace the join handle's waker - /// - /// Returns a snapshot of the state **after** the transition - pub(super) swap_join_waker: unsafe fn(*mut (), &Waker, Snapshot) -> Snapshot, - - /// The join handle has been dropped - pub(super) drop_join_handle_slow: unsafe fn(*mut ()), - - /// The task is being canceled - pub(super) cancel: unsafe fn(*mut (), bool), -} - -/// Get the vtable for the requested `T` and `S` generics. -pub(super) fn vtable() -> &'static Vtable { - &Vtable { - poll: poll::, - drop_task: drop_task::, - read_output: read_output::, - store_join_waker: store_join_waker::, - swap_join_waker: swap_join_waker::, - drop_join_handle_slow: drop_join_handle_slow::, - cancel: cancel::, - } -} - -cfg_rt_util! { - impl RawTask { - pub(super) fn new_joinable_local(task: T) -> RawTask - where - T: Future + 'static, - S: Schedule, - { - RawTask::new::<_, S>(task, State::new_joinable()) - } - } -} - -impl RawTask { - pub(super) fn new_joinable(task: T) -> RawTask - where - T: Future + Send + 'static, - S: ScheduleSendOnly, - { - RawTask::new::<_, S>(task, State::new_joinable()) - } - - fn new(task: T, state: State) -> RawTask - where - T: Future + 'static, - S: Schedule, - { - let ptr = Box::into_raw(Cell::new::(task, state)); - let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) }; - - RawTask { ptr } - } - - pub(super) unsafe fn from_raw(ptr: NonNull
) -> RawTask { - RawTask { ptr } - } - - /// Returns a reference to the task's meta structure. - /// - /// Safe as `Header` is `Sync`. - pub(super) fn header(&self) -> &Header { - unsafe { self.ptr.as_ref() } - } - - /// Returns a raw pointer to the task's meta structure. - pub(super) fn into_raw(self) -> NonNull
{ - self.ptr - } - - /// Safety: mutual exclusion is required to call this function. - /// - /// Returns `true` if the task needs to be scheduled again. - pub(super) unsafe fn poll(self, executor: &mut dyn FnMut() -> Option>) -> bool { - // Get the vtable without holding a ref to the meta struct. This is done - // because a mutable reference to the task is passed into the poll fn. - let vtable = self.header().vtable; - - (vtable.poll)(self.ptr.as_ptr() as *mut (), executor) - } - - pub(super) fn drop_task(self) { - let vtable = self.header().vtable; - unsafe { - (vtable.drop_task)(self.ptr.as_ptr() as *mut ()); - } - } - - pub(super) unsafe fn read_output(self, dst: *mut (), state: Snapshot) { - let vtable = self.header().vtable; - (vtable.read_output)(self.ptr.as_ptr() as *mut (), dst, state); - } - - pub(super) fn store_join_waker(self, waker: &Waker) -> Snapshot { - let vtable = self.header().vtable; - unsafe { (vtable.store_join_waker)(self.ptr.as_ptr() as *mut (), waker) } - } - - pub(super) fn swap_join_waker(self, waker: &Waker, prev: Snapshot) -> Snapshot { - let vtable = self.header().vtable; - unsafe { (vtable.swap_join_waker)(self.ptr.as_ptr() as *mut (), waker, prev) } - } - - pub(super) fn drop_join_handle_slow(self) { - let vtable = self.header().vtable; - unsafe { (vtable.drop_join_handle_slow)(self.ptr.as_ptr() as *mut ()) } - } - - pub(super) fn cancel_from_queue(self) { - let vtable = self.header().vtable; - unsafe { (vtable.cancel)(self.ptr.as_ptr() as *mut (), true) } - } -} - -impl Clone for RawTask { - fn clone(&self) -> Self { - RawTask { ptr: self.ptr } - } -} - -impl Copy for RawTask {} - -unsafe fn poll( - ptr: *mut (), - executor: &mut dyn FnMut() -> Option>, -) -> bool { - let harness = Harness::::from_raw(ptr); - harness.poll(executor) -} - -unsafe fn drop_task(ptr: *mut ()) { - let harness = Harness::::from_raw(ptr); - harness.drop_task(); -} - -unsafe fn read_output(ptr: *mut (), dst: *mut (), state: Snapshot) { - let harness = Harness::::from_raw(ptr); - harness.read_output(dst as *mut Track>, state); -} - -unsafe fn store_join_waker(ptr: *mut (), waker: &Waker) -> Snapshot { - let harness = Harness::::from_raw(ptr); - harness.store_join_waker(waker) -} - -unsafe fn swap_join_waker( - ptr: *mut (), - waker: &Waker, - prev: Snapshot, -) -> Snapshot { - let harness = Harness::::from_raw(ptr); - harness.swap_join_waker(waker, prev) -} - -unsafe fn drop_join_handle_slow(ptr: *mut ()) { - let harness = Harness::::from_raw(ptr); - harness.drop_join_handle_slow() -} - -unsafe fn cancel(ptr: *mut (), from_queue: bool) { - let harness = Harness::::from_raw(ptr); - harness.cancel(from_queue) -} diff --git a/tokio/src/task/stack.rs b/tokio/src/task/stack.rs deleted file mode 100644 index 36ebb797e..000000000 --- a/tokio/src/task/stack.rs +++ /dev/null @@ -1,88 +0,0 @@ -use crate::loom::sync::atomic::AtomicPtr; -use crate::task::{Header, Task}; - -use std::marker::PhantomData; -use std::ptr::{self, NonNull}; -use std::sync::atomic::Ordering::{Acquire, Relaxed, Release}; - -/// Concurrent stack of tasks, used to pass ownership of a task from one worker -/// to another. -pub(crate) struct TransferStack { - head: AtomicPtr
, - _p: PhantomData, -} - -impl TransferStack { - pub(crate) fn new() -> TransferStack { - TransferStack { - head: AtomicPtr::new(ptr::null_mut()), - _p: PhantomData, - } - } - - pub(crate) fn push(&self, task: Task) { - unsafe { - let task = task.into_raw(); - - let next = (*task.as_ref().queue_next.get()) as usize; - - // At this point, the queue_next field may also be used to track - // whether or not the task must drop the join waker. - debug_assert_eq!(0, next & !1); - - // We don't care about any memory associated w/ setting the `head` - // field, just the current value. - let mut curr = self.head.load(Relaxed); - - loop { - *task.as_ref().queue_next.get() = (next | curr as usize) as *const _; - - let res = - self.head - .compare_exchange(curr, task.as_ptr() as *mut _, Release, Relaxed); - - match res { - Ok(_) => return, - Err(actual) => { - curr = actual; - } - } - } - } - } - - pub(crate) fn drain(&self) -> impl Iterator> { - struct Iter(*mut Header, PhantomData); - - impl Iterator for Iter { - type Item = Task; - - fn next(&mut self) -> Option> { - let task = NonNull::new(self.0)?; - - unsafe { - let next = *task.as_ref().queue_next.get() as usize; - - // remove the data bit - self.0 = (next & !1) as *mut _; - - Some(Task::from_raw(task)) - } - } - } - - impl Drop for Iter { - fn drop(&mut self) { - use std::process; - - if !self.0.is_null() { - // we have bugs - process::abort(); - } - } - } - - let ptr = self.head.swap(ptr::null_mut(), Acquire); - Iter(ptr, PhantomData) - } -} diff --git a/tokio/src/task/state.rs b/tokio/src/task/state.rs deleted file mode 100644 index 89e529e9f..000000000 --- a/tokio/src/task/state.rs +++ /dev/null @@ -1,497 +0,0 @@ -use crate::loom::sync::atomic::AtomicUsize; - -use std::fmt; -use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; -use std::usize; - -pub(super) struct State { - val: AtomicUsize, -} - -/// Current state value -#[derive(Copy, Clone)] -pub(super) struct Snapshot(usize); - -/// The task is currently being run. -const RUNNING: usize = 0b00_0001; - -/// The task has been notified by a waker. -const NOTIFIED: usize = 0b00_0010; - -/// The task is complete. -/// -/// Once this bit is set, it is never unset -const COMPLETE: usize = 0b00_0100; - -/// The primary task handle has been dropped. -const RELEASED: usize = 0b00_1000; - -/// The join handle is still around -const JOIN_INTEREST: usize = 0b01_0000; - -/// A join handle waker has been set -const JOIN_WAKER: usize = 0b10_0000; - -/// The task has been forcibly canceled. -const CANCELLED: usize = 0b100_0000; - -/// All bits -const LIFECYCLE_MASK: usize = - RUNNING | NOTIFIED | COMPLETE | RELEASED | JOIN_INTEREST | JOIN_WAKER | CANCELLED; - -/// Bits used by the waker ref count portion of the state. -/// -/// Ref counts only cover **wakers**. Other handles are tracked with other state -/// bits. -const WAKER_COUNT_MASK: usize = usize::MAX - LIFECYCLE_MASK; - -/// Number of positions to shift the ref count -const WAKER_COUNT_SHIFT: usize = WAKER_COUNT_MASK.count_zeros() as usize; - -/// One ref count -const WAKER_ONE: usize = 1 << WAKER_COUNT_SHIFT; - -/// Initial state -const INITIAL_STATE: usize = NOTIFIED; - -/// All transitions are performed via RMW operations. This establishes an -/// unambiguous modification order. -impl State { - /// Starts with a ref count of 2 - pub(super) fn new_joinable() -> State { - State { - val: AtomicUsize::new(INITIAL_STATE | JOIN_INTEREST), - } - } - - /// Loads the current state, establishes `Acquire` ordering. - pub(super) fn load(&self) -> Snapshot { - Snapshot(self.val.load(Acquire)) - } - - /// Transitions a task to the `Running` state. - /// - /// Returns a snapshot of the state **after** the transition. - pub(super) fn transition_to_running(&self) -> Snapshot { - const DELTA: usize = RUNNING | NOTIFIED; - - let prev = Snapshot(self.val.fetch_xor(DELTA, Acquire)); - assert!(prev.is_notified()); - - if prev.is_running() { - // We were signalled to cancel - // - // Apply the state - let prev = self.val.fetch_or(CANCELLED, AcqRel); - return Snapshot(prev | CANCELLED); - } - - assert!(!prev.is_running()); - - let next = Snapshot(prev.0 ^ DELTA); - - assert!(next.is_running()); - assert!(!next.is_notified()); - - next - } - - /// Transitions the task from `Running` -> `Idle`. - /// - /// Returns a snapshot of the state **after** the transition. - pub(super) fn transition_to_idle(&self) -> Snapshot { - const DELTA: usize = RUNNING; - - let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel)); - - if !prev.is_running() { - // We were signaled to cancel. - // - // Apply the state - let prev = self.val.fetch_or(CANCELLED, AcqRel); - return Snapshot(prev | CANCELLED); - } - - let next = Snapshot(prev.0 ^ DELTA); - - assert!(!next.is_running()); - - next - } - - /// Transitions the task from `Running` -> `Complete`. - /// - /// Returns a snapshot of the state **after** the transition. - pub(super) fn transition_to_complete(&self) -> Snapshot { - const DELTA: usize = RUNNING | COMPLETE; - - let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel)); - - assert!(!prev.is_complete()); - - let next = Snapshot(prev.0 ^ DELTA); - - assert!(next.is_complete()); - - next - } - - /// Transitions the task from `Running` -> `Released`. - /// - /// Returns a snapshot of the state **after** the transition. - pub(super) fn transition_to_released(&self) -> Snapshot { - const DELTA: usize = RUNNING | COMPLETE | RELEASED; - - let prev = Snapshot(self.val.fetch_xor(DELTA, AcqRel)); - - assert!(prev.is_running()); - assert!(!prev.is_complete()); - assert!(!prev.is_released()); - - let next = Snapshot(prev.0 ^ DELTA); - - assert!(!next.is_running()); - assert!(next.is_complete()); - assert!(next.is_released()); - - next - } - - /// Transitions the task to the canceled state. - /// - /// Returns the snapshot of the state **after** the transition **if** the - /// transition was made successfully - /// - /// # States - /// - /// - Notifed: task may be in a queue, caller must not release. - /// - Running: cannot drop. The poll handle will handle releasing. - /// - Other prior states do not require cancellation. - /// - /// If the task has been notified, then it may still be in a queue. The - /// caller must not release the task. - pub(super) fn transition_to_canceled_from_queue(&self) -> Snapshot { - let prev = Snapshot(self.val.fetch_or(CANCELLED, AcqRel)); - - assert!(!prev.is_complete()); - assert!(!prev.is_running() || prev.is_notified()); - - Snapshot(prev.0 | CANCELLED) - } - - pub(super) fn transition_to_canceled_from_list(&self) -> Option { - let mut prev = self.load(); - - loop { - if !prev.is_active() { - return None; - } - - let mut next = prev; - - // Use the running flag to signal cancellation - if prev.is_running() { - next.0 -= RUNNING; - next.0 |= NOTIFIED; - } else if prev.is_notified() { - next.0 += RUNNING; - next.0 |= NOTIFIED; - } else { - next.0 |= CANCELLED; - } - - let res = self.val.compare_exchange(prev.0, next.0, AcqRel, Acquire); - - match res { - Ok(_) if next.is_canceled() => return Some(next), - Ok(_) => return None, - Err(actual) => { - prev = Snapshot(actual); - } - } - } - } - - /// Transitions to `Released`. Called when primary task handle is - /// dropped. This is roughly a "ref decrement" operation. - /// - /// Returns a snapshot of the state **after** the transition. - pub(super) fn release_task(&self) -> Snapshot { - use crate::loom::sync::atomic; - - const DELTA: usize = RELEASED; - - let prev = Snapshot(self.val.fetch_or(DELTA, Release)); - - assert!(!prev.is_released()); - assert!(prev.is_terminal(), "state = {:?}", prev); - - let next = Snapshot(prev.0 | DELTA); - - assert!(next.is_released()); - - if next.is_final_ref() || (next.has_join_waker() && !next.is_join_interested()) { - // The final reference to the task was dropped, the caller must free the - // memory. Establish an acquire ordering. - atomic::fence(Acquire); - } - - next - } - - /// Transitions the state to `Scheduled`. - /// - /// Returns `true` if the task needs to be submitted to the pool for - /// execution - pub(super) fn transition_to_notified(&self) -> bool { - const MASK: usize = RUNNING | NOTIFIED | COMPLETE | CANCELLED; - - let prev = self.val.fetch_or(NOTIFIED, Release); - prev & MASK == 0 - } - - /// Optimistically tries to swap the state assuming the join handle is - /// __immediately__ dropped on spawn - pub(super) fn drop_join_handle_fast(&self) -> bool { - use std::sync::atomic::Ordering::Relaxed; - - // Relaxed is acceptable as if this function is called and succeeds, - // then nothing has been done w/ the join handle. - // - // The moment the join handle is used (polled), the `JOIN_WAKER` flag is - // set, at which point the CAS will fail. - // - // Given this, there is no risk if this operation is reordered. - self.val - .compare_exchange_weak( - INITIAL_STATE | JOIN_INTEREST, - INITIAL_STATE, - Release, - Relaxed, - ) - .is_ok() - } - - /// The join handle has completed by reading the output. - /// - /// Returns a snapshot of the state **after** the transition. - pub(super) fn complete_join_handle(&self) -> Snapshot { - use crate::loom::sync::atomic; - - const DELTA: usize = JOIN_INTEREST; - - let prev = Snapshot(self.val.fetch_sub(DELTA, Release)); - - assert!(prev.is_join_interested()); - - let next = Snapshot(prev.0 - DELTA); - - if !next.is_final_ref() { - return next; - } - - atomic::fence(Acquire); - - next - } - - /// The join handle is being dropped, this fails if the task has been - /// completed and the output must be dropped first then - /// `complete_join_handle` should be called. - /// - /// Returns a snapshot of the state **after** the transition. - pub(super) fn drop_join_handle_slow(&self) -> Result { - const MASK: usize = COMPLETE | CANCELLED; - - let mut prev = self.val.load(Acquire); - - loop { - // Once the complete bit is set, it is never unset. - if prev & MASK != 0 { - return Err(Snapshot(prev)); - } - - assert!(prev & JOIN_INTEREST == JOIN_INTEREST); - - let next = (prev - JOIN_INTEREST) & !JOIN_WAKER; - - let res = self.val.compare_exchange(prev, next, AcqRel, Acquire); - - match res { - Ok(_) => { - return Ok(Snapshot(next)); - } - Err(actual) => { - prev = actual; - } - } - } - } - - /// Stores the join waker. - pub(super) fn store_join_waker(&self) -> Snapshot { - use crate::loom::sync::atomic; - - const DELTA: usize = JOIN_WAKER; - - let prev = Snapshot(self.val.fetch_xor(DELTA, Release)); - - assert!(!prev.has_join_waker()); - - let next = Snapshot(prev.0 ^ DELTA); - - assert!(next.has_join_waker()); - - if next.is_complete() { - atomic::fence(Acquire); - } - - next - } - - pub(super) fn unset_waker(&self) -> Snapshot { - const MASK: usize = COMPLETE | CANCELLED; - - let mut prev = self.val.load(Acquire); - - loop { - // Once the `COMPLETE` bit is set, it is never unset - if prev & MASK != 0 { - return Snapshot(prev); - } - - assert!(Snapshot(prev).has_join_waker()); - - let next = prev - JOIN_WAKER; - - let res = self.val.compare_exchange(prev, next, AcqRel, Acquire); - - match res { - Ok(_) => return Snapshot(next), - Err(actual) => { - prev = actual; - } - } - } - } - - pub(super) fn ref_inc(&self) { - use std::process; - use std::sync::atomic::Ordering::Relaxed; - - // Using a relaxed ordering is alright here, as knowledge of the - // original reference prevents other threads from erroneously deleting - // the object. - // - // As explained in the [Boost documentation][1], Increasing the - // reference counter can always be done with memory_order_relaxed: New - // references to an object can only be formed from an existing - // reference, and passing an existing reference from one thread to - // another must already provide any required synchronization. - // - // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html) - let prev = self.val.fetch_add(WAKER_ONE, Relaxed); - - // If the reference count overflowed, abort. - if prev > isize::max_value() as usize { - process::abort(); - } - } - - /// Returns `true` if the task should be released. - pub(super) fn ref_dec(&self) -> bool { - use crate::loom::sync::atomic; - - let prev = self.val.fetch_sub(WAKER_ONE, Release); - let next = Snapshot(prev - WAKER_ONE); - - if next.is_final_ref() { - atomic::fence(Acquire); - } - - next.is_final_ref() - } -} - -impl Snapshot { - pub(super) fn is_running(self) -> bool { - self.0 & RUNNING == RUNNING - } - - pub(super) fn is_notified(self) -> bool { - self.0 & NOTIFIED == NOTIFIED - } - - pub(super) fn is_released(self) -> bool { - self.0 & RELEASED == RELEASED - } - - pub(super) fn is_complete(self) -> bool { - self.0 & COMPLETE == COMPLETE - } - - pub(super) fn is_canceled(self) -> bool { - self.0 & CANCELLED == CANCELLED - } - - /// Used during normal runtime. - pub(super) fn is_active(self) -> bool { - self.0 & (COMPLETE | CANCELLED) == 0 - } - - /// Used before dropping the task - pub(super) fn is_terminal(self) -> bool { - // When both the notified & running flags are set, the task was canceled - // after being notified, before it was run. - // - // There is a race where: - // - The task state transitions to notified - // - The global queue is shutdown - // - The waker attempts to push into the global queue and fails. - // - The waker holds the last reference to the task, thus drops it. - // - // In this scenario, the cancelled bit will never get set. - !self.is_active() || (self.is_notified() && self.is_running()) - } - - pub(super) fn is_join_interested(self) -> bool { - self.0 & JOIN_INTEREST == JOIN_INTEREST - } - - pub(super) fn has_join_waker(self) -> bool { - self.0 & JOIN_WAKER == JOIN_WAKER - } - - pub(super) fn is_final_ref(self) -> bool { - const MASK: usize = WAKER_COUNT_MASK | RELEASED | JOIN_INTEREST; - - (self.0 & MASK) == RELEASED - } -} - -impl fmt::Debug for State { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - use std::sync::atomic::Ordering::SeqCst; - - let snapshot = Snapshot(self.val.load(SeqCst)); - - fmt.debug_struct("State") - .field("snapshot", &snapshot) - .finish() - } -} - -impl fmt::Debug for Snapshot { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt.debug_struct("Snapshot") - .field("is_running", &self.is_running()) - .field("is_notified", &self.is_notified()) - .field("is_released", &self.is_released()) - .field("is_complete", &self.is_complete()) - .field("is_canceled", &self.is_canceled()) - .field("is_join_interested", &self.is_join_interested()) - .field("has_join_waker", &self.has_join_waker()) - .field("is_final_ref", &self.is_final_ref()) - .finish() - } -} diff --git a/tokio/src/task/tests/loom.rs b/tokio/src/task/tests/loom.rs deleted file mode 100644 index f1ccace24..000000000 --- a/tokio/src/task/tests/loom.rs +++ /dev/null @@ -1,277 +0,0 @@ -use crate::task; -use crate::tests::loom_schedule::LoomSchedule; - -use tokio_test::{assert_err, assert_ok}; - -use loom::future::block_on; -use loom::sync::atomic::AtomicBool; -use loom::sync::atomic::Ordering::{Acquire, Release}; -use loom::thread; -use std::future::Future; - -#[test] -fn create_drop_join_handle() { - loom::model(|| { - let (task, join_handle) = task::joinable(async { "hello" }); - - let schedule = LoomSchedule::new(); - let schedule = &mut || Some(From::from(&schedule)); - - let th = thread::spawn(move || { - drop(join_handle); - }); - - assert_none!(task.run(schedule)); - - th.join().unwrap(); - }); -} - -#[test] -fn poll_drop_handle_then_drop() { - use futures::future::poll_fn; - use std::pin::Pin; - use std::task::Poll; - - loom::model(|| { - let (task, mut join_handle) = task::joinable(async { "hello" }); - - let schedule = LoomSchedule::new(); - let schedule = &mut || Some(From::from(&schedule)); - - let th = thread::spawn(move || { - block_on(poll_fn(|cx| { - let _ = Pin::new(&mut join_handle).poll(cx); - Poll::Ready(()) - })); - }); - - assert_none!(task.run(schedule)); - - th.join().unwrap(); - }); -} - -#[test] -fn join_output() { - loom::model(|| { - let (task, join_handle) = task::joinable(async { "hello world" }); - - let schedule = LoomSchedule::new(); - let schedule = &mut || Some(From::from(&schedule)); - - let th = thread::spawn(move || { - let out = assert_ok!(block_on(join_handle)); - assert_eq!("hello world", out); - }); - - assert_none!(task.run(schedule)); - th.join().unwrap(); - }); -} - -#[test] -fn wake_by_ref() { - loom::model(|| { - let (task, join_handle) = task::joinable(gated(2, true, false)); - - let schedule = LoomSchedule::new(); - let schedule = &schedule; - schedule.push_task(task); - - let th = join_one_task(join_handle); - - work(schedule); - - assert_ok!(th.join().unwrap()); - }); -} - -#[test] -fn wake_by_val() { - loom::model(|| { - let (task, join_handle) = task::joinable(gated(2, true, true)); - - let schedule = LoomSchedule::new(); - let schedule = &schedule; - schedule.push_task(task); - - let th = join_one_task(join_handle); - - work(schedule); - - assert_ok!(th.join().unwrap()); - }); -} - -#[test] -fn release_remote() { - loom::model(|| { - let (task, join_handle) = task::joinable(gated(1, false, true)); - - let s1 = LoomSchedule::new(); - let s2 = LoomSchedule::new(); - - // Join handle - let th = join_one_task(join_handle); - - let task = match task.run(&mut || Some(From::from(&s1))) { - Some(task) => task, - None => s1.recv().expect("released!"), - }; - - assert_none!(task.run(&mut || Some(From::from(&s2)))); - assert_none!(s1.recv()); - - assert_ok!(th.join().unwrap()); - }); -} - -#[test] -fn shutdown_task_before_poll() { - loom::model(|| { - let (task, join_handle) = task::joinable::<_, LoomSchedule>(async { "hello" }); - - let th = join_one_task(join_handle); - task.shutdown(); - - assert_err!(th.join().unwrap()); - }); -} - -#[test] -fn shutdown_from_list_after_poll() { - loom::model(|| { - let (task, join_handle) = task::joinable(gated(1, false, false)); - - let s1 = LoomSchedule::new(); - - let mut list = task::OwnedList::new(); - list.insert(&task); - - // Join handle - let th = join_two_tasks(join_handle); - - match task.run(&mut || Some(From::from(&s1))) { - Some(task) => { - // always drain the list before calling shutdown on tasks - list.shutdown(); - - // The task was scheduled, drain it explicitly. - task.shutdown(); - } - None => { - list.shutdown(); - } - }; - - match s1.recv() { - Some(task) => task.shutdown(), - None => {} - } - - assert_err!(th.join().unwrap()); - }); -} - -#[test] -fn shutdown_from_queue_after_poll() { - loom::model(|| { - let (task, join_handle) = task::joinable(gated(1, false, false)); - - let s1 = LoomSchedule::new(); - - // Join handle - let th = join_two_tasks(join_handle); - - let task = match task.run(&mut || Some(From::from(&s1))) { - Some(task) => task, - None => assert_some!(s1.recv()), - }; - - task.shutdown(); - - assert_err!(th.join().unwrap()); - }); -} - -fn gated(n: usize, complete_first_poll: bool, by_val: bool) -> impl Future { - use futures::future::poll_fn; - use std::sync::Arc; - use std::task::Poll; - - let gate = Arc::new(AtomicBool::new(false)); - let mut fired = false; - - poll_fn(move |cx| { - if !fired { - for _ in 0..n { - let gate = gate.clone(); - let waker = cx.waker().clone(); - thread::spawn(move || { - gate.store(true, Release); - - if by_val { - waker.wake() - } else { - waker.wake_by_ref(); - } - }); - } - - fired = true; - - if !complete_first_poll { - return Poll::Pending; - } - } - - if gate.load(Acquire) { - Poll::Ready("hello world") - } else { - Poll::Pending - } - }) -} - -fn work(schedule: &LoomSchedule) { - while let Some(task) = schedule.recv() { - let mut task = Some(task); - - while let Some(t) = task.take() { - task = t.run(&mut || Some(From::from(schedule))); - } - } -} - -/// Spawn a thread to wait on the join handle. Uses a single task. -fn join_one_task(join_handle: T) -> loom::thread::JoinHandle { - thread::spawn(move || block_on(join_handle)) -} - -/// Spawn a thread to wait on the join handle using two tasks. First, poll the -/// join handle on the first task. If the join handle is not ready, then use a -/// second task to wait on it. -fn join_two_tasks( - join_handle: T, -) -> loom::thread::JoinHandle { - use futures::future::poll_fn; - use std::task::Poll; - - // Join handle - thread::spawn(move || { - let mut join_handle = Some(join_handle); - block_on(poll_fn(move |cx| { - use std::pin::Pin; - - let res = Pin::new(join_handle.as_mut().unwrap()).poll(cx); - - if res.is_ready() { - return res; - } - - // Yes, we are nesting - Poll::Ready(block_on(join_handle.take().unwrap())) - })) - }) -} diff --git a/tokio/src/task/tests/mod.rs b/tokio/src/task/tests/mod.rs deleted file mode 100644 index 526e1e921..000000000 --- a/tokio/src/task/tests/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -#[cfg(loom)] -mod loom; - -#[cfg(not(loom))] -mod task; diff --git a/tokio/src/task/tests/task.rs b/tokio/src/task/tests/task.rs deleted file mode 100644 index 0cff42950..000000000 --- a/tokio/src/task/tests/task.rs +++ /dev/null @@ -1,661 +0,0 @@ -use crate::sync::oneshot; -use crate::task::{self, Header}; -use crate::tests::backoff::*; -use crate::tests::mock_schedule::{mock, Mock}; -use crate::tests::track_drop::track_drop; - -use tokio_test::task::spawn; -use tokio_test::{assert_pending, assert_ready_err, assert_ready_ok}; - -use futures::future::poll_fn; -use std::sync::mpsc; - -#[test] -fn header_lte_cache_line() { - use std::mem::size_of; - - assert!(size_of::
() <= 8 * size_of::<*const ()>()); -} - -#[test] -fn create_complete_drop() { - let (tx, rx) = mpsc::channel(); - - let (task, did_drop) = track_drop(async move { - tx.send(1).unwrap(); - }); - - let (task, _) = task::joinable(task); - - let mock = mock().bind(&task).release_local(); - let mock = &mut || Some(From::from(&mock)); - - // Nothing is returned - assert!(task.run(mock).is_none()); - - // The message was sent - assert!(rx.try_recv().is_ok()); - - // The future & output were dropped. - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -fn create_yield_complete_drop() { - let (tx, rx) = mpsc::channel(); - - let (task, did_drop) = track_drop(async move { - backoff(1).await; - tx.send(1).unwrap(); - }); - - let (task, _) = task::joinable(task); - - let mock = mock().bind(&task).release_local(); - let mock = || Some(From::from(&mock)); - - // Task is returned - let task = assert_some!(task.run(mock)); - - // The future was **not** dropped. - assert!(!did_drop.did_drop_future()); - - assert_none!(task.run(mock)); - - // The message was sent - assert!(rx.try_recv().is_ok()); - - // The future was dropped. - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -fn create_clone_yield_complete_drop() { - let (tx, rx) = mpsc::channel(); - - let (task, did_drop) = track_drop(async move { - backoff_clone(1).await; - tx.send(1).unwrap(); - }); - - let (task, _) = task::joinable(task); - - let mock = mock().bind(&task).release_local(); - let mock = || Some(From::from(&mock)); - - // Task is returned - let task = assert_some!(task.run(mock)); - - // The future was **not** dropped. - assert!(!did_drop.did_drop_future()); - - assert_none!(task.run(mock)); - - // The message was sent - assert!(rx.try_recv().is_ok()); - - // The future was dropped. - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -fn create_wake_drop() { - let (tx, rx) = oneshot::channel(); - - let (task, did_drop) = track_drop(async move { rx.await }); - - let (task, _) = task::joinable(task); - - let mock = mock().bind(&task).schedule().release_local(); - - assert_none!(task.run(&mut || Some(From::from(&mock)))); - assert_none!(mock.next_pending_run()); - - // The future was **not** dropped. - assert!(!did_drop.did_drop_future()); - - tx.send("hello").unwrap(); - - let task = assert_some!(mock.next_pending_run()); - - assert_none!(task.run(&mut || Some(From::from(&mock)))); - - // The future was dropped. - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -fn notify_complete() { - use std::task::Poll::Ready; - - let (task, did_drop) = track_drop(async move { - poll_fn(|cx| { - cx.waker().wake_by_ref(); - Ready(()) - }) - .await; - }); - - let (task, _) = task::joinable(task); - - let mock = mock().bind(&task).release_local(); - let mock = &mut || Some(From::from(&mock)); - - assert_none!(task.run(mock)); - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -fn complete_on_second_schedule_obj() { - let (tx, rx) = mpsc::channel(); - - let (task, did_drop) = track_drop(async move { - backoff(1).await; - tx.send(1).unwrap(); - }); - - let (task, _) = task::joinable(task); - - let mock1 = mock(); - let mock2 = mock().bind(&task).release(); - - // Task is returned - let task = assert_some!(task.run(&mut || Some(From::from(&mock2)))); - - assert_none!(task.run(&mut || Some(From::from(&mock1)))); - - // The message was sent - assert!(rx.try_recv().is_ok()); - - // The future was dropped. - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); - - let _ = assert_some!(mock2.next_pending_drop()); -} - -#[test] -fn join_task_immediate_drop_handle() { - let (task, did_drop) = track_drop(async move { "hello".to_string() }); - - let (task, _) = task::joinable(task); - - let mock = mock().bind(&task).release_local(); - - assert!(task.run(&mut || Some(From::from(&mock))).is_none()); - - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -fn join_task_immediate_complete_1() { - let (task, did_drop) = track_drop(async move { "hello".to_string() }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - let mock = mock().bind(&task).release_local(); - - assert!(task.run(&mut || Some(From::from(&mock))).is_none()); - - assert!(did_drop.did_drop_future()); - assert!(!did_drop.did_drop_output()); - assert!(!handle.is_woken()); - - let out = assert_ready_ok!(handle.poll()); - assert_eq!(out.get_ref(), "hello"); - - drop(out); - - assert!(did_drop.did_drop_output()); -} - -#[test] -fn join_task_immediate_complete_2() { - let (task, did_drop) = track_drop(async move { "hello".to_string() }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - let mock = mock().bind(&task).release_local(); - - assert_pending!(handle.poll()); - - assert!(task.run(&mut || Some(From::from(&mock))).is_none()); - - assert!(did_drop.did_drop_future()); - assert!(!did_drop.did_drop_output()); - assert!(handle.is_woken()); - - let out = assert_ready_ok!(handle.poll()); - assert_eq!(out.get_ref(), "hello"); - - drop(out); - - assert!(did_drop.did_drop_output()); -} - -#[test] -fn join_task_complete_later() { - let (task, did_drop) = track_drop(async move { - backoff(1).await; - "hello".to_string() - }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(async { handle.await }); - - let mock = mock().bind(&task).release_local(); - - let task = assert_some!(task.run(&mut || Some(From::from(&mock)))); - - assert!(!did_drop.did_drop_future()); - assert!(!did_drop.did_drop_output()); - - assert_pending!(handle.poll()); - - assert_none!(task.run(&mut || Some(From::from(&mock)))); - assert!(handle.is_woken()); - - let out = assert_ready_ok!(handle.poll()); - assert_eq!(out.get_ref(), "hello"); - - drop(out); - - assert!(did_drop.did_drop_output()); - - assert_eq!(1, handle.waker_ref_count()); -} - -#[test] -fn drop_join_after_poll() { - let (task, did_drop) = track_drop(async move { - backoff(1).await; - "hello".to_string() - }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(async { handle.await }); - - let mock = mock().bind(&task).release_local(); - - assert_pending!(handle.poll()); - drop(handle); - - let task = assert_some!(task.run(&mut || Some(From::from(&mock)))); - - assert!(!did_drop.did_drop_future()); - assert!(!did_drop.did_drop_output()); - - assert_none!(task.run(&mut || Some(From::from(&mock)))); - - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -fn join_handle_change_task_complete() { - use std::future::Future; - use std::pin::Pin; - - let (task, did_drop) = track_drop(async move { - backoff(1).await; - "hello".to_string() - }); - - let (task, mut handle) = task::joinable(task); - let mut t1 = spawn(poll_fn(|cx| Pin::new(&mut handle).poll(cx))); - - let mock = mock().bind(&task).release_local(); - - assert_pending!(t1.poll()); - drop(t1); - - let task = assert_some!(task.run(&mut || Some(From::from(&mock)))); - - let mut t2 = spawn(poll_fn(|cx| Pin::new(&mut handle).poll(cx))); - assert_pending!(t2.poll()); - - assert!(!did_drop.did_drop_future()); - assert!(!did_drop.did_drop_output()); - - assert_none!(task.run(&mut || Some(From::from(&mock)))); - - assert!(t2.is_woken()); - - let out = assert_ready_ok!(t2.poll()); - assert_eq!(out.get_ref(), "hello"); - - drop(out); - - assert!(did_drop.did_drop_output()); - - assert_eq!(1, t2.waker_ref_count()); -} - -#[test] -fn drop_handle_after_complete() { - let (task, did_drop) = track_drop(async move { "hello".to_string() }); - - let (task, handle) = task::joinable(task); - - let mock = mock().bind(&task).release_local(); - - assert!(task.run(&mut || Some(From::from(&mock))).is_none()); - - assert!(did_drop.did_drop_future()); - assert!(!did_drop.did_drop_output()); - - drop(handle); - - assert!(did_drop.did_drop_output()); -} - -#[test] -fn non_initial_task_state_drop_join_handle_without_polling() { - let (tx, rx) = oneshot::channel::<()>(); - - let (task, did_drop) = track_drop(async move { - rx.await.unwrap(); - "hello".to_string() - }); - - let (task, handle) = task::joinable(task); - - let mock = mock().bind(&task).schedule().release_local(); - - assert_none!(task.run(&mut || Some(From::from(&mock)))); - - drop(handle); - - assert!(!did_drop.did_drop_future()); - assert!(!did_drop.did_drop_output()); - - tx.send(()).unwrap(); - let task = assert_some!(mock.next_pending_run()); - - assert!(task.run(&mut || Some(From::from(&mock))).is_none()); - - assert!(did_drop.did_drop_future()); - assert!(did_drop.did_drop_output()); -} - -#[test] -#[cfg(not(miri))] -fn task_panic_background() { - let (task, did_drop) = track_drop(async move { - if true { - panic!() - } - "hello" - }); - - let (task, _) = task::joinable(task); - - let mock = mock().bind(&task).release_local(); - - assert!(task.run(&mut || Some(From::from(&mock))).is_none()); - - assert!(did_drop.did_drop_future()); -} - -#[test] -#[cfg(not(miri))] -fn task_panic_join() { - let (task, did_drop) = track_drop(async move { - if true { - panic!() - } - "hello" - }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - let mock = mock().bind(&task).release_local(); - - assert_pending!(handle.poll()); - - assert!(task.run(&mut || Some(From::from(&mock))).is_none()); - assert!(did_drop.did_drop_future()); - assert!(handle.is_woken()); - - assert_ready_err!(handle.poll()); -} - -#[test] -fn complete_second_schedule_obj_before_join() { - let (tx, rx) = oneshot::channel(); - - let (task, did_drop) = track_drop(async move { rx.await.unwrap() }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - let mock1 = mock(); - let mock2 = mock().bind(&task).schedule().release(); - - assert_pending!(handle.poll()); - - assert_none!(task.run(&mut || Some(From::from(&mock2)))); - - tx.send("hello").unwrap(); - - let task = assert_some!(mock2.next_pending_run()); - assert_none!(task.run(&mut || Some(From::from(&mock1)))); - assert!(did_drop.did_drop_future()); - - // The join handle was notified - assert!(handle.is_woken()); - - // Drop the task - let _ = assert_some!(mock2.next_pending_drop()); - - // Get the output - let out = assert_ready_ok!(handle.poll()); - assert_eq!(*out.get_ref(), "hello"); -} - -#[test] -fn complete_second_schedule_obj_after_join() { - let (tx, rx) = oneshot::channel(); - - let (task, did_drop) = track_drop(async move { rx.await.unwrap() }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - let mock1 = mock(); - let mock2 = mock().bind(&task).schedule().release(); - - assert_pending!(handle.poll()); - - assert_none!(task.run(&mut || Some(From::from(&mock2)))); - - tx.send("hello").unwrap(); - - let task = assert_some!(mock2.next_pending_run()); - assert_none!(task.run(&mut || Some(From::from(&mock1)))); - assert!(did_drop.did_drop_future()); - - // The join handle was notified - assert!(handle.is_woken()); - - // Get the output - let out = assert_ready_ok!(handle.poll()); - assert_eq!(*out.get_ref(), "hello"); - - // Drop the task - let _ = assert_some!(mock2.next_pending_drop()); - - assert_eq!(1, handle.waker_ref_count()); -} - -#[test] -fn shutdown_from_list_before_notified() { - let (tx, rx) = oneshot::channel::<()>(); - let mut list = task::OwnedList::new(); - - let (task, did_drop) = track_drop(async move { rx.await }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - list.insert(&task); - - let mock = mock().bind(&task).release(); - - assert_pending!(handle.poll()); - assert_none!(task.run(&mut || Some(From::from(&mock)))); - - list.shutdown(); - assert!(did_drop.did_drop_future()); - - assert!(handle.is_woken()); - - let task = assert_some!(mock.next_pending_drop()); - drop(task); - - assert_ready_err!(handle.poll()); - - drop(tx); -} - -#[test] -fn shutdown_from_list_after_notified() { - let (tx, rx) = oneshot::channel::<()>(); - let mut list = task::OwnedList::new(); - - let (task, did_drop) = track_drop(async move { rx.await }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - list.insert(&task); - - let mock = mock().bind(&task).schedule().release(); - - assert_pending!(handle.poll()); - assert_none!(task.run(&mut || Some(From::from(&mock)))); - - tx.send(()).unwrap(); - - let task = assert_some!(mock.next_pending_run()); - - list.shutdown(); - - assert_none!(mock.next_pending_drop()); - - assert_none!(task.run(&mut || Some(From::from(&mock)))); - assert!(did_drop.did_drop_future()); - assert!(handle.is_woken()); - - let task = assert_some!(mock.next_pending_drop()); - drop(task); - - assert_ready_err!(handle.poll()); -} - -#[test] -fn shutdown_from_list_after_complete() { - let mut list = task::OwnedList::new(); - - let (task, did_drop) = track_drop(async move { - backoff(1).await; - "hello" - }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - list.insert(&task); - - let m1 = mock().bind(&task).release(); - let m2 = mock(); - - assert_pending!(handle.poll()); - let task = assert_some!(task.run(&mut || Some(From::from(&m1)))); - assert_none!(task.run(&mut || Some(From::from(&m2)))); - assert!(did_drop.did_drop_future()); - assert!(handle.is_woken()); - - list.shutdown(); - - let task = assert_some!(m1.next_pending_drop()); - drop(task); - - let out = assert_ready_ok!(handle.poll()); - assert_eq!(*out.get_ref(), "hello"); -} - -#[test] -fn shutdown_from_task_before_notified() { - let (tx, rx) = oneshot::channel::<()>(); - - let (task, did_drop) = track_drop(async move { rx.await }); - - let (task, handle) = task::joinable::<_, Mock>(task); - let mut handle = spawn(handle); - - assert_pending!(handle.poll()); - - task.shutdown(); - assert!(did_drop.did_drop_future()); - assert!(handle.is_woken()); - - assert_ready_err!(handle.poll()); - - drop(tx); -} - -#[test] -fn shutdown_from_task_after_notified() { - let (tx, rx) = oneshot::channel::<()>(); - - let (task, did_drop) = track_drop(async move { rx.await }); - - let (task, handle) = task::joinable(task); - let mut handle = spawn(handle); - - let mock = mock().bind(&task).schedule().release(); - - assert_pending!(handle.poll()); - assert_none!(task.run(&mut || Some(From::from(&mock)))); - - tx.send(()).unwrap(); - - let task = assert_some!(mock.next_pending_run()); - - task.shutdown(); - assert!(did_drop.did_drop_future()); - assert!(handle.is_woken()); - - let task = assert_some!(mock.next_pending_drop()); - drop(task); - - assert_ready_err!(handle.poll()); -} - -#[test] -fn waker_ref_will_wake_clone() { - use std::task::Poll::Ready; - - let (task, handle) = task::joinable(poll_fn(|cx| { - let waker = cx.waker().clone(); - assert!(cx.waker().will_wake(&waker)); - Ready(()) - })); - let mut handle = spawn(handle); - - let mock = mock().bind(&task).release_local(); - let mock = &mut || Some(From::from(&mock)); - - assert_none!(task.run(mock)); - assert_ready_ok!(handle.poll()); -} diff --git a/tokio/src/tests/backoff.rs b/tokio/src/tests/backoff.rs deleted file mode 100644 index 358ab2dad..000000000 --- a/tokio/src/tests/backoff.rs +++ /dev/null @@ -1,32 +0,0 @@ -use std::future::Future; -use std::pin::Pin; -use std::task::{Context, Poll}; - -pub(crate) struct Backoff(usize, bool); - -pub(crate) fn backoff(n: usize) -> impl Future { - Backoff(n, false) -} - -/// Back off, but clone the waker each time -pub(crate) fn backoff_clone(n: usize) -> impl Future { - Backoff(n, true) -} - -impl Future for Backoff { - type Output = (); - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - if self.0 == 0 { - return Poll::Ready(()); - } - - self.0 -= 1; - if self.1 { - cx.waker().clone().wake(); - } else { - cx.waker().wake_by_ref(); - } - Poll::Pending - } -} diff --git a/tokio/src/tests/loom_schedule.rs b/tokio/src/tests/loom_schedule.rs deleted file mode 100644 index 15ec903b5..000000000 --- a/tokio/src/tests/loom_schedule.rs +++ /dev/null @@ -1,53 +0,0 @@ -use crate::task::{Schedule, ScheduleSendOnly, Task}; - -use loom::sync::Notify; -use std::collections::VecDeque; -use std::sync::Mutex; - -pub(crate) struct LoomSchedule { - notify: Notify, - pending: Mutex>>>, -} - -impl LoomSchedule { - pub(crate) fn new() -> LoomSchedule { - LoomSchedule { - notify: Notify::new(), - pending: Mutex::new(VecDeque::new()), - } - } - - pub(crate) fn push_task(&self, task: Task) { - self.schedule(task); - } - - pub(crate) fn recv(&self) -> Option> { - loop { - if let Some(task) = self.pending.lock().unwrap().pop_front() { - return task; - } - - self.notify.wait(); - } - } -} - -impl Schedule for LoomSchedule { - fn bind(&self, _task: &Task) {} - - fn release(&self, task: Task) { - self.release_local(&task); - } - - fn release_local(&self, _task: &Task) { - self.pending.lock().unwrap().push_back(None); - self.notify.notify(); - } - - fn schedule(&self, task: Task) { - self.pending.lock().unwrap().push_back(Some(task)); - self.notify.notify(); - } -} - -impl ScheduleSendOnly for LoomSchedule {} diff --git a/tokio/src/tests/mock_schedule.rs b/tokio/src/tests/mock_schedule.rs deleted file mode 100644 index a99415646..000000000 --- a/tokio/src/tests/mock_schedule.rs +++ /dev/null @@ -1,134 +0,0 @@ -#![allow(warnings)] -use crate::task::{Header, Schedule, ScheduleSendOnly, Task}; - -use std::collections::VecDeque; -use std::sync::Mutex; -use std::thread; - -pub(crate) struct Mock { - inner: Mutex, -} - -pub(crate) struct Noop; -pub(crate) static NOOP_SCHEDULE: Noop = Noop; - -struct Inner { - calls: VecDeque, - pending_run: VecDeque>, - pending_drop: VecDeque>, -} - -unsafe impl Send for Inner {} -unsafe impl Sync for Inner {} - -#[derive(Debug, Eq, PartialEq)] -enum Call { - Bind(*const Header), - Release, - ReleaseLocal, - Schedule, -} - -pub(crate) fn mock() -> Mock { - Mock { - inner: Mutex::new(Inner { - calls: VecDeque::new(), - pending_run: VecDeque::new(), - pending_drop: VecDeque::new(), - }), - } -} - -impl Mock { - pub(crate) fn bind(self, task: &Task) -> Self { - self.push(Call::Bind(task.header() as *const _)); - self - } - - pub(crate) fn release(self) -> Self { - self.push(Call::Release); - self - } - - pub(crate) fn release_local(self) -> Self { - self.push(Call::ReleaseLocal); - self - } - - pub(crate) fn schedule(self) -> Self { - self.push(Call::Schedule); - self - } - - pub(crate) fn next_pending_run(&self) -> Option> { - self.inner.lock().unwrap().pending_run.pop_front() - } - - pub(crate) fn next_pending_drop(&self) -> Option> { - self.inner.lock().unwrap().pending_drop.pop_front() - } - - fn push(&self, call: Call) { - self.inner.lock().unwrap().calls.push_back(call); - } - - fn next(&self, name: &str) -> Call { - self.inner - .lock() - .unwrap() - .calls - .pop_front() - .expect(&format!("received `{}`, but none expected", name)) - } -} - -impl Schedule for Mock { - fn bind(&self, task: &Task) { - match self.next("bind") { - Call::Bind(ptr) => { - assert!(ptr.eq(&(task.header() as *const _))); - } - call => panic!("expected `Bind`, was {:?}", call), - } - } - - fn release(&self, task: Task) { - match self.next("release") { - Call::Release => { - self.inner.lock().unwrap().pending_drop.push_back(task); - } - call => panic!("expected `Release`, was {:?}", call), - } - } - - fn release_local(&self, _task: &Task) { - assert_eq!(Call::ReleaseLocal, self.next("release_local")); - } - - fn schedule(&self, task: Task) { - self.inner.lock().unwrap().pending_run.push_back(task); - assert_eq!(Call::Schedule, self.next("schedule")); - } -} - -impl ScheduleSendOnly for Mock {} - -impl Drop for Mock { - fn drop(&mut self) { - if !thread::panicking() { - assert!(self.inner.lock().unwrap().calls.is_empty()); - } - } -} - -impl Schedule for Noop { - fn bind(&self, _task: &Task) {} - - fn release(&self, _task: Task) {} - - fn release_local(&self, _task: &Task) {} - - fn schedule(&self, _task: Task) {} -} - -impl ScheduleSendOnly for Noop {} diff --git a/tokio/src/tests/mod.rs b/tokio/src/tests/mod.rs deleted file mode 100644 index b326561b6..000000000 --- a/tokio/src/tests/mod.rs +++ /dev/null @@ -1,10 +0,0 @@ -#[cfg(not(loom))] -pub(crate) mod backoff; - -#[cfg(loom)] -pub(crate) mod loom_schedule; - -pub(crate) mod mock_schedule; - -#[cfg(not(loom))] -pub(crate) mod track_drop; diff --git a/tokio/src/tests/track_drop.rs b/tokio/src/tests/track_drop.rs deleted file mode 100644 index c3ded845f..000000000 --- a/tokio/src/tests/track_drop.rs +++ /dev/null @@ -1,57 +0,0 @@ -use std::future::Future; -use std::pin::Pin; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering::SeqCst; -use std::sync::Arc; -use std::task::{Context, Poll}; - -#[derive(Debug)] -pub(crate) struct TrackDrop(T, Arc); - -#[derive(Debug)] -pub(crate) struct DidDrop(Arc, Arc); - -pub(crate) fn track_drop( - future: T, -) -> (impl Future>, DidDrop) { - let did_drop_future = Arc::new(AtomicBool::new(false)); - let did_drop_output = Arc::new(AtomicBool::new(false)); - let did_drop = DidDrop(did_drop_future.clone(), did_drop_output.clone()); - - let future = async move { TrackDrop(future.await, did_drop_output) }; - - let future = TrackDrop(future, did_drop_future); - - (future, did_drop) -} - -impl TrackDrop { - pub(crate) fn get_ref(&self) -> &T { - &self.0 - } -} - -impl Future for TrackDrop { - type Output = T::Output; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let me = unsafe { Pin::map_unchecked_mut(self, |x| &mut x.0) }; - me.poll(cx) - } -} - -impl Drop for TrackDrop { - fn drop(&mut self) { - self.1.store(true, SeqCst); - } -} - -impl DidDrop { - pub(crate) fn did_drop_future(&self) -> bool { - self.0.load(SeqCst) - } - - pub(crate) fn did_drop_output(&self) -> bool { - self.1.load(SeqCst) - } -} diff --git a/tokio/src/util/linked_list.rs b/tokio/src/util/linked_list.rs index 57540c4a4..07c25fe98 100644 --- a/tokio/src/util/linked_list.rs +++ b/tokio/src/util/linked_list.rs @@ -4,6 +4,7 @@ //! structure's APIs are `unsafe` as they require the caller to ensure the //! specified node is actually contained by the list. +use core::mem::ManuallyDrop; use core::ptr::NonNull; /// An intrusive linked list. @@ -41,10 +42,8 @@ pub(crate) unsafe trait Link { /// Node type type Target; - /// Convert the handle to a raw pointer - /// - /// Consumes ownership of the handle. - fn to_raw(handle: Self::Handle) -> NonNull; + /// Convert the handle to a raw pointer without consuming the handle + fn as_raw(handle: &Self::Handle) -> NonNull; /// Convert the raw pointer to a handle unsafe fn from_raw(ptr: NonNull) -> Self::Handle; @@ -79,7 +78,9 @@ impl LinkedList { /// Adds an element first in the list. pub(crate) fn push_front(&mut self, val: T::Handle) { - let ptr = T::to_raw(val); + // The value should not be dropped, it is being inserted into the list + let val = ManuallyDrop::new(val); + let ptr = T::as_raw(&*val); unsafe { T::pointers(ptr).as_mut().next = self.head; @@ -133,13 +134,13 @@ impl LinkedList { /// /// The caller **must** ensure that `node` is currently contained by /// `self` or not contained by any other list. - pub(crate) unsafe fn remove(&mut self, node: NonNull) -> bool { + pub(crate) unsafe fn remove(&mut self, node: NonNull) -> Option { if let Some(prev) = T::pointers(node).as_ref().prev { debug_assert_eq!(T::pointers(prev).as_ref().next, Some(node)); T::pointers(prev).as_mut().next = T::pointers(node).as_ref().next; } else { if self.head != Some(node) { - return false; + return None; } self.head = T::pointers(node).as_ref().next; @@ -151,7 +152,7 @@ impl LinkedList { } else { // This might be the last item in the list if self.tail != Some(node) { - return false; + return None; } self.tail = T::pointers(node).as_ref().prev; @@ -160,7 +161,40 @@ impl LinkedList { T::pointers(node).as_mut().next = None; T::pointers(node).as_mut().prev = None; - true + Some(T::from_raw(node)) + } +} + +// ===== impl Iter ===== + +cfg_rt_threaded! { + use core::marker::PhantomData; + + pub(crate) struct Iter<'a, T: Link> { + curr: Option>, + _p: PhantomData<&'a T>, + } + + impl LinkedList { + pub(crate) fn iter(&self) -> Iter<'_, T> { + Iter { + curr: self.head, + _p: PhantomData, + } + } + } + + impl<'a, T: Link> Iterator for Iter<'a, T> { + type Item = &'a T::Target; + + fn next(&mut self) -> Option<&'a T::Target> { + let curr = self.curr?; + // safety: the pointer references data contained by the list + self.curr = unsafe { T::pointers(curr).as_ref() }.next; + + // safety: the value is still owned by the linked list. + Some(unsafe { &*curr.as_ptr() }) + } } } @@ -192,7 +226,7 @@ mod tests { type Handle = Pin<&'a Entry>; type Target = Entry; - fn to_raw(handle: Pin<&'_ Entry>) -> NonNull { + fn as_raw(handle: &Pin<&'_ Entry>) -> NonNull { NonNull::from(handle.get_ref()) } @@ -299,22 +333,22 @@ mod tests { let mut list = LinkedList::new(); push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); - assert!(list.remove(ptr(&a))); + assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); // `a` should be no longer there and can't be removed twice - assert!(!list.remove(ptr(&a))); + assert!(list.remove(ptr(&a)).is_none()); assert!(!list.is_empty()); - assert!(list.remove(ptr(&b))); + assert!(list.remove(ptr(&b)).is_some()); assert_clean!(b); // `b` should be no longer there and can't be removed twice - assert!(!list.remove(ptr(&b))); + assert!(list.remove(ptr(&b)).is_none()); assert!(!list.is_empty()); - assert!(list.remove(ptr(&c))); + assert!(list.remove(ptr(&c)).is_some()); assert_clean!(c); // `b` should be no longer there and can't be removed twice - assert!(!list.remove(ptr(&c))); + assert!(list.remove(ptr(&c)).is_none()); assert!(list.is_empty()); } @@ -324,7 +358,7 @@ mod tests { push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); - assert!(list.remove(ptr(&a))); + assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); assert_ptr_eq!(b, list.head); @@ -341,7 +375,7 @@ mod tests { push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); - assert!(list.remove(ptr(&b))); + assert!(list.remove(ptr(&b)).is_some()); assert_clean!(b); assert_ptr_eq!(c, a.pointers.next); @@ -358,7 +392,7 @@ mod tests { push_all(&mut list, &[c.as_ref(), b.as_ref(), a.as_ref()]); - assert!(list.remove(ptr(&c))); + assert!(list.remove(ptr(&c)).is_some()); assert_clean!(c); assert!(b.pointers.next.is_none()); @@ -374,12 +408,12 @@ mod tests { push_all(&mut list, &[b.as_ref(), a.as_ref()]); - assert!(list.remove(ptr(&a))); + assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); // a should be no longer there and can't be removed twice - assert!(!list.remove(ptr(&a))); + assert!(list.remove(ptr(&a)).is_none()); assert_ptr_eq!(b, list.head); assert_ptr_eq!(b, list.tail); @@ -397,7 +431,7 @@ mod tests { push_all(&mut list, &[b.as_ref(), a.as_ref()]); - assert!(list.remove(ptr(&b))); + assert!(list.remove(ptr(&b)).is_some()); assert_clean!(b); @@ -417,7 +451,7 @@ mod tests { push_all(&mut list, &[a.as_ref()]); - assert!(list.remove(ptr(&a))); + assert!(list.remove(ptr(&a)).is_some()); assert_clean!(a); assert!(list.head.is_none()); @@ -433,10 +467,28 @@ mod tests { list.push_front(b.as_ref()); list.push_front(a.as_ref()); - assert!(!list.remove(ptr(&c))); + assert!(list.remove(ptr(&c)).is_none()); } } + #[test] + fn iter() { + let a = entry(5); + let b = entry(7); + + let mut list = LinkedList::<&Entry>::new(); + + assert_eq!(0, list.iter().count()); + + list.push_front(a.as_ref()); + list.push_front(b.as_ref()); + + let mut i = list.iter(); + assert_eq!(7, i.next().unwrap().val); + assert_eq!(5, i.next().unwrap().val); + assert!(i.next().is_none()); + } + proptest::proptest! { #[test] fn fuzz_linked_list(ops: Vec) { @@ -493,10 +545,11 @@ mod tests { } let idx = n % reference.len(); - let v = reference.remove(idx).unwrap(); + let expect = reference.remove(idx).unwrap(); unsafe { - assert!(ll.remove(ptr(&entries[v as usize]))); + let entry = ll.remove(ptr(&entries[expect as usize])).unwrap(); + assert_eq!(expect, entry.val); } } } diff --git a/tokio/src/util/mod.rs b/tokio/src/util/mod.rs index 2761f7252..c2f572f1b 100644 --- a/tokio/src/util/mod.rs +++ b/tokio/src/util/mod.rs @@ -3,17 +3,18 @@ cfg_io_driver! { pub(crate) mod slab; } -cfg_sync! { - pub(crate) mod linked_list; -} +#[cfg(any(feature = "sync", feature = "rt-core"))] +pub(crate) mod linked_list; #[cfg(any(feature = "rt-threaded", feature = "macros", feature = "stream"))] mod rand; -cfg_rt_threaded! { - mod pad; - pub(crate) use pad::CachePadded; +cfg_rt_core! { + mod wake; + pub(crate) use wake::{waker_ref, Wake}; +} +cfg_rt_threaded! { pub(crate) use rand::FastRand; mod try_lock; diff --git a/tokio/src/util/try_lock.rs b/tokio/src/util/try_lock.rs index a42e750b0..8b0edb4a8 100644 --- a/tokio/src/util/try_lock.rs +++ b/tokio/src/util/try_lock.rs @@ -20,13 +20,26 @@ unsafe impl Sync for TryLock {} unsafe impl Sync for LockGuard<'_, T> {} -impl TryLock { - /// Create a new `TryLock` - pub(crate) fn new(data: T) -> TryLock { +macro_rules! new { + ($data:ident) => { TryLock { locked: AtomicBool::new(false), - data: UnsafeCell::new(data), + data: UnsafeCell::new($data), } + }; +} + +impl TryLock { + #[cfg(not(loom))] + /// Create a new `TryLock` + pub(crate) const fn new(data: T) -> TryLock { + new!(data) + } + + #[cfg(loom)] + /// Create a new `TryLock` + pub(crate) fn new(data: T) -> TryLock { + new!(data) } /// Attempt to acquire lock diff --git a/tokio/src/util/wake.rs b/tokio/src/util/wake.rs new file mode 100644 index 000000000..e49f1e895 --- /dev/null +++ b/tokio/src/util/wake.rs @@ -0,0 +1,83 @@ +use std::marker::PhantomData; +use std::mem::ManuallyDrop; +use std::ops::Deref; +use std::sync::Arc; +use std::task::{RawWaker, RawWakerVTable, Waker}; + +/// Simplfied waking interface based on Arcs +pub(crate) trait Wake: Send + Sync { + /// Wake by value + fn wake(self: Arc); + + /// Wake by reference + fn wake_by_ref(arc_self: &Arc); +} + +/// A `Waker` that is only valid for a given lifetime. +#[derive(Debug)] +pub(crate) struct WakerRef<'a> { + waker: ManuallyDrop, + _p: PhantomData<&'a ()>, +} + +impl Deref for WakerRef<'_> { + type Target = Waker; + + fn deref(&self) -> &Waker { + &self.waker + } +} + +/// Creates a reference to a `Waker` from a reference to `Arc`. +pub(crate) fn waker_ref(wake: &Arc) -> WakerRef<'_> { + let ptr = &**wake as *const _ as *const (); + + let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::())) }; + + WakerRef { + waker: ManuallyDrop::new(waker), + _p: PhantomData, + } +} + +fn waker_vtable() -> &'static RawWakerVTable { + &RawWakerVTable::new( + clone_arc_raw::, + wake_arc_raw::, + wake_by_ref_arc_raw::, + drop_arc_raw::, + ) +} + +unsafe fn inc_ref_count(data: *const ()) { + // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop + let arc = ManuallyDrop::new(Arc::::from_raw(data as *const T)); + + // Now increase refcount, but don't drop new refcount either + let arc_clone: ManuallyDrop<_> = arc.clone(); + + // Drop explicitly to avoid clippy warnings + drop(arc); + drop(arc_clone); +} + +unsafe fn clone_arc_raw(data: *const ()) -> RawWaker { + inc_ref_count::(data); + RawWaker::new(data, waker_vtable::()) +} + +unsafe fn wake_arc_raw(data: *const ()) { + let arc: Arc = Arc::from_raw(data as *const T); + Wake::wake(arc); +} + +// used by `waker_ref` +unsafe fn wake_by_ref_arc_raw(data: *const ()) { + // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop + let arc = ManuallyDrop::new(Arc::::from_raw(data as *const T)); + Wake::wake_by_ref(&arc); +} + +unsafe fn drop_arc_raw(data: *const ()) { + drop(Arc::::from_raw(data as *const T)) +} diff --git a/tokio/tests/rt_common.rs b/tokio/tests/rt_common.rs index 64dd36804..52d33a51d 100644 --- a/tokio/tests/rt_common.rs +++ b/tokio/tests/rt_common.rs @@ -307,7 +307,7 @@ rt_test! { } #[test] - fn spawn_from_other_thread() { + fn spawn_from_other_thread_idle() { let mut rt = rt(); let handle = rt.handle().clone(); @@ -326,6 +326,31 @@ rt_test! { }); } + #[test] + fn spawn_from_other_thread_under_load() { + let mut rt = rt(); + let handle = rt.handle().clone(); + + let (tx, rx) = oneshot::channel(); + + thread::spawn(move || { + handle.spawn(async move { + assert_ok!(tx.send(())); + }); + }); + + rt.block_on(async move { + // Spin hard + tokio::spawn(async { + loop { + yield_once().await; + } + }); + + assert_ok!(rx.await); + }); + } + #[test] fn delay_at_root() { let mut rt = rt(); @@ -680,7 +705,7 @@ rt_test! { fn io_notify_while_shutting_down() { use std::net::Ipv6Addr; - for _ in 1..100 { + for _ in 1..10 { let mut runtime = rt(); runtime.block_on(async { @@ -768,66 +793,61 @@ rt_test! { tx.send(()).unwrap(); } - mod local_set { - use tokio::task; - use super::*; + #[test] + fn local_set_block_on_socket() { + let mut rt = rt(); + let local = task::LocalSet::new(); - #[test] - fn block_on_socket() { - let mut rt = rt(); - let local = task::LocalSet::new(); + local.block_on(&mut rt, async move { + let (tx, rx) = oneshot::channel(); - local.block_on(&mut rt, async move { - let (tx, rx) = oneshot::channel(); + let mut listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); - let mut listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - - task::spawn_local(async move { - let _ = listener.accept().await; - tx.send(()).unwrap(); - }); - - TcpStream::connect(&addr).await.unwrap(); - rx.await.unwrap(); - }); - } - - #[test] - fn client_server_block_on() { - let mut rt = rt(); - let (tx, rx) = mpsc::channel(); - - let local = task::LocalSet::new(); - - local.block_on(&mut rt, async move { client_server_local(tx).await }); - - assert_ok!(rx.try_recv()); - assert_err!(rx.try_recv()); - } - - async fn client_server_local(tx: mpsc::Sender<()>) { - let mut server = assert_ok!(TcpListener::bind("127.0.0.1:0").await); - - // Get the assigned address - let addr = assert_ok!(server.local_addr()); - - // Spawn the server task::spawn_local(async move { - // Accept a socket - let (mut socket, _) = server.accept().await.unwrap(); - - // Write some data - socket.write_all(b"hello").await.unwrap(); + let _ = listener.accept().await; + tx.send(()).unwrap(); }); - let mut client = TcpStream::connect(&addr).await.unwrap(); + TcpStream::connect(&addr).await.unwrap(); + rx.await.unwrap(); + }); + } - let mut buf = vec![]; - client.read_to_end(&mut buf).await.unwrap(); + #[test] + fn local_set_client_server_block_on() { + let mut rt = rt(); + let (tx, rx) = mpsc::channel(); - assert_eq!(buf, b"hello"); - tx.send(()).unwrap(); - } + let local = task::LocalSet::new(); + + local.block_on(&mut rt, async move { client_server_local(tx).await }); + + assert_ok!(rx.try_recv()); + assert_err!(rx.try_recv()); + } + + async fn client_server_local(tx: mpsc::Sender<()>) { + let mut server = assert_ok!(TcpListener::bind("127.0.0.1:0").await); + + // Get the assigned address + let addr = assert_ok!(server.local_addr()); + + // Spawn the server + task::spawn_local(async move { + // Accept a socket + let (mut socket, _) = server.accept().await.unwrap(); + + // Write some data + socket.write_all(b"hello").await.unwrap(); + }); + + let mut client = TcpStream::connect(&addr).await.unwrap(); + + let mut buf = vec![]; + client.read_to_end(&mut buf).await.unwrap(); + + assert_eq!(buf, b"hello"); + tx.send(()).unwrap(); } } diff --git a/tokio/tests/task_local_set.rs b/tokio/tests/task_local_set.rs index 42bd4607a..1a10fefa6 100644 --- a/tokio/tests/task_local_set.rs +++ b/tokio/tests/task_local_set.rs @@ -1,20 +1,15 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "full")] -use std::{ - cell::Cell, - sync::atomic::{ - AtomicBool, AtomicUsize, - Ordering::{self, SeqCst}, - }, - time::Duration, -}; -use tokio::{ - runtime::{self, Runtime}, - sync::{mpsc, oneshot}, - task::{self, LocalSet}, - time, -}; +use tokio::runtime::{self, Runtime}; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::{self, LocalSet}; +use tokio::time; + +use std::cell::Cell; +use std::sync::atomic::Ordering::{self, SeqCst}; +use std::sync::atomic::{AtomicBool, AtomicUsize}; +use std::time::Duration; #[tokio::test(basic_scheduler)] async fn local_basic_scheduler() { @@ -285,15 +280,23 @@ fn join_local_future_elsewhere() { join2.await.unwrap() }); } + #[test] fn drop_cancels_tasks() { + use std::rc::Rc; + // This test reproduces issue #1842 let mut rt = rt(); + let rc1 = Rc::new(()); + let rc2 = rc1.clone(); let (started_tx, started_rx) = oneshot::channel(); let local = LocalSet::new(); local.spawn_local(async move { + // Move this in + let _rc2 = rc2; + started_tx.send(()).unwrap(); loop { time::delay_for(Duration::from_secs(3600)).await; @@ -305,6 +308,8 @@ fn drop_cancels_tasks() { }); drop(local); drop(rt); + + assert_eq!(1, Rc::strong_count(&rc1)); } #[test]