rt: cleanup and simplify scheduler (scheduler v2.5) (#2273)

A refactor of the scheduler internals focusing on simplifying and
reducing unsafety. There are no fundamental logic changes.

* The state transitions of the core task component are refined and
reduced.
* `basic_scheduler` has most unsafety removed.
* `local_set` has most unsafety removed.
* `threaded_scheduler` limits most unsafety to its queue implementation.

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

benches/Cargo.toml

@@ -17,3 +17,8 @@ harness = false
 name = "mpsc"
 path = "mpsc.rs"
 harness = false
+
+[[bench]]
+name = "scheduler"
+path = "scheduler.rs"
+harness = false

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);

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:

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

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

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

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:

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:

tokio/src/lib.rs

@@ -382,10 +382,6 @@ cfg_macros! {
     }
 }
 
-// Tests
-#[cfg(test)]
-mod tests;
-
 // TODO: rm
 #[cfg(feature = "io-util")]
 #[cfg(test)]

tokio/src/loom/std/alloc.rs

@@ -1,18 +0,0 @@
-#[derive(Debug)]
-pub(crate) struct Track<T> {
-    value: T,
-}
-
-impl<T> Track<T> {
-    pub(crate) fn new(value: T) -> Track<T> {
-        Track { value }
-    }
-
-    pub(crate) fn get_mut(&mut self) -> &mut T {
-        &mut self.value
-    }
-
-    pub(crate) fn into_inner(self) -> T {
-        self.value
-    }
-}

tokio/src/loom/std/causal_cell.rs

@@ -18,15 +18,6 @@ impl<T> CausalCell<T> {
         f(self.0.get())
     }
 
-    pub(crate) fn with_unchecked<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce(*const T) -> R,
-    {
-        f(self.0.get())
-    }
-
-    pub(crate) fn check(&self) {}
-
     pub(crate) fn with_deferred<F, R>(&self, f: F) -> (R, CausalCheck)
     where
         F: FnOnce(*const T) -> R,

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};
 }

tokio/src/macros/assert.rs

@@ -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);
-        }
-    }};
-}

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;

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<T> {
+    #[doc(hidden)]
+    pub(crate) inner: &'static LocalKey<Cell<*const ()>>,
+    #[doc(hidden)]
+    pub(crate) _marker: marker::PhantomData<T>,
+}
+
+unsafe impl<T> Sync for ScopedKey<T> {}
+
+impl<T> ScopedKey<T> {
+    /// Inserts a value into this scoped thread local storage slot for a
+    /// duration of a closure.
+    pub(crate) fn set<F, R>(&'static self, t: &T, f: F) -> R
+    where
+        F: FnOnce() -> R,
+    {
+        struct Reset {
+            key: &'static LocalKey<Cell<*const ()>>,
+            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<F, R>(&'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))) }
+        }
+    }
+}

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) {

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<P>
 where
     P: Park,
 {
-    /// Scheduler component
-    scheduler: Arc<SchedulerPriv>,
+    /// 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<Tasks>,
 
-    /// Local state
-    local: LocalState<P>,
-}
+    /// Sendable task spawner
+    spawner: Spawner,
 
-#[derive(Debug, Clone)]
-pub(crate) struct Spawner {
-    scheduler: Arc<SchedulerPriv>,
-}
-
-/// The scheduler component.
-pub(super) struct SchedulerPriv {
-    queues: MpscQueues<Self>,
-    /// Unpark the blocked thread
-    unpark: Box<dyn Unpark>,
-}
-
-unsafe impl Send for SchedulerPriv {}
-unsafe impl Sync for SchedulerPriv {}
-
-/// Local state
-#[derive(Debug)]
-struct LocalState<P> {
     /// Current tick
     tick: u8,
 
@@ -48,33 +35,76 @@ struct LocalState<P> {
     park: P,
 }
 
+#[derive(Clone)]
+pub(crate) struct Spawner {
+    shared: Arc<Shared>,
+}
+
+struct Tasks {
+    /// Collection of all active tasks spawned onto this executor.
+    owned: LinkedList<Task<Arc<Shared>>>,
+
+    /// Local run queue.
+    ///
+    /// Tasks notified from the current thread are pushed into this queue.
+    queue: VecDeque<task::Notified<Arc<Shared>>>,
+}
+
+/// Scheduler state shared between threads.
+struct Shared {
+    /// Remote run queue
+    queue: Mutex<VecDeque<task::Notified<Arc<Shared>>>>,
+
+    /// Unpark the blocked thread
+    unpark: Box<dyn Unpark>,
+}
+
+/// Thread-local context
+struct Context {
+    /// Shared scheduler state
+    shared: Arc<Shared>,
+
+    /// Local queue
+    tasks: RefCell<Tasks>,
+}
+
+/// 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<P> BasicScheduler<P>
 where
     P: Park,
 {
     pub(crate) fn new(park: P) -> BasicScheduler<P> {
-        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<dyn Unpark>,
+                }),
+            },
+            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<F>(&mut self, mut future: F) -> F::Output
+    pub(crate) fn block_on<F>(&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<F, R, P>(scheduler: &mut BasicScheduler<P>, f: F) -> R
+where
+    F: FnOnce(&mut BasicScheduler<P>, &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<Context>,
+        scheduler: &'a mut BasicScheduler<P>,
+    }
+
+    impl<P: Park> 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<P> Drop for BasicScheduler<P>
+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<P: Park> fmt::Debug for BasicScheduler<P> {
+    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<F>(&self, future: F) -> JoinHandle<F::Output>
@@ -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<impl Park>) {
-        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<Self>, 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<task::Notified<Arc<Shared>>> {
+        self.shared.queue.lock().unwrap().pop_front()
     }
 }
 
-impl Schedule for SchedulerPriv {
-    fn bind(&self, task: &Task<Self>) {
-        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>) {
-        self.queues.release_remote(task);
-    }
-
-    fn release_local(&self, task: &Task<Self>) {
-        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<Self>) {
-        SchedulerPriv::schedule(self, task, false);
-    }
-}
-
-impl ScheduleSendOnly for SchedulerPriv {}
-
-impl<P> Drop for BasicScheduler<P>
-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<Shared> {
+    fn bind(task: Task<Self>) -> Arc<Shared> {
+        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<Self>) -> Option<Task<Self>> {
+        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<Self>) {
+        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<Self>) {
+        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<Self>) {
+        arc_self.unpark.unpark();
+    }
 }

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<shutdown::Sender>,
 }
 
-type Task = task::Task<NoopSchedule>;
+type Task = task::Notified<NoopSchedule>;
 
 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());
-}

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<Self>) {}
+impl task::Schedule for NoopSchedule {
+    fn bind(_task: Task<Self>) -> NoopSchedule {
+        // Do nothing w/ the task
+        NoopSchedule
+    }
 
-    fn release(&self, _task: Task<Self>) {}
+    fn release(&self, _task: &Task<Self>) -> Option<Task<Self>> {
+        None
+    }
 
-    fn release_local(&self, _task: &Task<Self>) {}
-
-    fn schedule(&self, _task: Task<Self>) {
+    fn schedule(&self, _task: task::Notified<Self>) {
         unreachable!();
     }
 }
-
-impl ScheduleSendOnly for NoopSchedule {}

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),

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<Self> {
+    pub fn new() -> io::Result<Runtime> {
         #[cfg(feature = "rt-threaded")]
         let ret = Builder::new().threaded_scheduler().enable_all().build();
 

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<T: Future, S> {
+    /// Hot task state data
+    pub(super) header: Header,
+
+    /// Either the future or output, depending on the execution stage.
+    pub(super) core: Core<T, S>,
+
+    /// Cold data
+    pub(super) trailer: Trailer,
+}
+
+/// The core of the task.
+///
+/// Holds the future or output, depending on the stage of execution.
+pub(super) struct Core<T: Future, S> {
+    /// Scheduler used to drive this future
+    pub(super) scheduler: CausalCell<Option<S>>,
+
+    /// Either the future or the output
+    pub(super) stage: CausalCell<Stage<T>>,
+}
+
+/// 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<linked_list::Pointers<Header>>,
+
+    /// Pointer to next task, used with the injection queue
+    pub(crate) queue_next: UnsafeCell<Option<NonNull<Header>>>,
+
+    /// Pointer to the next task in the transfer stack
+    pub(super) stack_next: UnsafeCell<Option<NonNull<Header>>>,
+
+    /// 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<Option<Waker>>,
+}
+
+/// Either the future or the output.
+pub(super) enum Stage<T: Future> {
+    Running(T),
+    Finished(super::Result<T::Output>),
+    Consumed,
+}
+
+impl<T: Future, S: Schedule> Cell<T, S> {
+    /// Allocates a new task cell, containing the header, trailer, and core
+    /// structures.
+    pub(super) fn new(future: T, state: State) -> Box<Cell<T, S>> {
+        Box::new(Cell {
+            header: Header {
+                state,
+                owned: UnsafeCell::new(linked_list::Pointers::new()),
+                queue_next: UnsafeCell::new(None),
+                stack_next: UnsafeCell::new(None),
+                vtable: raw::vtable::<T, S>(),
+            },
+            core: Core {
+                scheduler: CausalCell::new(None),
+                stage: CausalCell::new(Stage::Running(future)),
+            },
+            trailer: Trailer {
+                waker: CausalCell::new(None),
+            },
+        })
+    }
+}
+
+impl<T: Future, S: Schedule> Core<T, S> {
+    /// If needed, bind a scheduler to the task.
+    ///
+    /// This only happens on the first poll.
+    pub(super) fn bind_scheduler(&self, task: Task<S>) {
+        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<T::Output> {
+        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::<T, S>(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<T::Output>) {
+        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<T::Output> {
+        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<S>) {
+        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<S>) {
+        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<S>) -> Option<Task<S>> {
+        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::<Header>() <= 8 * size_of::<*const ()>());
+}

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<T: Future, S: 'static> {
+    cell: NonNull<Cell<T, S>>,
+}
+
+impl<T, S> Harness<T, S>
+where
+    T: Future,
+    S: 'static,
+{
+    pub(super) unsafe fn from_raw(ptr: NonNull<Header>) -> Harness<T, S> {
+        Harness {
+            cell: ptr.cast::<Cell<T, S>>(),
+        }
+    }
+
+    fn header(&self) -> &Header {
+        unsafe { &self.cell.as_ref().header }
+    }
+
+    fn trailer(&self) -> &Trailer {
+        unsafe { &self.cell.as_ref().trailer }
+    }
+
+    fn core(&self) -> &Core<T, S> {
+        unsafe { &self.cell.as_ref().core }
+    }
+}
+
+impl<T, S> Harness<T, S>
+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<T, S>,
+                polled: bool,
+            }
+
+            impl<T: Future, S: Schedule> 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<super::Result<T::Output>>, 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<Snapshot, Snapshot> {
+        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<T::Output>, 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<S> {
+        unsafe { Task::from_raw(self.header().into()) }
+    }
+}

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<T> Unpin for JoinHandle<T> {}
 impl<T> Future for JoinHandle<T> {
     type Output = super::Result<T>;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        use std::mem::MaybeUninit;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        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::<Track<Self::Output>>::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<T> Drop for JoinHandle<T> {
     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;
             }
 

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<S: 'static> {
+    raw: RawTask,
+    _p: PhantomData<S>,
+}
+
+unsafe impl<S> Send for Task<S> {}
+unsafe impl<S> Sync for Task<S> {}
+
+/// A task was notified
+#[repr(transparent)]
+pub(crate) struct Notified<S: 'static>(Task<S>);
+
+unsafe impl<S: Schedule> Send for Notified<S> {}
+unsafe impl<S: Schedule> Sync for Notified<S> {}
+
+/// Task result sent back
+pub(crate) type Result<T> = std::result::Result<T, JoinError>;
+
+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>) -> 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<Self>) -> Option<Task<Self>>;
+
+    /// Schedule the task
+    fn schedule(&self, task: Notified<Self>);
+
+    /// Schedule the task to run in the near future, yielding the thread to
+    /// other tasks.
+    fn yield_now(&self, task: Notified<Self>) {
+        self.schedule(task);
+    }
+}
+
+/// Create a new task with an associated join handle
+pub(crate) fn joinable<T, S>(task: T) -> (Notified<S>, JoinHandle<T::Output>)
+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<T, S>(task: T) -> (Notified<S>, JoinHandle<T::Output>)
+    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<S: 'static> Task<S> {
+    pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
+        Task {
+            raw: RawTask::from_raw(ptr),
+            _p: PhantomData,
+        }
+    }
+
+    pub(crate) fn header(&self) -> &Header {
+        self.raw.header()
+    }
+}
+
+cfg_rt_threaded! {
+    impl<S: 'static> Notified<S> {
+        pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> {
+            Notified(Task::from_raw(ptr))
+        }
+
+        pub(crate) fn header(&self) -> &Header {
+            self.0.header()
+        }
+    }
+
+    impl<S: 'static> Task<S> {
+        pub(crate) fn into_raw(self) -> NonNull<Header> {
+            let ret = self.header().into();
+            mem::forget(self);
+            ret
+        }
+    }
+
+    impl<S: 'static> Notified<S> {
+        pub(crate) fn into_raw(self) -> NonNull<Header> {
+            self.0.into_raw()
+        }
+    }
+}
+
+impl<S: Schedule> Task<S> {
+    /// Pre-emptively cancel the task as part of the shutdown process.
+    pub(crate) fn shutdown(&self) {
+        self.raw.shutdown();
+    }
+}
+
+impl<S: Schedule> Notified<S> {
+    /// 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<S: 'static> Drop for Task<S> {
+    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<S> fmt::Debug for Task<S> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(fmt, "Task({:p})", self.header())
+    }
+}
+
+impl<S> fmt::Debug for Notified<S> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(fmt, "task::Notified({:p})", self.0.header())
+    }
+}
+
+/// # Safety
+///
+/// Tasks are pinned
+unsafe impl<S> linked_list::Link for Task<S> {
+    type Handle = Task<S>;
+    type Target = Header;
+
+    fn as_raw(handle: &Task<S>) -> NonNull<Header> {
+        handle.header().into()
+    }
+
+    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
+        Task::from_raw(ptr)
+    }
+
+    unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
+        NonNull::from(&mut *target.as_ref().owned.get())
+    }
+}

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<Header>,
+}
+
+pub(super) struct Vtable {
+    /// Poll the future
+    pub(super) poll: unsafe fn(NonNull<Header>),
+
+    /// Deallocate the memory
+    pub(super) dealloc: unsafe fn(NonNull<Header>),
+
+    /// Read the task output, if complete
+    pub(super) try_read_output: unsafe fn(NonNull<Header>, *mut (), &Waker),
+
+    /// The join handle has been dropped
+    pub(super) drop_join_handle_slow: unsafe fn(NonNull<Header>),
+
+    /// Scheduler is being shutdown
+    pub(super) shutdown: unsafe fn(NonNull<Header>),
+}
+
+/// Get the vtable for the requested `T` and `S` generics.
+pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
+    &Vtable {
+        poll: poll::<T, S>,
+        dealloc: dealloc::<T, S>,
+        try_read_output: try_read_output::<T, S>,
+        drop_join_handle_slow: drop_join_handle_slow::<T, S>,
+        shutdown: shutdown::<T, S>,
+    }
+}
+
+impl RawTask {
+    pub(super) fn new<T, S>(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<Header>) -> 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<super::Result<T::Output>>` 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<T: Future, S: Schedule>(ptr: NonNull<Header>) {
+    let harness = Harness::<T, S>::from_raw(ptr);
+    harness.poll();
+}
+
+unsafe fn dealloc<T: Future, S: Schedule>(ptr: NonNull<Header>) {
+    let harness = Harness::<T, S>::from_raw(ptr);
+    harness.dealloc();
+}
+
+unsafe fn try_read_output<T: Future, S: Schedule>(
+    ptr: NonNull<Header>,
+    dst: *mut (),
+    waker: &Waker,
+) {
+    let out = &mut *(dst as *mut Poll<super::Result<T::Output>>);
+
+    let harness = Harness::<T, S>::from_raw(ptr);
+    harness.try_read_output(out, waker);
+}
+
+unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: NonNull<Header>) {
+    let harness = Harness::<T, S>::from_raw(ptr);
+    harness.drop_join_handle_slow()
+}
+
+unsafe fn shutdown<T: Future, S: Schedule>(ptr: NonNull<Header>) {
+    let harness = Harness::<T, S>::from_raw(ptr);
+    harness.shutdown()
+}

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<T: 'static> {
+    head: AtomicPtr<Header>,
+    _p: PhantomData<T>,
+}
+
+impl<T: 'static> TransferStack<T> {
+    pub(crate) fn new() -> TransferStack<T> {
+        TransferStack {
+            head: AtomicPtr::new(ptr::null_mut()),
+            _p: PhantomData,
+        }
+    }
+
+    pub(crate) fn push(&self, task: Task<T>) {
+        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<Item = Task<T>> {
+        struct Iter<T: 'static>(Option<NonNull<Header>>, PhantomData<T>);
+
+        impl<T: 'static> Iterator for Iter<T> {
+            type Item = Task<T>;
+
+            fn next(&mut self) -> Option<Task<T>> {
+                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<T: 'static> Drop for Iter<T> {
+            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)
+    }
+}

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<Snapshot, Snapshot>;
+
+/// 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<F>(&self, mut f: F) -> Result<Snapshot, Snapshot>
+    where
+        F: FnMut(Snapshot) -> Option<Snapshot>,
+    {
+        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()
+    }
+}

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<T, S>(meta: &Header) -> WakerRef<'_, S>
+pub(super) fn waker_ref<T, S>(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::<T, S>(meta))) };
+    let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker::<T, S>(header))) };
 
     WakerRef {
         waker,
@@ -48,9 +49,9 @@ where
     T: Future,
     S: Schedule,
 {
-    let meta = ptr as *const Header;
-    (*meta).state.ref_inc();
-    raw_waker::<T, S>(meta)
+    let header = ptr as *const Header;
+    (*header).state.ref_inc();
+    raw_waker::<T, S>(header)
 }
 
 unsafe fn drop_waker<T, S>(ptr: *const ())
@@ -58,8 +59,9 @@ where
     T: Future,
     S: Schedule,
 {
-    let harness = Harness::<T, S>::from_raw(ptr as *mut _);
-    harness.drop_waker();
+    let ptr = NonNull::new_unchecked(ptr as *mut Header);
+    let harness = Harness::<T, S>::from_raw(ptr);
+    harness.drop_reference();
 }
 
 unsafe fn wake_by_val<T, S>(ptr: *const ())
@@ -67,7 +69,8 @@ where
     T: Future,
     S: Schedule,
 {
-    let harness = Harness::<T, S>::from_raw(ptr as *mut _);
+    let ptr = NonNull::new_unchecked(ptr as *mut Header);
+    let harness = Harness::<T, S>::from_raw(ptr);
     harness.wake_by_val();
 }
 
@@ -77,16 +80,17 @@ where
     T: Future,
     S: Schedule,
 {
-    let harness = Harness::<T, S>::from_raw(ptr as *mut _);
+    let ptr = NonNull::new_unchecked(ptr as *mut Header);
+    let harness = Harness::<T, S>::from_raw(ptr);
     harness.wake_by_ref();
 }
 
-fn raw_waker<T, S>(meta: *const Header) -> RawWaker
+fn raw_waker<T, S>(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::<T, S>,
         wake_by_val::<T, S>,

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<Output = &'static str> {
+    gated2(false)
+}
+
+fn gated2(thread: bool) -> impl Future<Output = &'static str> {
+    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<T: Future>(f: T) -> Track<T> {
+    Track {
+        inner: f,
+        arc: Arc::new(()),
+    }
+}
+
+pin_project! {
+    struct Track<T> {
+        #[pin]
+        inner: T,
+        // Arc is used to hook into loom's leak tracking.
+        arc: Arc<()>,
+    }
+}
+
+impl<T> Track<T> {
+    fn into_inner(self) -> T {
+        self.inner
+    }
+}
+
+impl<T: Future> Future for Track<T> {
+    type Output = Track<T::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let me = self.project();
+
+        Poll::Ready(Track {
+            inner: ready!(me.inner.poll(cx)),
+            arc: me.arc.clone(),
+        })
+    }
+}

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;

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<Inner>);
+
+struct Inner {
+    released: task::TransferStack<Runtime>,
+    core: TryLock<Core>,
+}
+
+struct Core {
+    queue: VecDeque<task::Notified<Runtime>>,
+    tasks: LinkedList<Task<Runtime>>,
+}
+
+static CURRENT: TryLock<Option<Runtime>> = 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<Runtime> {
+        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<Self>) -> 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<Self>) -> Option<Task<Self>> {
+        // 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>) {
+        self.0.core.try_lock().unwrap().queue.push_back(task);
+    }
+}

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<T> {
+    data: AtomicPtr<T>,
+}
+
+unsafe impl<T: Send> Send for AtomicCell<T> {}
+unsafe impl<T: Send> Sync for AtomicCell<T> {}
+
+impl<T> AtomicCell<T> {
+    pub(super) fn new(data: Option<Box<T>>) -> AtomicCell<T> {
+        AtomicCell {
+            data: AtomicPtr::new(to_raw(data)),
+        }
+    }
+
+    pub(super) fn swap(&self, val: Option<Box<T>>) -> Option<Box<T>> {
+        let old = self.data.swap(to_raw(val), AcqRel);
+        from_raw(old)
+    }
+
+    #[cfg(feature = "blocking")]
+    pub(super) fn set(&self, val: Box<T>) {
+        let _ = self.swap(Some(val));
+    }
+
+    pub(super) fn take(&self) -> Option<Box<T>> {
+        self.swap(None)
+    }
+}
+
+fn to_raw<T>(data: Option<Box<T>>) -> *mut T {
+    data.map(Box::into_raw).unwrap_or(ptr::null_mut())
+}
+
+fn from_raw<T>(val: *mut T) -> Option<Box<T>> {
+    if val.is_null() {
+        None
+    } else {
+        Some(unsafe { Box::from_raw(val) })
+    }
+}
+
+impl<T> Drop for AtomicCell<T> {
+    fn drop(&mut self) {
+        // Free any data still held by the cell
+        let _ = self.take();
+    }
+}

tokio/src/runtime/thread_pool/current.rs

@@ -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<Inner> = Cell::new(Inner::new()));
-
-pub(super) fn set<F, R>(pool: &Arc<slice::Set>, index: usize, f: F) -> R
-where
-    F: FnOnce() -> R,
-{
-    CURRENT_WORKER.with(|cell| {
-        assert!(cell.get().workers.is_null());
-
-        struct Guard<'a>(&'a Cell<Inner>);
-
-        impl Drop for Guard<'_> {
-            fn drop(&mut self) {
-                self.0.set(Inner::new());
-            }
-        }
-
-        cell.set(Inner {
-            workers: pool.shared() as *const _ as *const (),
-            idx: index,
-        });
-
-        let _g = Guard(cell);
-
-        f()
-    })
-}
-
-pub(super) fn clear() {
-    CURRENT_WORKER.with(|cell| cell.set(Inner::new()))
-}
-
-pub(super) fn get<F, R>(f: F) -> R
-where
-    F: FnOnce(&Current) -> R,
-{
-    CURRENT_WORKER.with(|cell| {
-        let current = Current { inner: cell.get() };
-        f(&current)
-    })
-}
-
-impl Current {
-    pub(super) fn as_member<'a>(&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,
-        }
-    }
-}

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<Worker>,
+/// 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<worker::Shared>,
 }
 
+// ===== 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<F>(&self, future: F) -> JoinHandle<F::Output>
+    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()
     }
 }

tokio/src/runtime/thread_pool/owned.rs

@@ -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<u16>,
-
-    /// Caches the pool run state.
-    pub(super) is_running: Cell<bool>,
-
-    /// `true` if the worker is currently searching for more work.
-    pub(super) is_searching: Cell<bool>,
-
-    /// `true` when worker notification should be delayed.
-    ///
-    /// This is used to batch notifications triggered by the parker.
-    pub(super) defer_notification: Cell<bool>,
-
-    /// `true` if a task was submitted while `defer_notification` was set
-    pub(super) did_submit_task: Cell<bool>,
-
-    /// Fast random number generator
-    pub(super) rand: FastRand,
-
-    /// Work queue
-    pub(super) work_queue: queue::Worker<Shared>,
-
-    /// List of tasks owned by the worker
-    pub(super) owned_tasks: task::OwnedList<Shared>,
-}
-
-impl Owned {
-    pub(super) fn new(work_queue: queue::Worker<Shared>, 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<Shared>) -> bool {
-        let ret = self.work_queue.push(task);
-
-        if self.defer_notification.get() {
-            self.did_submit_task.set(true);
-            false
-        } else {
-            ret
-        }
-    }
-
-    pub(super) fn submit_local_yield(&self, task: Task<Shared>) {
-        self.work_queue.push_yield(task);
-    }
-
-    pub(super) fn bind_task(&mut self, task: &Task<Shared>) {
-        self.owned_tasks.insert(task);
-    }
-
-    pub(super) fn release_task(&mut self, task: &Task<Shared>) {
-        self.owned_tasks.remove(task);
-    }
-}

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<T: 'static> {
+    inner: Arc<Inner<T>>,
+
+    /// LIFO slot. Cannot be stolen.
+    next: Option<task::Notified<T>>,
+}
+
+/// Consumer handle. May be used from many threads.
+pub(super) struct Steal<T: 'static>(Arc<Inner<T>>);
+
+/// 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<T: 'static> {
+    /// Pointers to the head and tail of the queue
+    pointers: Mutex<Pointers>,
+
+    /// Number of pending tasks in the queue. This helps prevent unnecessary
+    /// locking in the hot path.
+    len: AtomicUsize,
+
+    _p: PhantomData<T>,
+}
+
+pub(super) struct Inner<T: 'static> {
+    /// Concurrently updated by many threads.
+    head: AtomicU32,
+
+    /// Only updated by producer thread but read by many threads.
+    tail: AtomicU32,
+
+    /// Elements
+    buffer: Box<[CausalCell<MaybeUninit<task::Notified<T>>>]>,
+}
+
+struct Pointers {
+    /// True if the queue is closed
+    is_closed: bool,
+
+    /// Linked-list head
+    head: Option<NonNull<task::Header>>,
+
+    /// Linked-list tail
+    tail: Option<NonNull<task::Header>>,
+}
+
+unsafe impl<T> Send for Inner<T> {}
+unsafe impl<T> Sync for Inner<T> {}
+unsafe impl<T> Send for Inject<T> {}
+unsafe impl<T> Sync for Inject<T> {}
+
+#[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<T: 'static>() -> (Steal<T>, Local<T>) {
+    debug_assert!(LOCAL_QUEUE_CAPACITY >= 2 && LOCAL_QUEUE_CAPACITY.is_power_of_two());
+
+    let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY);
+
+    for _ in 0..LOCAL_QUEUE_CAPACITY {
+        buffer.push(CausalCell::new(MaybeUninit::uninit()));
+    }
+
+    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<T> Local<T> {
+    /// 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<T>, inject: &Inject<T>) -> 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<T>, inject: &Inject<T>) {
+        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<T>,
+        head: u32,
+        tail: u32,
+        inject: &Inject<T>,
+    ) -> Result<(), task::Notified<T>> {
+        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<task::Notified<T>> {
+        // 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<T> Steal<T> {
+    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<T>) -> Option<task::Notified<T>> {
+        // 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<T>, 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<T> Drop for Local<T> {
+    fn drop(&mut self) {
+        if !std::thread::panicking() {
+            assert!(self.pop().is_none(), "queue not empty");
+        }
+    }
+}
+
+impl<T> Inner<T> {
+    fn is_empty(&self) -> bool {
+        let head = self.head.load(Acquire);
+        let tail = self.tail.load(Acquire);
+
+        head == tail
+    }
+}
+
+impl<T: 'static> Inject<T> {
+    pub(super) fn new() -> Inject<T> {
+        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<T>) {
+        // 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<T>,
+        batch_tail: task::Notified<T>,
+        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<task::Notified<T>> {
+        // 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<T: 'static> Drop for Inject<T> {
+    fn drop(&mut self) {
+        if !std::thread::panicking() {
+            assert!(self.pop().is_none(), "queue not empty");
+        }
+    }
+}
+
+fn get_next(header: NonNull<task::Header>) -> Option<NonNull<task::Header>> {
+    unsafe { *header.as_ref().queue_next.get() }
+}
+
+fn set_next(header: NonNull<task::Header>, val: Option<NonNull<task::Header>>) {
+    unsafe {
+        *header.as_ref().queue_next.get() = val;
+    }
+}

tokio/src/runtime/thread_pool/queue/global.rs

@@ -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<T: 'static> {
-    /// Pointers to the head and tail of the queue
-    pointers: Mutex<Pointers>,
-
-    /// 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<T>,
-}
-
-struct Pointers {
-    head: *const Header,
-    tail: *const Header,
-}
-
-const CLOSED: usize = 1;
-const MAX_LEN: usize = usize::MAX >> 1;
-
-impl<T: 'static> Queue<T> {
-    pub(super) fn new() -> Queue<T> {
-        Queue {
-            pointers: Mutex::new(Pointers {
-                head: ptr::null(),
-                tail: ptr::null(),
-            }),
-            len: AtomicUsize::new(0),
-            _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<F>(&self, task: Task<T>, f: F)
-    where
-        F: FnOnce(Result<(), Task<T>>),
-    {
-        unsafe {
-            // Acquire queue lock
-            let mut p = self.pointers.lock().unwrap();
-
-            // Check if the queue is closed. This must happen in the lock.
-            let len = self.len.unsync_load();
-            if len & CLOSED == CLOSED {
-                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<T>, batch_tail: Task<T>, num: usize) {
-        unsafe {
-            let batch_head = batch_head.into_raw().as_ptr();
-            let batch_tail = batch_tail.into_raw();
-
-            debug_assert!(get_next(batch_tail).is_null());
-
-            let mut p = self.pointers.lock().unwrap();
-
-            if let Some(tail) = NonNull::new(p.tail as *mut _) {
-                set_next(tail, batch_head);
-            } else {
-                p.head = batch_head;
-            }
-
-            p.tail = batch_tail.as_ptr();
-
-            // Increment the count.
-            //
-            // All updates to the len atomic are guarded by the mutex. As such,
-            // a non-atomic load followed by a store is safe.
-            //
-            // Left shift by 1 to avoid touching the shutdown flag.
-            let len = self.len.unsync_load();
-
-            if (len >> 1) >= (MAX_LEN - num) {
-                std::process::abort();
-            }
-
-            self.len.store(len + (num << 1), Release);
-
-            drop(p);
-        }
-    }
-
-    pub(super) fn pop(&self) -> Option<Task<T>> {
-        // Fast path, if len == 0, then there are no values
-        if self.is_empty() {
-            return None;
-        }
-
-        unsafe {
-            let mut p = self.pointers.lock().unwrap();
-
-            // It is possible to hit null here if another thread poped the last
-            // task between us checking `len` and acquiring the lock.
-            let task = NonNull::new(p.head as *mut _)?;
-
-            p.head = get_next(task);
-
-            if p.head.is_null() {
-                p.tail = ptr::null();
-            }
-
-            set_next(task, ptr::null());
-
-            // Decrement the count.
-            //
-            // All updates to the len atomic are guarded by the mutex. As such,
-            // a non-atomic load followed by a store is safe.
-            //
-            // Decrement by 2 to avoid touching the shutdown flag
-            self.len.store(self.len.unsync_load() - 2, Release);
-
-            drop(p);
-
-            Some(Task::from_raw(task))
-        }
-    }
-}
-
-unsafe fn get_next(meta: NonNull<Header>) -> *const Header {
-    *meta.as_ref().queue_next.get()
-}
-
-unsafe fn set_next(meta: NonNull<Header>, val: *const Header) {
-    *meta.as_ref().queue_next.get() = val;
-}

tokio/src/runtime/thread_pool/queue/inject.rs

@@ -1,41 +0,0 @@
-use crate::loom::sync::Arc;
-use crate::runtime::thread_pool::queue::Cluster;
-use crate::task::Task;
-
-pub(crate) struct Inject<T: 'static> {
-    cluster: Arc<Cluster<T>>,
-}
-
-impl<T: 'static> Inject<T> {
-    pub(super) fn new(cluster: Arc<Cluster<T>>) -> Inject<T> {
-        Inject { cluster }
-    }
-
-    /// Pushes a value onto the queue
-    pub(crate) fn push<F>(&self, task: Task<T>, f: F)
-    where
-        F: FnOnce(Result<(), Task<T>>),
-    {
-        self.cluster.global.push(task, f)
-    }
-
-    /// 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();
-    }
-}

tokio/src/runtime/thread_pool/queue/local.rs

@@ -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<T: 'static> {
-    /// Concurrently updated by many threads.
-    head: AtomicU32,
-
-    /// Only updated by producer thread but read by many threads.
-    tail: AtomicU32,
-
-    /// Elements
-    buffer: Box<[CausalCell<MaybeUninit<Task<T>>>]>,
-}
-
-const MASK: usize = LOCAL_QUEUE_CAPACITY - 1;
-
-impl<T: 'static> Queue<T> {
-    pub(super) fn new() -> Queue<T> {
-        debug_assert!(LOCAL_QUEUE_CAPACITY >= 2 && LOCAL_QUEUE_CAPACITY.is_power_of_two());
-
-        let mut buffer = Vec::with_capacity(LOCAL_QUEUE_CAPACITY);
-
-        for _ in 0..LOCAL_QUEUE_CAPACITY {
-            buffer.push(CausalCell::new(MaybeUninit::uninit()));
-        }
-
-        Queue {
-            head: AtomicU32::new(0),
-            tail: AtomicU32::new(0),
-            buffer: buffer.into(),
-        }
-    }
-}
-
-impl<T> Queue<T> {
-    /// Pushes a task onto the local queue.
-    ///
-    /// This **must** be called by the producer thread.
-    pub(super) unsafe fn push(&self, mut task: Task<T>, global: &global::Queue<T>) {
-        loop {
-            let head = self.head.load(Acquire);
-
-            // safety: this is the **only** thread that updates this cell.
-            let tail = self.tail.unsync_load();
-
-            if tail.wrapping_sub(head) < LOCAL_QUEUE_CAPACITY as u32 {
-                // Map the position to a slot index.
-                let idx = tail as usize & MASK;
-
-                self.buffer[idx].with_mut(|ptr| {
-                    // Write the task to the slot
-                    ptr::write((*ptr).as_mut_ptr(), task);
-                });
-
-                // Make the task available
-                self.tail.store(tail.wrapping_add(1), Release);
-
-                return;
-            }
-
-            // The local buffer is full. Push a batch of work to the global
-            // queue.
-            match self.push_overflow(task, head, tail, global) {
-                Ok(_) => return,
-                // Lost the race, try again
-                Err(v) => task = v,
-            }
-
-            atomic::spin_loop_hint();
-        }
-    }
-
-    /// 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<T>,
-        head: u32,
-        tail: u32,
-        global: &global::Queue<T>,
-    ) -> Result<(), Task<T>> {
-        const BATCH_LEN: usize = LOCAL_QUEUE_CAPACITY / 2 + 1;
-
-        let n = tail.wrapping_sub(head) / 2;
-        assert_eq!(n as usize, LOCAL_QUEUE_CAPACITY / 2, "queue is not full");
-
-        // Claim a bunch of tasks
-        //
-        // We are claiming the tasks **before** reading them out of the buffer.
-        // This is safe because only the **current** thread is able to push new
-        // tasks.
-        //
-        // There isn't really any need for memory ordering... Relaxed would
-        // work. This is because all tasks are pushed into the queue from the
-        // current thread (or memory has been acquired if the local queue handle
-        // moved).
-        let actual = self.head.compare_and_swap(head, head + n, Release);
-        if actual != head {
-            // We failed to claim the tasks, losing the race. Return out of
-            // this function and try the full `push` routine again. The queue
-            // may not be full anymore.
-            return Err(task);
-        }
-
-        // link the tasks
-        for i in 0..n {
-            let j = i + 1;
-
-            let i_idx = (i + head) as usize & MASK;
-            let j_idx = (j + head) as usize & MASK;
-
-            // Get the next pointer
-            let next = if j == n {
-                // The last task in the local queue being moved
-                task.header() as *const _
-            } else {
-                self.buffer[j_idx].with(|ptr| {
-                    let value = (*ptr).as_ptr();
-                    (*value).header() as *const _
-                })
-            };
-
-            self.buffer[i_idx].with_mut(|ptr| {
-                let ptr = (*ptr).as_ptr();
-                debug_assert!((*(*ptr).header().queue_next.get()).is_null());
-                *(*ptr).header().queue_next.get() = next;
-            });
-        }
-
-        let head = self.buffer[head as usize & MASK].with(|ptr| ptr::read((*ptr).as_ptr()));
-
-        // Push the tasks onto the global queue
-        global.push_batch(head, task, BATCH_LEN);
-
-        Ok(())
-    }
-
-    /// Pops a task from the local queue.
-    ///
-    /// This **must** be called by the producer thread
-    pub(super) unsafe fn pop(&self) -> Option<Task<T>> {
-        loop {
-            let head = self.head.load(Acquire);
-
-            // safety: this is the **only** thread that updates this cell.
-            let tail = self.tail.unsync_load();
-
-            if head == tail {
-                // queue is empty
-                return None;
-            }
-
-            // Map the head position to a slot index.
-            let idx = head as usize & MASK;
-
-            let task = self.buffer[idx].with(|ptr| {
-                // Tentatively read the task at the head position. Note that we
-                // have not yet claimed the task.
-                //
-                ptr::read(ptr)
-            });
-
-            // Attempt to claim the task read above.
-            let actual = self
-                .head
-                .compare_and_swap(head, head.wrapping_add(1), Release);
-
-            if actual == head {
-                return Some(task.assume_init());
-            }
-
-            atomic::spin_loop_hint();
-        }
-    }
-
-    pub(super) fn is_empty(&self) -> bool {
-        let head = self.head.load(Acquire);
-        let tail = self.tail.load(Acquire);
-
-        head == tail
-    }
-
-    /// Steals half the tasks from self and place them into `dst`.
-    pub(super) unsafe fn steal(&self, dst: &Queue<T>) -> Option<Task<T>> {
-        let dst_tail = dst.tail.unsync_load();
-
-        // Steal the tasks into `dst`'s buffer. This does not yet expose the
-        // tasks in `dst`.
-        let mut n = self.steal2(dst, dst_tail);
-
-        if n == 0 {
-            // No tasks were stolen
-            return None;
-        }
-
-        // We are returning a task here
-        n -= 1;
-
-        let ret_pos = dst_tail.wrapping_add(n);
-        let ret_idx = ret_pos as usize & MASK;
-
-        let ret = dst.buffer[ret_idx].with(|ptr| ptr::read((*ptr).as_ptr()));
-
-        if n == 0 {
-            // The `dst` queue is empty, but a single task was stolen
-            return Some(ret);
-        }
-
-        // Synchronize with stealers
-        let dst_head = dst.head.load(Acquire);
-
-        assert!(dst_tail.wrapping_sub(dst_head) + n <= LOCAL_QUEUE_CAPACITY as u32);
-
-        // Make the stolen items available to consumers
-        dst.tail.store(dst_tail.wrapping_add(n), Release);
-
-        Some(ret)
-    }
-
-    unsafe fn steal2(&self, dst: &Queue<T>, dst_tail: u32) -> u32 {
-        loop {
-            let src_head = self.head.load(Acquire);
-            let src_tail = self.tail.load(Acquire);
-
-            // Number of available tasks to steal
-            let n = src_tail.wrapping_sub(src_head);
-            let n = n - n / 2;
-
-            if n == 0 {
-                return 0;
-            }
-
-            if n > LOCAL_QUEUE_CAPACITY as u32 / 2 {
-                atomic::spin_loop_hint();
-                // inconsistent, try again
-                continue;
-            }
-
-            // Track CausalCell causality checks. The check is deferred until
-            // the compare_and_swap claims ownership of the tasks.
-            let mut check = CausalCheck::default();
-
-            for i in 0..n {
-                // Compute the positions
-                let src_pos = src_head.wrapping_add(i);
-                let dst_pos = dst_tail.wrapping_add(i);
-
-                // Map to slots
-                let src_idx = src_pos as usize & MASK;
-                let dst_idx = dst_pos as usize & MASK;
-
-                // Read the task
-                let (task, ch) =
-                    self.buffer[src_idx].with_deferred(|ptr| ptr::read((*ptr).as_ptr()));
-
-                check.join(ch);
-
-                // Write the task to the new slot
-                dst.buffer[dst_idx].with_mut(|ptr| ptr::write((*ptr).as_mut_ptr(), task));
-            }
-
-            // Claim all of those tasks!
-            let actual = self
-                .head
-                .compare_and_swap(src_head, src_head.wrapping_add(n), Release);
-
-            if actual == src_head {
-                check.check();
-                return n;
-            }
-
-            atomic::spin_loop_hint();
-        }
-    }
-}
-
-impl<T> fmt::Debug for Queue<T> {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("local::Queue")
-            .field("head", &self.head)
-            .field("tail", &self.tail)
-            .field("buffer", &"[...]")
-            .finish()
-    }
-}

tokio/src/runtime/thread_pool/queue/mod.rs

@@ -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<T: 'static>(workers: usize) -> Vec<Worker<T>> {
-    let local: Vec<_> = (0..workers).map(|_| local::Queue::new()).collect();
-
-    let cluster = Arc::new(Cluster {
-        local: local.into_boxed_slice(),
-        global: global::Queue::new(),
-    });
-
-    (0..workers)
-        .map(|index| Worker::new(cluster.clone(), index))
-        .collect()
-}
-
-struct Cluster<T: 'static> {
-    /// per-worker local queues
-    local: Box<[local::Queue<T>]>,
-    global: global::Queue<T>,
-}
-
-impl<T: 'static> Drop for Cluster<T> {
-    fn drop(&mut self) {
-        // Drain all the queues
-        for queue in &self.local[..] {
-            while let Some(_) = unsafe { queue.pop() } {}
-        }
-
-        while let Some(_) = self.global.pop() {}
-    }
-}

tokio/src/runtime/thread_pool/queue/worker.rs

@@ -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<T: 'static> {
-    cluster: Arc<Cluster<T>>,
-    index: u16,
-    /// Task to pop next
-    next: Cell<Option<Task<T>>>,
-}
-
-impl<T: 'static> Worker<T> {
-    pub(super) fn new(cluster: Arc<Cluster<T>>, index: usize) -> Worker<T> {
-        Worker {
-            cluster,
-            index: index as u16,
-            next: Cell::new(None),
-        }
-    }
-
-    pub(crate) fn injector(&self) -> Inject<T> {
-        Inject::new(self.cluster.clone())
-    }
-
-    /// Returns `true` if the queue is closed
-    pub(crate) fn is_closed(&self) -> bool {
-        self.cluster.global.is_closed()
-    }
-
-    /// 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<T>) -> bool {
-        let prev = self.next.take();
-        let ret = prev.is_some();
-
-        if let Some(prev) = prev {
-            // safety: we guarantee that only one thread pushes to this local
-            // queue at a time.
-            unsafe {
-                self.local().push(prev, &self.cluster.global);
-            }
-        }
-
-        self.next.set(Some(task));
-
-        ret
-    }
-
-    pub(crate) fn push_yield(&self, task: Task<T>) {
-        unsafe { self.local().push(task, &self.cluster.global) }
-    }
-
-    /// Pops a task checking the local queue first.
-    pub(crate) fn pop_local_first(&self) -> Option<Task<T>> {
-        self.local_pop().or_else(|| self.cluster.global.pop())
-    }
-
-    /// Pops a task checking the global queue first.
-    pub(crate) fn pop_global_first(&self) -> Option<Task<T>> {
-        self.cluster.global.pop().or_else(|| self.local_pop())
-    }
-
-    /// Steals from other local queues.
-    ///
-    /// `start` specifies the queue from which to start stealing.
-    pub(crate) fn steal(&self, start: usize) -> Option<Task<T>> {
-        let num_queues = self.cluster.local.len();
-
-        for i in 0..num_queues {
-            let i = (start + i) % num_queues;
-
-            if i == self.index as usize {
-                continue;
-            }
-
-            // safety: we own the dst queue
-            let ret = unsafe { self.cluster.local[i].steal(self.local()) };
-
-            if ret.is_some() {
-                return ret;
-            }
-        }
-
-        None
-    }
-
-    /// An approximation of whether or not the queue is empty.
-    pub(crate) fn is_empty(&self) -> bool {
-        for local_queue in &self.cluster.local[..] {
-            if !local_queue.is_empty() {
-                return false;
-            }
-        }
-
-        self.cluster.global.is_empty()
-    }
-
-    fn local_pop(&self) -> Option<Task<T>> {
-        if let Some(task) = self.next.take() {
-            return Some(task);
-        }
-        // safety: we guarantee that only one thread pushes to this local queue
-        // at a time.
-        unsafe { self.local().pop() }
-    }
-
-    fn local(&self) -> &local::Queue<T> {
-        &self.cluster.local[self.index as usize]
-    }
-}
-
-impl<T: 'static> fmt::Debug for Worker<T> {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("queue::Worker")
-            .field("cluster", &"...")
-            .field("index", &self.index)
-            .finish()
-    }
-}

tokio/src/runtime/thread_pool/shared.rs

@@ -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<Self>,
-
-    /// 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>) {
-        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<Self>) {
-        // 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<Self>) {
-        // This stores the task with the owning worker. The worker is not
-        // notified. Instead, the worker will clean up the tasks "eventually".
-        //
-        self.pending_drop.push(task);
-    }
-
-    fn release_local(&self, task: &Task<Self>) {
-        // Get access to the Owned component. This function can only be called
-        // when on the worker.
-        unsafe {
-            let index = self.slices().index_of(self);
-            let owned = &mut *self.slices().owned()[index].get();
-
-            owned.release_task(task);
-        }
-    }
-
-    fn schedule(&self, task: Task<Self>) {
-        Self::schedule(self, task);
-    }
-}
-
-impl ScheduleSendOnly for Shared {}

tokio/src/runtime/thread_pool/shutdown.rs

@@ -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<oneshot::Sender<()>>,
-}
-
-#[derive(Debug)]
-pub(super) struct Receiver {
-    rx: oneshot::Receiver<()>,
-}
-
-pub(super) fn channel() -> (Sender, Receiver) {
-    let (tx, rx) = oneshot::channel();
-    let tx = Sender { tx: Arc::new(tx) };
-    let rx = Receiver { rx };
-
-    (tx, rx)
-}
-
-impl Receiver {
-    /// 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);
-    }
-}

tokio/src/runtime/thread_pool/slice.rs

@@ -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<CachePadded<Owned>>]>,
-
-    /// Submit work to the pool while *not* currently on a worker thread.
-    inject: queue::Inject<Shared>,
-
-    /// 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<F>(&self, future: F) -> JoinHandle<F::Output>
-    where
-        F: Future + Send + 'static,
-        F::Output: Send + 'static,
-    {
-        let (task, handle) = task::joinable(future);
-        self.schedule(task);
-        handle
-    }
-
-    fn inject_task(&self, task: Task<Shared>) {
-        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<Shared>) {
-        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>();
-
-        ((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<CachePadded<Owned>>] {
-        &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();
-    }
-}

tokio/src/runtime/thread_pool/spawner.rs

@@ -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<slice::Set>,
-}
-
-impl Spawner {
-    pub(super) fn new(workers: Arc<slice::Set>) -> Spawner {
-        Spawner { workers }
-    }
-
-    /// Spawns a future onto the thread pool
-    pub(crate) fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
-    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()
-    }
-}

tokio/src/runtime/thread_pool/tests/loom_pool.rs

@@ -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<Output = &'static str> {
-    gated2(false)
-}
-
-fn gated2(thread: bool) -> impl Future<Output = &'static str> {
-    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
-        }
-    })
-}

tokio/src/runtime/thread_pool/tests/loom_queue.rs

@@ -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<Noop>, rem: Rc<Cell<usize>>) {
-        let mut rem_local = PER_THREAD;
-
-        while rem.get() != 0 {
-            for _ in 0..3 {
-                if rem_local > 0 {
-                    q.push(val(0));
-                    rem_local -= 1;
-                }
-            }
-
-            // Try to work
-            while let Some(task) = q.pop_local_first() {
-                assert!(task.run(&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<Noop> {
-    let (task, _) = task::joinable(async move { num });
-    task
-}

tokio/src/runtime/thread_pool/tests/mod.rs

@@ -1,8 +0,0 @@
-#[cfg(loom)]
-mod loom_pool;
-
-#[cfg(loom)]
-mod loom_queue;
-
-#[cfg(not(loom))]
-mod queue;

tokio/src/runtime/thread_pool/tests/queue.rs

@@ -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<Noop> = &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<Noop>, queue::Worker<Noop>) {
-    let mut qs = queue::build(2);
-    (qs.remove(0), qs.remove(0))
-}
-
-// pretty big hack to track tasks
-use std::cell::RefCell;
-use std::collections::HashMap;
-thread_local! {
-    static TASKS: RefCell<HashMap<u32, task::JoinHandle<u32>>> = RefCell::new(HashMap::new())
-}
-
-fn val(num: u32) -> Task<Noop> {
-    let (task, join) = task::joinable(async move { num });
-    let prev = TASKS.with(|t| t.borrow_mut().insert(num, join));
-    assert!(prev.is_none());
-    task
-}
-
-fn num(task: Task<Noop>) -> u32 {
-    use futures::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<T, E>(r: Result<T, E>) {
-    assert!(r.is_ok())
-}

tokio/src/sync/notify.rs

@@ -306,10 +306,10 @@ fn notify_locked(waiters: &mut LinkedList<Waiter>, 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<Waiter>;
     type Target = Waiter;
 
-    fn to_raw(handle: *mut Waiter) -> NonNull<Waiter> {
-        debug_assert!(!handle.is_null());
-        unsafe { NonNull::new_unchecked(handle) }
+    fn as_raw(handle: &NonNull<Waiter>) -> NonNull<Waiter> {
+        *handle
     }
 
-    unsafe fn from_raw(ptr: NonNull<Waiter>) -> *mut Waiter {
-        ptr.as_ptr()
+    unsafe fn from_raw(ptr: NonNull<Waiter>) -> NonNull<Waiter> {
+        ptr
     }
 
     unsafe fn pointers(mut target: NonNull<Waiter>) -> NonNull<linked_list::Pointers<Waiter>> {

tokio/src/task/core.rs

@@ -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<T: Future> {
-    /// Hot task state data
-    pub(super) header: Header,
-
-    /// Either the future or output, depending on the execution stage.
-    pub(super) core: Core<T>,
-
-    /// Cold data
-    pub(super) trailer: Trailer,
-}
-
-/// The core of the task.
-///
-/// Holds the future or output, depending on the stage of execution.
-pub(super) struct Core<T: Future> {
-    stage: Stage<T>,
-}
-
-/// Crate public as this is also needed by the pool.
-#[repr(C)]
-pub(crate) struct Header {
-    /// Task state
-    pub(super) state: State,
-
-    /// Pointer to the executor owned by the task
-    pub(super) executor: CausalCell<Option<NonNull<()>>>,
-
-    /// 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<Option<NonNull<Header>>>,
-
-    /// Pointer to the previous task in the ownership list.
-    pub(crate) owned_prev: UnsafeCell<Option<NonNull<Header>>>,
-
-    /// 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<MaybeUninit<Option<Waker>>>,
-}
-
-/// Either the future or the output.
-enum Stage<T: Future> {
-    Running(Track<T>),
-    Finished(Track<super::Result<T::Output>>),
-    Consumed,
-}
-
-impl<T: Future> Cell<T> {
-    /// Allocates a new task cell, containing the header, trailer, and core
-    /// structures.
-    pub(super) fn new<S>(future: T, state: State) -> Box<Cell<T>>
-    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::<T, S>(),
-                future_causality: CausalCell::new(()),
-            },
-            core: Core {
-                stage: Stage::Running(Track::new(future)),
-            },
-            trailer: Trailer {
-                waker: CausalCell::new(MaybeUninit::new(None)),
-            },
-        })
-    }
-}
-
-impl<T: Future> Core<T> {
-    pub(super) fn transition_to_consumed(&mut self) {
-        self.stage = Stage::Consumed
-    }
-
-    pub(super) fn poll<S>(&mut self, header: &Header) -> Poll<T::Output>
-    where
-        S: Schedule,
-    {
-        let res = {
-            let future = match &mut self.stage {
-                Stage::Running(tracked) => tracked.get_mut(),
-                _ => unreachable!("unexpected stage"),
-            };
-
-            // The future is pinned within the task. The above state transition
-            // has ensured the safety of this action.
-            let future = unsafe { Pin::new_unchecked(future) };
-
-            // The waker passed into the `poll` function does not require a ref
-            // count increment.
-            let waker_ref = waker_ref::<T, S>(header);
-            let mut cx = Context::from_waker(&*waker_ref);
-
-            future.poll(&mut cx)
-        };
-
-        if res.is_ready() {
-            self.stage = Stage::Consumed;
-        }
-
-        res
-    }
-
-    pub(super) fn store_output(&mut self, output: super::Result<T::Output>) {
-        self.stage = Stage::Finished(Track::new(output));
-    }
-
-    pub(super) unsafe fn read_output(&mut self, dst: *mut Track<super::Result<T::Output>>) {
-        use std::mem;
-
-        dst.write(match mem::replace(&mut self.stage, Stage::Consumed) {
-            Stage::Finished(output) => output,
-            _ => unreachable!("unexpected state"),
-        });
-    }
-}
-
-impl Header {
-    pub(super) fn executor(&self) -> Option<NonNull<()>> {
-        unsafe { self.executor.with(|ptr| *ptr) }
-    }
-}

tokio/src/task/harness.rs

@@ -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<T: Future, S: 'static> {
-    cell: NonNull<Cell<T>>,
-    _p: PhantomData<S>,
-}
-
-impl<T, S> Harness<T, S>
-where
-    T: Future,
-    S: 'static,
-{
-    pub(super) unsafe fn from_raw(ptr: *mut ()) -> Harness<T, S> {
-        debug_assert!(!ptr.is_null());
-
-        Harness {
-            cell: NonNull::new_unchecked(ptr as *mut Cell<T>),
-            _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<T> {
-        unsafe { &mut self.cell.as_mut().core }
-    }
-}
-
-impl<T, S> Harness<T, S>
-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<NonNull<()>>) -> 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::<S>().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<T>,
-                    polled: bool,
-                }
-
-                impl<T: Future> Drop for Guard<'_, T> {
-                    fn drop(&mut self) {
-                        if !self.polled {
-                            self.core.transition_to_consumed();
-                        }
-                    }
-                }
-
-                let mut guard = Guard {
-                    core,
-                    polled: false,
-                };
-
-                let res = guard.core.poll::<S>(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<super::Result<T::Output>>,
-        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::<S>().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<NonNull<()>>,
-        join_interest: bool,
-        output: super::Result<T::Output>,
-    ) {
-        if join_interest {
-            // Store the output. The future has already been dropped
-            self.core().store_output(output);
-        }
-
-        let 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::<S>()
-                    .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::<S>().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<Option<Waker>> {
-        self.trailer().waker.with(|ptr| ptr.read())
-    }
-
-    unsafe fn to_task(&self) -> Task<S> {
-        let ptr = self.cell.as_ptr() as *mut Header;
-        Task::from_raw(NonNull::new_unchecked(ptr))
-    }
-}

tokio/src/task/list.rs

@@ -1,96 +0,0 @@
-use crate::task::{Header, Task};
-
-use std::fmt;
-use std::marker::PhantomData;
-use std::ptr::NonNull;
-
-pub(crate) struct OwnedList<T: 'static> {
-    head: Option<NonNull<Header>>,
-    _p: PhantomData<T>,
-}
-
-impl<T: 'static> OwnedList<T> {
-    pub(crate) fn new() -> OwnedList<T> {
-        OwnedList {
-            head: None,
-            _p: PhantomData,
-        }
-    }
-
-    pub(crate) fn insert(&mut self, task: &Task<T>) {
-        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<T>) {
-        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<T>) -> 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<T: 'static> fmt::Debug for OwnedList<T> {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("OwnedList").finish()
-    }
-}

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<Scheduler>,
+        /// Current scheduler tick
+        tick: Cell<u8>,
+
+        /// State available from thread-local
+        context: Context,
     }
 }
 
-#[derive(Debug)]
-struct Scheduler {
-    tick: Cell<u8>,
+/// State available from the thread-local
+struct Context {
+    /// Owned task set and local run queue
+    tasks: RefCell<Tasks>,
 
-    queues: MpscQueues<Self>,
+    /// State shared between threads.
+    shared: Arc<Shared>,
+}
 
-    /// 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<Task<Arc<Shared>>>,
+
+    /// Local run queue sender and receiver.
+    queue: VecDeque<task::Notified<Arc<Shared>>>,
+}
+
+/// LocalSet state shared between threads.
+struct Shared {
+    /// Remote run queue sender
+    queue: Mutex<VecDeque<task::Notified<Arc<Shared>>>>,
+
+    /// Wake the `LocalSet` task
     waker: AtomicWaker,
 }
 
 pin_project! {
     #[derive(Debug)]
-    struct LocalFuture<F> {
-        scheduler: Rc<Scheduler>,
+    struct RunUntil<'a, F> {
+        local_set: &'a LocalSet,
         #[pin]
         future: F,
     }
 }
 
-thread_local! {
-    static CURRENT_TASK_SET: Cell<Option<NonNull<Scheduler>>> = 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<task::Notified<Arc<Shared>>> {
+        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<T>(&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<Self::Output> {
-        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<Self::Output> {
+        // 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<F: Future> Future for LocalFuture<F> {
-    type Output = F::Output;
+impl<T: Future> Future for RunUntil<'_, T> {
+    type Output = T::Output;
 
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        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<Self::Output> {
+        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<F: Future> Future for LocalFuture<F> {
     }
 }
 
-// === impl Scheduler ===
-
-impl Schedule for Scheduler {
-    fn bind(&self, task: &Task<Self>) {
-        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<Arc<Self>>) {
+        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<Self>) {
-        // This will be called when dropping the local runtime.
-        self.queues.release_remote(task);
-    }
-
-    fn release_local(&self, task: &Task<Self>) {
-        debug_assert!(self.is_current());
-        unsafe {
-            self.queues.release_local(task);
-        }
-    }
-
-    fn schedule(&self, task: Task<Self>) {
-        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<F>(&self, f: impl FnOnce() -> F) -> F {
-        struct Entered<'a> {
-            current: &'a Cell<Option<NonNull<Scheduler>>>,
-        }
-
-        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<Shared> {
+    fn bind(task: Task<Self>) -> Arc<Shared> {
+        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<Self>) -> Option<Task<Self>> {
+        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<Self>) {
+        Shared::schedule(self, task);
     }
 }

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<S: 'static> {
-        raw: RawTask,
-        _p: PhantomData<S>,
-    }
-
-    unsafe impl<S: ScheduleSendOnly + 'static> Send for Task<S> {}
-
-    /// Task result sent back
-    pub(crate) type Result<T> = std::result::Result<T, JoinError>;
-
-    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<Self>);
-
-        /// The task has completed work and is ready to be released. The scheduler
-        /// is free to drop it whenever.
-        fn release(&self, task: Task<Self>);
-
-        /// The has been completed by the executor it was bound to.
-        fn release_local(&self, task: &Task<Self>);
-
-        /// Schedule the task
-        fn schedule(&self, task: Task<Self>);
-    }
-
-    /// 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<T, S>(task: T) -> (Task<S>, JoinHandle<T::Output>)
-    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<T, S>(task: T) -> (Task<S>, JoinHandle<T::Output>)
-        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<S: 'static> Task<S> {
-        pub(crate) unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
-            Task {
-                raw: RawTask::from_raw(ptr),
-                _p: PhantomData,
-            }
-        }
-
-        pub(crate) fn header(&self) -> &Header {
-            self.raw.header()
-        }
-
-        pub(crate) fn into_raw(self) -> NonNull<Header> {
-            let raw = self.raw.into_raw();
-            mem::forget(self);
-            raw
-        }
-    }
-
-    impl<S: Schedule> Task<S> {
-        /// Returns `self` when the task needs to be immediately re-scheduled
-        pub(crate) fn run<F>(self, mut executor: F) -> Option<Self>
-        where
-            F: FnMut() -> Option<NonNull<S>>,
-        {
-            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<S: 'static> Drop for Task<S> {
-        fn drop(&mut self) {
-            self.raw.drop_task();
-        }
-    }
-
-    impl<S> fmt::Debug for Task<S> {
-        fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-            fmt.debug_struct("Task").finish()
-        }
-    }
-}

tokio/src/task/queue.rs

@@ -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<S: 'static> {
-    /// List of all active tasks spawned onto this executor.
-    ///
-    /// # Safety
-    ///
-    /// Must only be accessed from the primary thread
-    owned_tasks: UnsafeCell<OwnedList<S>>,
-
-    /// 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<VecDeque<Task<S>>>,
-
-    /// Remote run queue.
-    ///
-    /// Tasks notified from another thread are pushed into this queue.
-    remote_queue: Mutex<RemoteQueue<S>>,
-
-    /// Tasks pending drop
-    pending_drop: TransferStack<S>,
-}
-
-pub(crate) struct RemoteQueue<S: 'static> {
-    /// FIFO list of tasks
-    queue: VecDeque<Task<S>>,
-
-    /// `true` when a task can be pushed into the queue, `false` otherwise.
-    open: bool,
-}
-
-// === impl Queues ===
-
-impl<S> MpscQueues<S>
-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<S>) {
-        (*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<S>) {
-        (*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<S>) {
-        (*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<S>> {
-        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<S>) {
-        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<Task<S>> {
-        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<Task<S>> {
-        (*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<Task<S>> {
-        // 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<S> fmt::Debug for MpscQueues<S> {
-    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<S> RemoteQueue<S>
-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<S>, spawn: bool) {
-        if !spawn || self.open {
-            self.queue.push_back(task);
-        } else {
-            task.shutdown();
-        }
-    }
-}
-
-impl<S> fmt::Debug for RemoteQueue<S> {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("RemoteQueue")
-            .field("queue", &self.queue)
-            .field("open", &self.open)
-            .finish()
-    }
-}

tokio/src/task/raw.rs

@@ -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<Header>,
-}
-
-pub(super) struct Vtable {
-    /// Poll the future
-    pub(super) poll: unsafe fn(*mut (), &mut dyn FnMut() -> Option<NonNull<()>>) -> bool,
-
-    /// The task handle has been dropped and the join waker needs to be dropped
-    /// or the task struct needs to be deallocated
-    pub(super) drop_task: unsafe fn(*mut ()),
-
-    /// Read the task output
-    pub(super) read_output: unsafe fn(*mut (), *mut (), Snapshot),
-
-    /// Store the join handle's waker
-    ///
-    /// Returns a snapshot of the state **after** the transition
-    pub(super) store_join_waker: unsafe fn(*mut (), &Waker) -> Snapshot,
-
-    /// Replace the join handle's waker
-    ///
-    /// Returns a snapshot of the state **after** the transition
-    pub(super) swap_join_waker: unsafe fn(*mut (), &Waker, Snapshot) -> Snapshot,
-
-    /// The join handle has been dropped
-    pub(super) drop_join_handle_slow: unsafe fn(*mut ()),
-
-    /// The task is being canceled
-    pub(super) cancel: unsafe fn(*mut (), bool),
-}
-
-/// Get the vtable for the requested `T` and `S` generics.
-pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
-    &Vtable {
-        poll: poll::<T, S>,
-        drop_task: drop_task::<T, S>,
-        read_output: read_output::<T, S>,
-        store_join_waker: store_join_waker::<T, S>,
-        swap_join_waker: swap_join_waker::<T, S>,
-        drop_join_handle_slow: drop_join_handle_slow::<T, S>,
-        cancel: cancel::<T, S>,
-    }
-}
-
-cfg_rt_util! {
-    impl RawTask {
-        pub(super) fn new_joinable_local<T, S>(task: T) -> RawTask
-        where
-            T: Future + 'static,
-            S: Schedule,
-        {
-            RawTask::new::<_, S>(task, State::new_joinable())
-        }
-    }
-}
-
-impl RawTask {
-    pub(super) fn new_joinable<T, S>(task: T) -> RawTask
-    where
-        T: Future + Send + 'static,
-        S: ScheduleSendOnly,
-    {
-        RawTask::new::<_, S>(task, State::new_joinable())
-    }
-
-    fn new<T, S>(task: T, state: State) -> RawTask
-    where
-        T: Future + 'static,
-        S: Schedule,
-    {
-        let ptr = Box::into_raw(Cell::new::<S>(task, state));
-        let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) };
-
-        RawTask { ptr }
-    }
-
-    pub(super) unsafe fn from_raw(ptr: NonNull<Header>) -> 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<Header> {
-        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<NonNull<()>>) -> 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<T: Future, S: Schedule>(
-    ptr: *mut (),
-    executor: &mut dyn FnMut() -> Option<NonNull<()>>,
-) -> bool {
-    let harness = Harness::<T, S>::from_raw(ptr);
-    harness.poll(executor)
-}
-
-unsafe fn drop_task<T: Future, S: Schedule>(ptr: *mut ()) {
-    let harness = Harness::<T, S>::from_raw(ptr);
-    harness.drop_task();
-}
-
-unsafe fn read_output<T: Future, S: Schedule>(ptr: *mut (), dst: *mut (), state: Snapshot) {
-    let harness = Harness::<T, S>::from_raw(ptr);
-    harness.read_output(dst as *mut Track<super::Result<T::Output>>, state);
-}
-
-unsafe fn store_join_waker<T: Future, S: Schedule>(ptr: *mut (), waker: &Waker) -> Snapshot {
-    let harness = Harness::<T, S>::from_raw(ptr);
-    harness.store_join_waker(waker)
-}
-
-unsafe fn swap_join_waker<T: Future, S: Schedule>(
-    ptr: *mut (),
-    waker: &Waker,
-    prev: Snapshot,
-) -> Snapshot {
-    let harness = Harness::<T, S>::from_raw(ptr);
-    harness.swap_join_waker(waker, prev)
-}
-
-unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: *mut ()) {
-    let harness = Harness::<T, S>::from_raw(ptr);
-    harness.drop_join_handle_slow()
-}
-
-unsafe fn cancel<T: Future, S: Schedule>(ptr: *mut (), from_queue: bool) {
-    let harness = Harness::<T, S>::from_raw(ptr);
-    harness.cancel(from_queue)
-}

tokio/src/task/stack.rs

@@ -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<T: 'static> {
-    head: AtomicPtr<Header>,
-    _p: PhantomData<T>,
-}
-
-impl<T: 'static> TransferStack<T> {
-    pub(crate) fn new() -> TransferStack<T> {
-        TransferStack {
-            head: AtomicPtr::new(ptr::null_mut()),
-            _p: PhantomData,
-        }
-    }
-
-    pub(crate) fn push(&self, task: Task<T>) {
-        unsafe {
-            let task = task.into_raw();
-
-            let next = (*task.as_ref().queue_next.get()) as usize;
-
-            // At this point, the queue_next field may also be used to track
-            // whether or not the task must drop the join waker.
-            debug_assert_eq!(0, next & !1);
-
-            // We don't care about any memory associated w/ setting the `head`
-            // field, just the current value.
-            let mut curr = self.head.load(Relaxed);
-
-            loop {
-                *task.as_ref().queue_next.get() = (next | curr as usize) as *const _;
-
-                let res =
-                    self.head
-                        .compare_exchange(curr, task.as_ptr() as *mut _, Release, Relaxed);
-
-                match res {
-                    Ok(_) => return,
-                    Err(actual) => {
-                        curr = actual;
-                    }
-                }
-            }
-        }
-    }
-
-    pub(crate) fn drain(&self) -> impl Iterator<Item = Task<T>> {
-        struct Iter<T: 'static>(*mut Header, PhantomData<T>);
-
-        impl<T: 'static> Iterator for Iter<T> {
-            type Item = Task<T>;
-
-            fn next(&mut self) -> Option<Task<T>> {
-                let task = NonNull::new(self.0)?;
-
-                unsafe {
-                    let next = *task.as_ref().queue_next.get() as usize;
-
-                    // remove the data bit
-                    self.0 = (next & !1) as *mut _;
-
-                    Some(Task::from_raw(task))
-                }
-            }
-        }
-
-        impl<T: 'static> Drop for Iter<T> {
-            fn drop(&mut self) {
-                use std::process;
-
-                if !self.0.is_null() {
-                    // we have bugs
-                    process::abort();
-                }
-            }
-        }
-
-        let ptr = self.head.swap(ptr::null_mut(), Acquire);
-        Iter(ptr, PhantomData)
-    }
-}

tokio/src/task/state.rs

@@ -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<Snapshot> {
-        let mut prev = self.load();
-
-        loop {
-            if !prev.is_active() {
-                return None;
-            }
-
-            let mut next = prev;
-
-            // Use the running flag to signal cancellation
-            if prev.is_running() {
-                next.0 -= RUNNING;
-                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<Snapshot, Snapshot> {
-        const MASK: usize = COMPLETE | CANCELLED;
-
-        let mut prev = self.val.load(Acquire);
-
-        loop {
-            // Once the complete bit is set, it is never unset.
-            if prev & MASK != 0 {
-                return Err(Snapshot(prev));
-            }
-
-            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()
-    }
-}

tokio/src/task/tests/loom.rs

@@ -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<Output = &'static str> {
-    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<T: Future + 'static>(join_handle: T) -> loom::thread::JoinHandle<T::Output> {
-    thread::spawn(move || block_on(join_handle))
-}
-
-/// Spawn a thread to wait on the join handle using two tasks. First, poll the
-/// join handle on the first task. If the join handle is not ready, then use a
-/// second task to wait on it.
-fn join_two_tasks<T: Future + Unpin + 'static>(
-    join_handle: T,
-) -> loom::thread::JoinHandle<T::Output> {
-    use futures::future::poll_fn;
-    use std::task::Poll;
-
-    // Join handle
-    thread::spawn(move || {
-        let mut join_handle = Some(join_handle);
-        block_on(poll_fn(move |cx| {
-            use std::pin::Pin;
-
-            let res = Pin::new(join_handle.as_mut().unwrap()).poll(cx);
-
-            if res.is_ready() {
-                return res;
-            }
-
-            // Yes, we are nesting
-            Poll::Ready(block_on(join_handle.take().unwrap()))
-        }))
-    })
-}

tokio/src/task/tests/mod.rs

@@ -1,5 +0,0 @@
-#[cfg(loom)]
-mod loom;
-
-#[cfg(not(loom))]
-mod task;

tokio/src/task/tests/task.rs

@@ -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::<Header>() <= 8 * size_of::<*const ()>());
-}
-
-#[test]
-fn create_complete_drop() {
-    let (tx, rx) = mpsc::channel();
-
-    let (task, did_drop) = track_drop(async move {
-        tx.send(1).unwrap();
-    });
-
-    let (task, _) = task::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());
-}

tokio/src/tests/backoff.rs

@@ -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<Output = ()> {
-    Backoff(n, false)
-}
-
-/// Back off, but clone the waker each time
-pub(crate) fn backoff_clone(n: usize) -> impl Future<Output = ()> {
-    Backoff(n, true)
-}
-
-impl Future for Backoff {
-    type Output = ();
-
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        if self.0 == 0 {
-            return Poll::Ready(());
-        }
-
-        self.0 -= 1;
-        if self.1 {
-            cx.waker().clone().wake();
-        } else {
-            cx.waker().wake_by_ref();
-        }
-        Poll::Pending
-    }
-}

tokio/src/tests/loom_schedule.rs

@@ -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<VecDeque<Option<Task<Self>>>>,
-}
-
-impl LoomSchedule {
-    pub(crate) fn new() -> LoomSchedule {
-        LoomSchedule {
-            notify: Notify::new(),
-            pending: Mutex::new(VecDeque::new()),
-        }
-    }
-
-    pub(crate) fn push_task(&self, task: Task<Self>) {
-        self.schedule(task);
-    }
-
-    pub(crate) fn recv(&self) -> Option<Task<Self>> {
-        loop {
-            if let Some(task) = self.pending.lock().unwrap().pop_front() {
-                return task;
-            }
-
-            self.notify.wait();
-        }
-    }
-}
-
-impl Schedule for LoomSchedule {
-    fn bind(&self, _task: &Task<Self>) {}
-
-    fn release(&self, task: Task<Self>) {
-        self.release_local(&task);
-    }
-
-    fn release_local(&self, _task: &Task<Self>) {
-        self.pending.lock().unwrap().push_back(None);
-        self.notify.notify();
-    }
-
-    fn schedule(&self, task: Task<Self>) {
-        self.pending.lock().unwrap().push_back(Some(task));
-        self.notify.notify();
-    }
-}
-
-impl ScheduleSendOnly for LoomSchedule {}

tokio/src/tests/mock_schedule.rs

@@ -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<Inner>,
-}
-
-pub(crate) struct Noop;
-pub(crate) static NOOP_SCHEDULE: Noop = Noop;
-
-struct Inner {
-    calls: VecDeque<Call>,
-    pending_run: VecDeque<Task<Mock>>,
-    pending_drop: VecDeque<Task<Mock>>,
-}
-
-unsafe impl Send for Inner {}
-unsafe impl Sync for Inner {}
-
-#[derive(Debug, Eq, PartialEq)]
-enum Call {
-    Bind(*const Header),
-    Release,
-    ReleaseLocal,
-    Schedule,
-}
-
-pub(crate) fn mock() -> Mock {
-    Mock {
-        inner: Mutex::new(Inner {
-            calls: VecDeque::new(),
-            pending_run: VecDeque::new(),
-            pending_drop: VecDeque::new(),
-        }),
-    }
-}
-
-impl Mock {
-    pub(crate) fn bind(self, task: &Task<Mock>) -> Self {
-        self.push(Call::Bind(task.header() as *const _));
-        self
-    }
-
-    pub(crate) fn release(self) -> Self {
-        self.push(Call::Release);
-        self
-    }
-
-    pub(crate) fn release_local(self) -> Self {
-        self.push(Call::ReleaseLocal);
-        self
-    }
-
-    pub(crate) fn schedule(self) -> Self {
-        self.push(Call::Schedule);
-        self
-    }
-
-    pub(crate) fn next_pending_run(&self) -> Option<Task<Self>> {
-        self.inner.lock().unwrap().pending_run.pop_front()
-    }
-
-    pub(crate) fn next_pending_drop(&self) -> Option<Task<Self>> {
-        self.inner.lock().unwrap().pending_drop.pop_front()
-    }
-
-    fn push(&self, call: Call) {
-        self.inner.lock().unwrap().calls.push_back(call);
-    }
-
-    fn next(&self, name: &str) -> Call {
-        self.inner
-            .lock()
-            .unwrap()
-            .calls
-            .pop_front()
-            .expect(&format!("received `{}`, but none expected", name))
-    }
-}
-
-impl Schedule for Mock {
-    fn bind(&self, task: &Task<Self>) {
-        match self.next("bind") {
-            Call::Bind(ptr) => {
-                assert!(ptr.eq(&(task.header() as *const _)));
-            }
-            call => panic!("expected `Bind`, was {:?}", call),
-        }
-    }
-
-    fn release(&self, task: Task<Self>) {
-        match self.next("release") {
-            Call::Release => {
-                self.inner.lock().unwrap().pending_drop.push_back(task);
-            }
-            call => panic!("expected `Release`, was {:?}", call),
-        }
-    }
-
-    fn release_local(&self, _task: &Task<Self>) {
-        assert_eq!(Call::ReleaseLocal, self.next("release_local"));
-    }
-
-    fn schedule(&self, task: Task<Self>) {
-        self.inner.lock().unwrap().pending_run.push_back(task);
-        assert_eq!(Call::Schedule, self.next("schedule"));
-    }
-}
-
-impl 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<Self>) {}
-
-    fn release(&self, _task: Task<Self>) {}
-
-    fn release_local(&self, _task: &Task<Self>) {}
-
-    fn schedule(&self, _task: Task<Self>) {}
-}
-
-impl ScheduleSendOnly for Noop {}

tokio/src/tests/mod.rs

@@ -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;

tokio/src/tests/track_drop.rs

@@ -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>(T, Arc<AtomicBool>);
-
-#[derive(Debug)]
-pub(crate) struct DidDrop(Arc<AtomicBool>, Arc<AtomicBool>);
-
-pub(crate) fn track_drop<T: Future>(
-    future: T,
-) -> (impl Future<Output = TrackDrop<T::Output>>, DidDrop) {
-    let did_drop_future = Arc::new(AtomicBool::new(false));
-    let did_drop_output = Arc::new(AtomicBool::new(false));
-    let did_drop = DidDrop(did_drop_future.clone(), did_drop_output.clone());
-
-    let future = async move { TrackDrop(future.await, did_drop_output) };
-
-    let future = TrackDrop(future, did_drop_future);
-
-    (future, did_drop)
-}
-
-impl<T> TrackDrop<T> {
-    pub(crate) fn get_ref(&self) -> &T {
-        &self.0
-    }
-}
-
-impl<T: Future> Future for TrackDrop<T> {
-    type Output = T::Output;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let me = unsafe { Pin::map_unchecked_mut(self, |x| &mut x.0) };
-        me.poll(cx)
-    }
-}
-
-impl<T> Drop for TrackDrop<T> {
-    fn drop(&mut self) {
-        self.1.store(true, SeqCst);
-    }
-}
-
-impl DidDrop {
-    pub(crate) fn did_drop_future(&self) -> bool {
-        self.0.load(SeqCst)
-    }
-
-    pub(crate) fn did_drop_output(&self) -> bool {
-        self.1.load(SeqCst)
-    }
-}

tokio/src/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<Self::Target>;
+    /// Convert the handle to a raw pointer without consuming the handle
+    fn as_raw(handle: &Self::Handle) -> NonNull<Self::Target>;
 
     /// Convert the raw pointer to a handle
     unsafe fn from_raw(ptr: NonNull<Self::Target>) -> Self::Handle;
@@ -79,7 +78,9 @@ impl<T: Link> LinkedList<T> {
 
     /// 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<T: Link> LinkedList<T> {
     ///
     /// 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<T::Target>) -> bool {
+    pub(crate) unsafe fn remove(&mut self, node: NonNull<T::Target>) -> Option<T::Handle> {
         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<T: Link> LinkedList<T> {
         } 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<T: Link> LinkedList<T> {
         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<NonNull<T::Target>>,
+        _p: PhantomData<&'a T>,
+    }
+
+    impl<T: Link> LinkedList<T> {
+        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<Entry> {
+        fn as_raw(handle: &Pin<&'_ Entry>) -> NonNull<Entry> {
             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<usize>) {
@@ -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);
                     }
                 }
             }

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;

tokio/src/util/try_lock.rs

@@ -20,13 +20,26 @@ unsafe impl<T: Send> Sync for TryLock<T> {}
 
 unsafe impl<T: Sync> Sync for LockGuard<'_, T> {}
 
-impl<T> TryLock<T> {
-    /// Create a new `TryLock`
-    pub(crate) fn new(data: T) -> TryLock<T> {
+macro_rules! new {
+    ($data:ident) => {
         TryLock {
             locked: AtomicBool::new(false),
-            data: UnsafeCell::new(data),
+            data: UnsafeCell::new($data),
         }
+    };
+}
+
+impl<T> TryLock<T> {
+    #[cfg(not(loom))]
+    /// Create a new `TryLock`
+    pub(crate) const fn new(data: T) -> TryLock<T> {
+        new!(data)
+    }
+
+    #[cfg(loom)]
+    /// Create a new `TryLock`
+    pub(crate) fn new(data: T) -> TryLock<T> {
+        new!(data)
     }
 
     /// Attempt to acquire lock

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<Self>);
+
+    /// Wake by reference
+    fn wake_by_ref(arc_self: &Arc<Self>);
+}
+
+/// A `Waker` that is only valid for a given lifetime.
+#[derive(Debug)]
+pub(crate) struct WakerRef<'a> {
+    waker: ManuallyDrop<Waker>,
+    _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<impl Wake>`.
+pub(crate) fn waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_> {
+    let ptr = &**wake as *const _ as *const ();
+
+    let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>())) };
+
+    WakerRef {
+        waker: ManuallyDrop::new(waker),
+        _p: PhantomData,
+    }
+}
+
+fn waker_vtable<W: Wake>() -> &'static RawWakerVTable {
+    &RawWakerVTable::new(
+        clone_arc_raw::<W>,
+        wake_arc_raw::<W>,
+        wake_by_ref_arc_raw::<W>,
+        drop_arc_raw::<W>,
+    )
+}
+
+unsafe fn inc_ref_count<T: Wake>(data: *const ()) {
+    // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
+    let arc = ManuallyDrop::new(Arc::<T>::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<T: Wake>(data: *const ()) -> RawWaker {
+    inc_ref_count::<T>(data);
+    RawWaker::new(data, waker_vtable::<T>())
+}
+
+unsafe fn wake_arc_raw<T: Wake>(data: *const ()) {
+    let arc: Arc<T> = Arc::from_raw(data as *const T);
+    Wake::wake(arc);
+}
+
+// used by `waker_ref`
+unsafe fn wake_by_ref_arc_raw<T: Wake>(data: *const ()) {
+    // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
+    let arc = ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
+    Wake::wake_by_ref(&arc);
+}
+
+unsafe fn drop_arc_raw<T: Wake>(data: *const ()) {
+    drop(Arc::<T>::from_raw(data as *const T))
+}

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();
     }
 }

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]