tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-24 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

thread-pool: in-place blocking with new scheduler (#1681)

Browse Source 
The initial new scheduler PR omitted in-place blocking
support. This patch brings it back.
This commit is contained in:
Jon Gjengset
2019-10-30 08:58:49 -07:00
committed by Carl Lerche
parent e3261440e5
commit 109fd3086b
17 changed files with 496 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ci/azure-loom.yml
+1 -1
View File
@@ -10,7 +10,7 @@ jobs:
rust_version: ${{ parameters.rust }}
- ${{ each crate in parameters.crates }}:
- script: RUSTFLAGS="--cfg loom" cargo test --lib --release
- script: RUSTFLAGS="--cfg loom" cargo test --lib --release -- --test-threads=1 --nocapture
env:
LOOM_MAX_PREEMPTIONS: 2
CI: 'True'
tokio/src/executor/task/harness.rs
+6 -4
View File
@@ -51,7 +51,7 @@ where
/// Panics raised while polling the future are handled.
///
/// Returns `true` if the task needs to be scheduled again
pub(super) fn poll(mut self, executor: NonNull<S>) -> bool {
pub(super) fn poll(mut self, executor: &mut dyn FnMut() -> Option<NonNull<S>>) -> bool {
use std::panic;
// Transition the task to the running state.
@@ -81,6 +81,7 @@ where
// 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.as_ref().bind(&task);
header.executor.with_mut(|ptr| *ptr = Some(executor));
}
@@ -393,7 +394,7 @@ where
fn complete(
mut self,
executor: NonNull<S>,
executor: &mut dyn FnMut() -> Option<NonNull<S>>,
join_interest: bool,
output: super::Result<T::Output>,
) {
@@ -402,15 +403,16 @@ where
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 Some(executor) == bound_executor {
if executor.is_some() && executor == bound_executor {
unsafe {
// perform a local release
let task = ManuallyDrop::new(self.to_task());
executor.as_ref().release_local(&task);
executor.as_ref().unwrap().as_ref().release_local(&task);
if self.transition_to_released(join_interest).is_final_ref() {
self.dealloc();
tokio/src/executor/task/mod.rs
+1 -1
View File
@@ -99,7 +99,7 @@ impl<S: 'static> Task<S> {
impl<S: Schedule> Task<S> {
/// Returns `self` when the task needs to be immediately re-scheduled
pub(crate) fn run(self, executor: NonNull<S>) -> Option<Self> {
pub(crate) fn run(self, executor: &mut dyn FnMut() -> Option<NonNull<S>>) -> Option<Self> {
if unsafe { self.raw.poll(executor) } {
Some(self)
} else {
tokio/src/executor/task/raw.rs
+6 -3
View File
@@ -15,7 +15,7 @@ pub(super) struct RawTask<S: 'static> {
pub(super) struct Vtable<S: 'static> {
/// Poll the future
pub(super) poll: unsafe fn(*mut (), NonNull<S>) -> bool,
pub(super) poll: unsafe fn(*mut (), &mut dyn FnMut() -> Option<NonNull<S>>) -> bool,
/// The task handle has been dropped and the join waker needs to be dropped
/// or the task struct needs to be deallocated
@@ -101,7 +101,7 @@ impl<S> RawTask<S> {
/// Safety: mutual exclusion is required to call this function.
///
/// Returns `true` if the task needs to be scheduled again.
pub(super) unsafe fn poll(self, executor: NonNull<S>) -> bool {
pub(super) unsafe fn poll(self, executor: &mut dyn FnMut() -> Option<NonNull<S>>) -> bool {
// Get the vtable without holding a ref to the meta struct. This is done
// because a mutable reference to the task is passed into the poll fn.
let vtable = self.header().vtable;
@@ -150,7 +150,10 @@ impl<S: 'static> Clone for RawTask<S> {
impl<S: 'static> Copy for RawTask<S> {}
unsafe fn poll<T: Future, S: Schedule>(ptr: *mut (), executor: NonNull<S>) -> bool {
unsafe fn poll<T: Future, S: Schedule>(
ptr: *mut (),
executor: &mut dyn FnMut() -> Option<NonNull<S>>,
) -> bool {
let harness = Harness::<T, S>::from_raw(ptr);
harness.poll(executor)
}
tokio/src/executor/task/tests/loom.rs
+8 -8
View File
@@ -15,7 +15,7 @@ fn create_drop_join_handle() {
let (task, join_handle) = task::joinable(async { "hello" });
let schedule = LoomSchedule::new();
let schedule = From::from(&schedule);
let schedule = &mut || Some(From::from(&schedule));
let th = thread::spawn(move || {
drop(join_handle);
@@ -37,7 +37,7 @@ fn poll_drop_handle_then_drop() {
let (task, mut join_handle) = task::joinable(async { "hello" });
let schedule = LoomSchedule::new();
let schedule = From::from(&schedule);
let schedule = &mut || Some(From::from(&schedule));
let th = thread::spawn(move || {
block_on(poll_fn(|cx| {
@@ -58,7 +58,7 @@ fn join_output() {
let (task, join_handle) = task::joinable(async { "hello world" });
let schedule = LoomSchedule::new();
let schedule = From::from(&schedule);
let schedule = &mut || Some(From::from(&schedule));
let th = thread::spawn(move || {
let out = assert_ok!(block_on(join_handle));
@@ -115,12 +115,12 @@ fn release_remote() {
// Join handle
let th = join_one_task(join_handle);
let task = match task.run(From::from(&s1)) {
let task = match task.run(&mut || Some(From::from(&s1))) {
Some(task) => task,
None => s1.recv().expect("released!"),
};
assert_none!(task.run(From::from(&s2)));
assert_none!(task.run(&mut || Some(From::from(&s2))));
assert_none!(s1.recv());
assert_ok!(th.join().unwrap());
@@ -152,7 +152,7 @@ fn shutdown_from_list_after_poll() {
// Join handle
let th = join_two_tasks(join_handle);
match task.run(From::from(&s1)) {
match task.run(&mut || Some(From::from(&s1))) {
Some(task) => {
// always drain the list before calling shutdown on tasks
list.shutdown();
@@ -184,7 +184,7 @@ fn shutdown_from_queue_after_poll() {
// Join handle
let th = join_two_tasks(join_handle);
let task = match task.run(From::from(&s1)) {
let task = match task.run(&mut || Some(From::from(&s1))) {
Some(task) => task,
None => assert_some!(s1.recv()),
};
@@ -239,7 +239,7 @@ fn work(schedule: &LoomSchedule) {
let mut task = Some(task);
while let Some(t) = task.take() {
task = t.run(From::from(schedule));
task = t.run(&mut || Some(From::from(schedule)));
}
}
}
tokio/src/executor/task/tests/task.rs
+32 -32
View File
@@ -28,7 +28,7 @@ fn create_complete_drop() {
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
let mock = &mut || Some(From::from(&mock));
// Nothing is returned
assert!(task.run(mock).is_none());
@@ -53,7 +53,7 @@ fn create_yield_complete_drop() {
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
let mock = &mut || Some(From::from(&mock));
// Task is returned
let task = assert_some!(task.run(mock));
@@ -83,7 +83,7 @@ fn create_clone_yield_complete_drop() {
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
let mock = &mut || Some(From::from(&mock));
// Task is returned
let task = assert_some!(task.run(mock));
@@ -111,7 +111,7 @@ fn create_wake_drop() {
let mock = mock().bind(&task).schedule().release_local();
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
assert_none!(mock.next_pending_run());
// The future was **not** dropped.
@@ -121,7 +121,7 @@ fn create_wake_drop() {
let task = assert_some!(mock.next_pending_run());
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
// The future was dropped.
assert!(did_drop.did_drop_future());
@@ -143,7 +143,7 @@ fn notify_complete() {
let task = task::background(task);
let mock = mock().bind(&task).release_local();
let mock = From::from(&mock);
let mock = &mut || Some(From::from(&mock));
assert_none!(task.run(mock));
assert!(did_drop.did_drop_future());
@@ -165,9 +165,9 @@ fn complete_on_second_schedule_obj() {
let mock2 = mock().bind(&task).release();
// Task is returned
let task = assert_some!(task.run(From::from(&mock2)));
let task = assert_some!(task.run(&mut || Some(From::from(&mock2))));
assert_none!(task.run(From::from(&mock1)));
assert_none!(task.run(&mut || Some(From::from(&mock1))));
// The message was sent
assert!(rx.try_recv().is_ok());
@@ -187,7 +187,7 @@ fn join_task_immediate_drop_handle() {
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(task.run(&mut || Some(From::from(&mock))).is_none());
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
@@ -202,7 +202,7 @@ fn join_task_immediate_complete_1() {
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(task.run(&mut || Some(From::from(&mock))).is_none());
assert!(did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
@@ -227,7 +227,7 @@ fn join_task_immediate_complete_2() {
assert_pending!(handle.poll());
assert!(task.run(From::from(&mock)).is_none());
assert!(task.run(&mut || Some(From::from(&mock))).is_none());
assert!(did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
@@ -253,14 +253,14 @@ fn join_task_complete_later() {
let mock = mock().bind(&task).release_local();
let task = assert_some!(task.run(From::from(&mock)));
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(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
assert!(handle.is_woken());
let out = assert_ready_ok!(handle.poll());
@@ -288,12 +288,12 @@ fn drop_join_after_poll() {
assert_pending!(handle.poll());
drop(handle);
let task = assert_some!(task.run(From::from(&mock)));
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(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
@@ -317,7 +317,7 @@ fn join_handle_change_task_complete() {
assert_pending!(t1.poll());
drop(t1);
let task = assert_some!(task.run(From::from(&mock)));
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());
@@ -325,7 +325,7 @@ fn join_handle_change_task_complete() {
assert!(!did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
assert!(t2.is_woken());
@@ -347,7 +347,7 @@ fn drop_handle_after_complete() {
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(task.run(&mut || Some(From::from(&mock))).is_none());
assert!(did_drop.did_drop_future());
assert!(!did_drop.did_drop_output());
@@ -370,7 +370,7 @@ fn non_initial_task_state_drop_join_handle_without_polling() {
let mock = mock().bind(&task).schedule().release_local();
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
drop(handle);
@@ -380,7 +380,7 @@ fn non_initial_task_state_drop_join_handle_without_polling() {
tx.send(()).unwrap();
let task = assert_some!(mock.next_pending_run());
assert!(task.run(From::from(&mock)).is_none());
assert!(task.run(&mut || Some(From::from(&mock))).is_none());
assert!(did_drop.did_drop_future());
assert!(did_drop.did_drop_output());
@@ -400,7 +400,7 @@ fn task_panic_background() {
let mock = mock().bind(&task).release_local();
assert!(task.run(From::from(&mock)).is_none());
assert!(task.run(&mut || Some(From::from(&mock))).is_none());
assert!(did_drop.did_drop_future());
}
@@ -422,7 +422,7 @@ fn task_panic_join() {
assert_pending!(handle.poll());
assert!(task.run(From::from(&mock)).is_none());
assert!(task.run(&mut || Some(From::from(&mock))).is_none());
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
@@ -443,12 +443,12 @@ fn complete_second_schedule_obj_before_join() {
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock2)));
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(From::from(&mock1)));
assert_none!(task.run(&mut || Some(From::from(&mock1))));
assert!(did_drop.did_drop_future());
// The join handle was notified
@@ -476,12 +476,12 @@ fn complete_second_schedule_obj_after_join() {
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock2)));
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(From::from(&mock1)));
assert_none!(task.run(&mut || Some(From::from(&mock1))));
assert!(did_drop.did_drop_future());
// The join handle was notified
@@ -512,7 +512,7 @@ fn shutdown_from_list_before_notified() {
let mock = mock().bind(&task).release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
list.shutdown();
assert!(did_drop.did_drop_future());
@@ -542,7 +542,7 @@ fn shutdown_from_list_after_notified() {
let mock = mock().bind(&task).schedule().release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
tx.send(()).unwrap();
@@ -552,7 +552,7 @@ fn shutdown_from_list_after_notified() {
assert_none!(mock.next_pending_drop());
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
assert!(did_drop.did_drop_future());
assert!(handle.is_woken());
@@ -580,8 +580,8 @@ fn shutdown_from_list_after_complete() {
let m2 = mock();
assert_pending!(handle.poll());
let task = assert_some!(task.run(From::from(&m1)));
assert_none!(task.run(From::from(&m2)));
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());
@@ -626,7 +626,7 @@ fn shutdown_from_task_after_notified() {
let mock = mock().bind(&task).schedule().release();
assert_pending!(handle.poll());
assert_none!(task.run(From::from(&mock)));
assert_none!(task.run(&mut || Some(From::from(&mock))));
tx.send(()).unwrap();
tokio/src/executor/thread_pool/builder.rs
+12 -3
View File
@@ -155,7 +155,14 @@ impl Builder {
let (shutdown_tx, shutdown_rx) = shutdown::channel();
let around_worker = self.around_worker.as_ref().map(Arc::clone);
let launch_worker = move |worker: Worker<BoxedPark<P>>| {
let launch_worker = Arc::new(Box::new(move |worker: Worker<BoxedPark<P>>| {
// NOTE: It might seem like the shutdown_tx that's moved into this Arc is never
// dropped, and that shutdown_rx will therefore never see EOF, but that is not actually
// the case. Only `build_with_park` and each worker hold onto a copy of this Arc.
// `build_with_park` drops it immediately, and the workers drop theirs when their `run`
// method returns (and their copy of the Arc are dropped). In fact, we don't actually
// _need_ a copy of `shutdown_tx` for each worker thread; having them all hold onto
// this Arc, which in turn holds the last `shutdown_tx` would have been sufficient.
let shutdown_tx = shutdown_tx.clone();
let around_worker = around_worker.as_ref().map(Arc::clone);
Box::new(move || {
@@ -186,7 +193,8 @@ impl Builder {
// Dropping the handle must happen __after__ the callback
drop(shutdown_tx);
}) as Box<dyn FnOnce() + Send + 'static>
};
})
as Box<dyn Fn(Worker<BoxedPark<P>>) -> Box<dyn FnOnce() + Send> + Send + Sync>);
let mut blocking = crate::executor::blocking::Builder::default();
blocking.name(self.name.clone());
@@ -197,7 +205,8 @@ impl Builder {
let (pool, workers) = worker::create_set::<_, BoxedPark<P>>(
self.pool_size,
|i| BoxedPark::new(build_park(i)),
|i| Box::new(BoxedPark::new(build_park(i))),
Arc::clone(&launch_worker),
blocking.clone(),
);
tokio/src/executor/thread_pool/current.rs
+4
View File
@@ -50,6 +50,10 @@ where
})
}
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,
tokio/src/executor/thread_pool/mod.rs
+3
View File
@@ -37,6 +37,9 @@ mod worker;
#[cfg(test)]
mod tests;
#[cfg(feature = "blocking")]
pub use worker::blocking;
// Re-export `task::Error`
pub use crate::executor::task::Error;
tokio/src/executor/thread_pool/queue/inject.rs
+5
View File
@@ -19,6 +19,11 @@ impl<T: 'static> Inject<T> {
self.cluster.global.push(task, f)
}
/// Check if the queue has been closed
pub(crate) fn is_closed(&self) -> bool {
self.cluster.global.is_closed()
}
/// Close the queue
///
/// Returns `true` if the channel was closed. `false` indicates the pool was
tokio/src/executor/thread_pool/set.rs
+14 -1
View File
@@ -150,6 +150,10 @@ where
}
}
pub(crate) fn is_closed(&self) -> bool {
self.inject.is_closed()
}
pub(crate) fn len(&self) -> usize {
self.shared.len()
}
@@ -175,10 +179,19 @@ where
}
}
impl<P: 'static> Set<P> {
/// Wait 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<P: 'static> Drop for Set<P> {
fn drop(&mut self) {
// Before proceeding, wait for all concurrent wakers to exit
self.inject.wait_for_unlocked();
self.wait_for_unlocked();
}
}
tokio/src/executor/thread_pool/tests/loom_pool.rs
+54 -2
View File
@@ -1,9 +1,9 @@
use crate::spawn;
use crate::executor::loom::sync::atomic::Ordering::{Acquire, Relaxed, Release};
use crate::executor::loom::sync::atomic::{AtomicBool, AtomicUsize};
use crate::executor::loom::sync::{Arc, Mutex};
use crate::executor::tests::loom_oneshot as oneshot;
use crate::executor::thread_pool::ThreadPool;
use crate::executor::thread_pool::{self, Builder, ThreadPool};
use crate::spawn;
use std::future::Future;
@@ -41,6 +41,58 @@ fn pool_multi_spawn() {
});
}
#[test]
fn only_blocking() {
loom::model(|| {
let mut pool = Builder::new().num_threads(1).build();
let (block_tx, block_rx) = oneshot::channel();
pool.spawn(async move {
thread_pool::blocking(move || {
block_tx.send(());
})
});
block_rx.recv();
pool.shutdown_now();
});
}
#[test]
fn blocking_and_regular() {
const NUM: usize = 3;
loom::model(|| {
let mut pool = Builder::new().num_threads(1).build();
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 {
thread_pool::blocking(move || {
block_tx.send(());
})
});
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();
pool.shutdown_now();
});
}
#[test]
fn pool_multi_notify() {
loom::model(|| {
tokio/src/executor/thread_pool/tests/loom_queue.rs
+2 -2
View File
@@ -25,7 +25,7 @@ fn multi_worker() {
// Try to work
while let Some(task) = q.pop_local_first() {
assert!(task.run(From::from(&NOOP_SCHEDULE)).is_none());
assert!(task.run(&mut || Some(From::from(&NOOP_SCHEDULE))).is_none());
let r = rem.get();
assert!(r > 0);
rem.set(r - 1);
@@ -33,7 +33,7 @@ fn multi_worker() {
// Try to steal
if let Some(task) = q.steal(0) {
assert!(task.run(From::from(&NOOP_SCHEDULE)).is_none());
assert!(task.run(&mut || Some(From::from(&NOOP_SCHEDULE))).is_none());
let r = rem.get();
assert!(r > 0);
rem.set(r - 1);
tokio/src/executor/thread_pool/tests/queue.rs
+1 -1
View File
@@ -252,7 +252,7 @@ fn num(task: Task<Noop>) -> u32 {
use std::task::Context;
use std::task::Poll::*;
assert!(task.run(From::from(&NOOP_SCHEDULE)).is_none());
assert!(task.run(&mut || Some(From::from(&NOOP_SCHEDULE))).is_none());
// Find the task that completed
TASKS.with(|c| {
tokio/src/executor/thread_pool/tests/worker.rs
+11 -3
View File
@@ -3,6 +3,8 @@ use crate::executor::thread_pool;
use tokio_test::assert_ok;
use std::sync::Arc;
macro_rules! pool {
(2) => {{
let (pool, mut w, mock_park) = pool!(!2);
@@ -11,8 +13,14 @@ macro_rules! pool {
(! $n:expr) => {{
let mut mock_park = crate::executor::tests::mock_park::MockPark::new();
let blocking = std::sync::Arc::new(crate::executor::blocking::Pool::default());
let (pool, workers) =
thread_pool::create_pool($n, |index| mock_park.mk_park(index), blocking);
let (pool, workers) = thread_pool::create_pool(
$n,
|index| Box::new(mock_park.mk_park(index)),
Arc::new(Box::new(|_| {
unreachable!("attempted to move worker during non-blocking test")
})),
blocking,
);
(pool, workers, mock_park)
}};
}
@@ -39,8 +47,8 @@ fn execute_single_task() {
#[test]
fn task_migrates() {
use std::sync::mpsc;
use crate::sync::oneshot;
use std::sync::mpsc;
let (p, mut w0, mut w1, ..) = pool!(2);
let (tx1, rx1) = oneshot::channel();
tokio/src/executor/thread_pool/worker.rs
+289 -25
View File
@@ -3,8 +3,57 @@ use crate::executor::park::{Park, Unpark};
use crate::executor::task::Task;
use crate::executor::thread_pool::{current, Owned, Shared};
use std::cell::Cell;
use std::ops::{Deref, DerefMut};
use std::time::Duration;
// The Arc<Box<_>> is needed because loom doesn't support Arc<T> where T: !Sized
// loom doesn't support that because it requires CoerceUnsized, which is unstable
type LaunchWorker<P> = Arc<Box<dyn Fn(Worker<P>) -> Box<dyn FnOnce() + Send> + Send + Sync>>;
thread_local! {
/// Thread-local tracking the current executor
static ON_BLOCK: Cell<Option<*mut dyn FnMut()>> = Cell::new(None)
}
/// Run the provided blocking function without blocking the executor.
///
/// In general, issuing a blocking call or performing a lot of compute in a future without
/// yielding is not okay, as it may prevent the executor from driving other futures forward.
/// If you run a closure through this method, the current executor thread will relegate all its
/// executor duties to another (possibly new) thread, and only then poll the task. Note that this
/// requires additional synchronization.
///
/// # Examples
///
/// ```
/// # async fn docs() {
/// tokio::executor::thread_pool::blocking(move || {
/// // do some compute-heavy work or call synchronous code
/// });
/// # }
/// ```
#[cfg(feature = "blocking")]
pub fn blocking<F, R>(f: F) -> R
where
F: FnOnce() -> R,
{
// Make the current worker give away its Worker to another thread so that we can safely block
// this one without preventing progress on other futures the worker owns.
ON_BLOCK.with(|ob| {
let allow_blocking = ob
.get()
.expect("can only call blocking when on Tokio runtime");
// This is safe, because ON_BLOCK was set from an &mut dyn FnMut in the worker that wraps
// the worker's operation, and is unset just prior to when the FnMut is dropped.
let allow_blocking = unsafe { &mut *allow_blocking };
allow_blocking();
f()
})
}
// TODO: remove this re-export
pub(super) use crate::executor::thread_pool::set::Set;
@@ -13,17 +62,24 @@ pub(crate) struct Worker<P: Park + 'static> {
entry: Entry<P::Unpark>,
/// Park the thread
park: P,
park: Box<P>,
/// Fn for launching another Worker should we need it
launch_worker: LaunchWorker<P>,
/// To indicate that the Worker has been given away and should no longer be used
gone: Cell<bool>,
}
pub(crate) fn create_set<F, P>(
pool_size: usize,
mk_park: F,
launch_worker: LaunchWorker<P>,
blocking: Arc<crate::executor::blocking::Pool>,
) -> (Arc<Set<P::Unpark>>, Vec<Worker<P>>)
where
P: Send + Park,
F: FnMut(usize) -> P,
F: FnMut(usize) -> Box<P>,
{
// Create the parks...
let parks: Vec<_> = (0..pool_size).map(mk_park).collect();
@@ -40,7 +96,7 @@ where
.enumerate()
.map(|(index, park)| {
// unsafe is safe because we call Worker::new only once with each index in the pool
unsafe { Worker::new(pool.clone(), index, park) }
unsafe { Worker::new(pool.clone(), index, park, Arc::clone(&launch_worker)) }
})
.collect();
@@ -58,10 +114,17 @@ where
P: Send + Park,
{
// unsafe because new may only be called once for each index in pool's set
pub(super) unsafe fn new(pool: Arc<Set<P::Unpark>>, index: usize, park: P) -> Self {
pub(super) unsafe fn new(
pool: Arc<Set<P::Unpark>>,
index: usize,
park: Box<P>,
launch_worker: LaunchWorker<P>,
) -> Self {
Worker {
entry: Entry::new(pool, index),
park,
launch_worker,
gone: Cell::new(false),
}
}
@@ -72,18 +135,149 @@ where
let mut executor = &**pool;
let entry = &mut self.entry;
let park = &mut self.park;
let launch_worker = &self.launch_worker;
let blocking = &executor.blocking;
let gone = &self.gone;
let mut park = DropNotGone::new(self.park, gone);
// Track the current worker
current::set(&pool, index, || {
let _enter = crate::executor::enter().expect("executor already running on thread");
crate::executor::with_default(&mut executor, || {
crate::executor::blocking::with_pool(blocking, || entry.run(park))
crate::executor::blocking::with_pool(blocking, || {
ON_BLOCK.with(|ob| {
// Ensure that the ON_BLOCK is removed from the thread-local context
// when leaving the scope. This handles cases that involve panicking.
struct Reset<'a>(&'a Cell<Option<*mut dyn FnMut()>>);
impl<'a> Drop for Reset<'a> {
fn drop(&mut self) {
self.0.set(None);
}
}
let _reset = Reset(ob);
let park_ptr = &mut **park as *mut _;
let mut allow_blocking = move || {
// If our Worker has already been given away, then blocking is fine!
if gone.get() {
return;
}
// If this method is called, we need to move the entire worker onto a
// separate (blocking) thread before returning. Once we return, the
// caller is going to execute some blocking code which would otherwise
// block our reactor from making progress. Since we are _in the middle_
// of running a task, this isn't trivial, as the Worker is "active".
// We do have the luxury of knowing that we are on the worker thread,
// so we can assert exclusive access to any Worker-specific state.
//
// More specifically, the caller is _currently_ "stuck" in
// Entry::run_task at:
//
// if let Some(task) = task.run(self.shared().into()) {
//
// And _we_ get to decide when it continues (specifically, by choosing
// when we return from the second callback (i.e., after the FnOnce
// passed to blocking has returned).
//
// Here's what we'll have to do:
//
// - Reconstruct our `Worker` struct
// - Notably, this includes `park`, which we're passing in below.
// - Spawn the reconstructed `Worker` on another blocking thread
// - Clear any state indicating what worker we are on, since at this
// point we are effectively no longer "on" that worker.
// - Allow the caller of `blocking` to continue.
//
// TODO: should we also undo the enter()?
//
// Once the caller completes the blocking operations, we need to ensure
// that async code can continue running in that context. Luckily, since
// `Arc<Set>` has a fallback for when current::get() is None, we can
// just let the task run until it yields, and then put it back into the
// pool.
// We know that the code we're about to execute (inside
// Entry::run_task) has no way to reach the park passed to entry.run.
// therefore, it's fine for us to take ownership of it here _as long as
// we don't drop `park` later_! The DropNotGone wrapper around `park`
// takes care of that.
let park = unsafe { Box::from_raw(park_ptr) };
let worker = Worker {
entry: unsafe {
// The same argument applies here. Since we unset `current`,
// the task's execution won't assume that it owns a worker any
// more. When the task yields, entry will use its `Arc<Set>`
// (which is fine and safe), and then immediately return,
// without calling any code that assumes there is only one
// Entry with the given index (namely it won't call
// Entry::owned).
Entry::new(Arc::clone(&pool), index)
},
park,
launch_worker: Arc::clone(launch_worker),
gone: Cell::new(false),
};
// Give away the worker
//
// TODO: it would be _really_ nice if we had a way to _not_ spawn a
// thread and hand off the worker if the blocking routine ran only for
// a short amount of time. maybe push the Worker onto a "stealing
// queue" somehow? or maybe keep a shared "active" AtomicBool in both
// instances of the Worker, and compare_exchange it to true afterwards
// in an attempt to take it back. if it succeeds, we just resume where
// we were. if it fails, another thread has already stolen the Worker.
crate::executor::blocking::Pool::spawn(
&pool.blocking,
launch_worker(worker),
);
// make sure no subsequent code thinks that it is on a worker
current::clear();
// and make sure that when Entry finishes running the current task,
// it immediately returns all the way up to the worker.
gone.set(true);
};
let allow_blocking: &mut dyn FnMut() = &mut allow_blocking;
ob.set(Some(unsafe {
// NOTE: We cannot use a safe cast to raw pointer here, since we are
// _also_ erasing the lifetime of these pointers. That is safe here,
// because we know that ob will set back to None before allow_blocking
// is dropped.
#[allow(clippy::useless_transmute)]
std::mem::transmute::<_, *mut dyn FnMut()>(allow_blocking)
}));
let _ = entry.run(&mut **park, gone);
// Ensure that we reset ob before allow_blocking is dropped.
drop(_reset);
});
})
})
});
if gone.get() {
// Synchronize with the pool for load(Acquire) in is_closed to get up-to-date value.
pool.wait_for_unlocked();
if pool.is_closed() {
// If the pool is shutting down, some other thread may be waiting to clean up after
// the task that we were holding on to. If we completed that task, we did nothing
// (because task.run() returned None), and so crucially we did not wait up any such
// thread.
//
// So, we have to do that here.
pool.notify_all();
}
}
}
pub(super) fn id(&self) -> usize {
@@ -102,10 +296,12 @@ where
#[cfg(test)]
#[allow(warnings)]
pub(crate) fn tick(&mut self) {
self.entry.tick(&mut self.park);
self.entry.tick(&mut *self.park, &self.gone);
}
}
struct WorkerGone;
struct Entry<P: 'static> {
pool: Arc<Set<P>>,
index: usize,
@@ -120,14 +316,19 @@ where
Entry { pool, index }
}
fn run(&mut self, park: &mut impl Park<Unpark = P>) {
fn run(
&mut self,
park: &mut impl Park<Unpark = P>,
gone: &Cell<bool>,
) -> Result<(), WorkerGone> {
while self.is_running() {
if self.tick(park) {
if self.tick(park, gone)? {
self.park(park);
}
}
self.shutdown(park);
Ok(())
}
fn is_running(&mut self) -> bool {
@@ -135,10 +336,14 @@ where
}
/// Returns `true` if the worker needs to park
fn tick(&mut self, park: &mut impl Park<Unpark = P>) -> bool {
fn tick(
&mut self,
park: &mut impl Park<Unpark = P>,
gone: &Cell<bool>,
) -> Result<bool, WorkerGone> {
// Process all pending tasks in the local queue.
if !self.process_local_queue(park) {
return false;
if !self.process_local_queue(park, gone)? {
return Ok(false);
}
// No more **local** work to process, try transitioning to searching
@@ -147,12 +352,12 @@ where
// On `false`, the worker has entered the parked state
if self.transition_to_searching() {
// If `true` then work was found
if self.search_for_work() {
return false;
if self.search_for_work(gone)? {
return Ok(false);
}
}
true
Ok(true)
}
/// Process all pending tasks in the local queue, occasionally checking the
@@ -160,7 +365,11 @@ where
///
/// Returns `false` if processing was interrupted due to the pool shutting
/// down.
fn process_local_queue(&mut self, park: &mut impl Park<Unpark = P>) -> bool {
fn process_local_queue(
&mut self,
park: &mut impl Park<Unpark = P>,
gone: &Cell<bool>,
) -> Result<bool, WorkerGone> {
debug_assert!(self.is_running());
loop {
@@ -174,7 +383,7 @@ where
self.maintenance();
if !self.is_running() {
return false;
return Ok(false);
}
// Check the global queue
@@ -184,9 +393,9 @@ where
};
if let Some(task) = task {
self.run_task(task);
self.run_task(task, gone)?;
} else {
return true;
return Ok(true);
}
}
}
@@ -214,16 +423,16 @@ where
self.owned().is_running.set(!closed)
}
fn search_for_work(&mut self) -> bool {
fn search_for_work(&mut self, gone: &Cell<bool>) -> Result<bool, WorkerGone> {
debug_assert!(self.is_searching());
if let Some(task) = self.steal_work() {
self.run_task(task);
true
self.run_task(task, gone)?;
Ok(true)
} else {
// Perform some routine work
self.drain_tasks_pending_drop();
false
Ok(false)
}
}
@@ -291,15 +500,34 @@ where
}
}
fn run_task(&mut self, task: Task<Shared<P>>) {
fn run_task(&mut self, task: Task<Shared<P>>, gone: &Cell<bool>) -> Result<(), WorkerGone> {
if self.is_searching() {
self.transition_from_searching();
}
if let Some(task) = task.run(self.shared().into()) {
let executor = self.shared();
let task = task.run(&mut || {
if gone.get() {
None
} else {
Some(executor.into())
}
});
if gone.get() {
// The Worker disappeared from under us.
// We need to return, because we no longer own all of our state!
// Make sure the task gets picked up again eventually.
if let Some(task) = task {
self.pool.schedule(task);
}
return Err(WorkerGone);
}
if let Some(task) = task {
self.owned().submit_local_yield(task);
self.set().notify_work();
}
Ok(())
}
fn final_work_sweep(&mut self) {
@@ -413,3 +641,39 @@ where
unsafe { &*self.set().owned()[self.index].get() }
}
}
struct DropNotGone<'a, T> {
gone: &'a Cell<bool>,
inner: Option<T>,
}
impl<'a, T> DropNotGone<'a, T> {
fn new(inner: T, gone: &'a Cell<bool>) -> Self {
DropNotGone {
gone,
inner: Some(inner),
}
}
}
impl<'a, T> Drop for DropNotGone<'a, T> {
fn drop(&mut self) {
if self.gone.get() {
let inner = self.inner.take().unwrap();
std::mem::forget(inner);
}
}
}
impl<'a, T> Deref for DropNotGone<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
self.inner.as_ref().unwrap()
}
}
impl<'a, T> DerefMut for DropNotGone<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
self.inner.as_mut().unwrap()
}
}
tokio/tests/thread_pool.rs
+47 -1
View File
@@ -1,7 +1,7 @@
#![warn(rust_2018_idioms)]
use tokio::executor::park::{Park, Unpark};
use tokio::executor::thread_pool::*;
use tokio::executor::thread_pool::{self, *};
use futures_util::future::poll_fn;
use std::cell::Cell;
@@ -126,6 +126,52 @@ fn drop_threadpool_drops_futures() {
}
}
#[test]
fn blocking() {
// used for notifying the main thread
const NUM: usize = 10_000;
for _ in 0..50 {
let (tx, rx) = mpsc::channel();
let mut pool = new_pool();
let cnt = Arc::new(AtomicUsize::new(0));
// there are four workers in the pool
// so, if we run 4 blocking tasks, we know that handoff must have happened
let block = Arc::new(std::sync::Barrier::new(5));
for _ in 0..4 {
let block = block.clone();
pool.spawn(async move {
thread_pool::blocking(move || {
block.wait();
block.wait();
})
});
}
block.wait();
for _ in 0..NUM {
let cnt = cnt.clone();
let tx = tx.clone();
pool.spawn(async move {
let num = cnt.fetch_add(1, Relaxed) + 1;
if num == NUM {
tx.send(()).unwrap();
}
});
}
rx.recv().unwrap();
// Wait for the pool to shutdown
block.wait();
pool.shutdown_now();
}
}
#[test]
fn many_oneshot_futures() {
// used for notifying the main thread