From 109fd3086b97896c422d9565e7d1ddbe4b6f300b Mon Sep 17 00:00:00 2001 From: Jon Gjengset Date: Wed, 30 Oct 2019 11:58:49 -0400 Subject: [PATCH] thread-pool: in-place blocking with new scheduler (#1681) The initial new scheduler PR omitted in-place blocking support. This patch brings it back. --- ci/azure-loom.yml | 2 +- tokio/src/executor/task/harness.rs | 10 +- tokio/src/executor/task/mod.rs | 2 +- tokio/src/executor/task/raw.rs | 9 +- tokio/src/executor/task/tests/loom.rs | 16 +- tokio/src/executor/task/tests/task.rs | 64 ++-- tokio/src/executor/thread_pool/builder.rs | 15 +- tokio/src/executor/thread_pool/current.rs | 4 + tokio/src/executor/thread_pool/mod.rs | 3 + .../src/executor/thread_pool/queue/inject.rs | 5 + tokio/src/executor/thread_pool/set.rs | 15 +- .../executor/thread_pool/tests/loom_pool.rs | 56 +++- .../executor/thread_pool/tests/loom_queue.rs | 4 +- tokio/src/executor/thread_pool/tests/queue.rs | 2 +- .../src/executor/thread_pool/tests/worker.rs | 14 +- tokio/src/executor/thread_pool/worker.rs | 314 ++++++++++++++++-- tokio/tests/thread_pool.rs | 48 ++- 17 files changed, 496 insertions(+), 87 deletions(-) diff --git a/ci/azure-loom.yml b/ci/azure-loom.yml index c08b68ca1..99e03da27 100644 --- a/ci/azure-loom.yml +++ b/ci/azure-loom.yml @@ -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' diff --git a/tokio/src/executor/task/harness.rs b/tokio/src/executor/task/harness.rs index ef2978dc4..e5355e4fe 100644 --- a/tokio/src/executor/task/harness.rs +++ b/tokio/src/executor/task/harness.rs @@ -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) -> bool { + pub(super) fn poll(mut self, executor: &mut dyn FnMut() -> Option>) -> 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, + executor: &mut dyn FnMut() -> Option>, join_interest: bool, output: super::Result, ) { @@ -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(); diff --git a/tokio/src/executor/task/mod.rs b/tokio/src/executor/task/mod.rs index dfc6628e7..3e4e500ba 100644 --- a/tokio/src/executor/task/mod.rs +++ b/tokio/src/executor/task/mod.rs @@ -99,7 +99,7 @@ impl Task { impl Task { /// Returns `self` when the task needs to be immediately re-scheduled - pub(crate) fn run(self, executor: NonNull) -> Option { + pub(crate) fn run(self, executor: &mut dyn FnMut() -> Option>) -> Option { if unsafe { self.raw.poll(executor) } { Some(self) } else { diff --git a/tokio/src/executor/task/raw.rs b/tokio/src/executor/task/raw.rs index a9b048fce..d0ffac3ab 100644 --- a/tokio/src/executor/task/raw.rs +++ b/tokio/src/executor/task/raw.rs @@ -15,7 +15,7 @@ pub(super) struct RawTask { pub(super) struct Vtable { /// Poll the future - pub(super) poll: unsafe fn(*mut (), NonNull) -> bool, + pub(super) poll: unsafe fn(*mut (), &mut dyn FnMut() -> Option>) -> bool, /// The task handle has been dropped and the join waker needs to be dropped /// or the task struct needs to be deallocated @@ -101,7 +101,7 @@ impl RawTask { /// 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) -> bool { + pub(super) unsafe fn poll(self, executor: &mut dyn FnMut() -> Option>) -> bool { // Get the vtable without holding a ref to the meta struct. This is done // because a mutable reference to the task is passed into the poll fn. let vtable = self.header().vtable; @@ -150,7 +150,10 @@ impl Clone for RawTask { impl Copy for RawTask {} -unsafe fn poll(ptr: *mut (), executor: NonNull) -> bool { +unsafe fn poll( + ptr: *mut (), + executor: &mut dyn FnMut() -> Option>, +) -> bool { let harness = Harness::::from_raw(ptr); harness.poll(executor) } diff --git a/tokio/src/executor/task/tests/loom.rs b/tokio/src/executor/task/tests/loom.rs index 53539987f..b377aeb99 100644 --- a/tokio/src/executor/task/tests/loom.rs +++ b/tokio/src/executor/task/tests/loom.rs @@ -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))); } } } diff --git a/tokio/src/executor/task/tests/task.rs b/tokio/src/executor/task/tests/task.rs index 76319b8df..66fbae59a 100644 --- a/tokio/src/executor/task/tests/task.rs +++ b/tokio/src/executor/task/tests/task.rs @@ -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(); diff --git a/tokio/src/executor/thread_pool/builder.rs b/tokio/src/executor/thread_pool/builder.rs index 7955e2864..c1298a5ac 100644 --- a/tokio/src/executor/thread_pool/builder.rs +++ b/tokio/src/executor/thread_pool/builder.rs @@ -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>| { + let launch_worker = Arc::new(Box::new(move |worker: Worker>| { + // 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 - }; + }) + as Box>) -> Box + 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

>( self.pool_size, - |i| BoxedPark::new(build_park(i)), + |i| Box::new(BoxedPark::new(build_park(i))), + Arc::clone(&launch_worker), blocking.clone(), ); diff --git a/tokio/src/executor/thread_pool/current.rs b/tokio/src/executor/thread_pool/current.rs index 6910dca1e..f02be1018 100644 --- a/tokio/src/executor/thread_pool/current.rs +++ b/tokio/src/executor/thread_pool/current.rs @@ -50,6 +50,10 @@ where }) } +pub(super) fn clear() { + CURRENT_WORKER.with(|cell| cell.set(Inner::new())) +} + pub(super) fn get(f: F) -> R where F: FnOnce(&Current) -> R, diff --git a/tokio/src/executor/thread_pool/mod.rs b/tokio/src/executor/thread_pool/mod.rs index bab594ef7..091792452 100644 --- a/tokio/src/executor/thread_pool/mod.rs +++ b/tokio/src/executor/thread_pool/mod.rs @@ -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; diff --git a/tokio/src/executor/thread_pool/queue/inject.rs b/tokio/src/executor/thread_pool/queue/inject.rs index dbad65fd3..5f8caf1ac 100644 --- a/tokio/src/executor/thread_pool/queue/inject.rs +++ b/tokio/src/executor/thread_pool/queue/inject.rs @@ -19,6 +19,11 @@ impl Inject { 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 diff --git a/tokio/src/executor/thread_pool/set.rs b/tokio/src/executor/thread_pool/set.rs index dddd50e88..5ee2f544a 100644 --- a/tokio/src/executor/thread_pool/set.rs +++ b/tokio/src/executor/thread_pool/set.rs @@ -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 Set

{ + /// 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 Drop for Set

{ fn drop(&mut self) { // Before proceeding, wait for all concurrent wakers to exit - self.inject.wait_for_unlocked(); + self.wait_for_unlocked(); } } diff --git a/tokio/src/executor/thread_pool/tests/loom_pool.rs b/tokio/src/executor/thread_pool/tests/loom_pool.rs index 55da54d09..df7540701 100644 --- a/tokio/src/executor/thread_pool/tests/loom_pool.rs +++ b/tokio/src/executor/thread_pool/tests/loom_pool.rs @@ -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(|| { diff --git a/tokio/src/executor/thread_pool/tests/loom_queue.rs b/tokio/src/executor/thread_pool/tests/loom_queue.rs index 8b0214a3d..cc9ae4494 100644 --- a/tokio/src/executor/thread_pool/tests/loom_queue.rs +++ b/tokio/src/executor/thread_pool/tests/loom_queue.rs @@ -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); diff --git a/tokio/src/executor/thread_pool/tests/queue.rs b/tokio/src/executor/thread_pool/tests/queue.rs index be89d94d2..94a3c7a18 100644 --- a/tokio/src/executor/thread_pool/tests/queue.rs +++ b/tokio/src/executor/thread_pool/tests/queue.rs @@ -252,7 +252,7 @@ fn num(task: Task) -> 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| { diff --git a/tokio/src/executor/thread_pool/tests/worker.rs b/tokio/src/executor/thread_pool/tests/worker.rs index dc132ab76..f5f9bace0 100644 --- a/tokio/src/executor/thread_pool/tests/worker.rs +++ b/tokio/src/executor/thread_pool/tests/worker.rs @@ -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(); diff --git a/tokio/src/executor/thread_pool/worker.rs b/tokio/src/executor/thread_pool/worker.rs index 7fbe6589b..b78c214ab 100644 --- a/tokio/src/executor/thread_pool/worker.rs +++ b/tokio/src/executor/thread_pool/worker.rs @@ -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> is needed because loom doesn't support Arc where T: !Sized +// loom doesn't support that because it requires CoerceUnsized, which is unstable +type LaunchWorker

= Arc) -> Box + Send + Sync>>; + +thread_local! { + /// Thread-local tracking the current executor + static ON_BLOCK: Cell> = 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: 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 { entry: Entry, /// Park the thread - park: P, + park: Box

, + + /// Fn for launching another Worker should we need it + launch_worker: LaunchWorker

, + + /// To indicate that the Worker has been given away and should no longer be used + gone: Cell, } pub(crate) fn create_set( pool_size: usize, mk_park: F, + launch_worker: LaunchWorker

, blocking: Arc, ) -> (Arc>, Vec>) where P: Send + Park, - F: FnMut(usize) -> P, + F: FnMut(usize) -> Box

, { // 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>, index: usize, park: P) -> Self { + pub(super) unsafe fn new( + pool: Arc>, + index: usize, + park: Box

, + launch_worker: LaunchWorker

, + ) -> 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>); + + 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` 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` + // (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 { pool: Arc>, index: usize, @@ -120,14 +316,19 @@ where Entry { pool, index } } - fn run(&mut self, park: &mut impl Park) { + fn run( + &mut self, + park: &mut impl Park, + gone: &Cell, + ) -> 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) -> bool { + fn tick( + &mut self, + park: &mut impl Park, + gone: &Cell, + ) -> Result { // 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) -> bool { + fn process_local_queue( + &mut self, + park: &mut impl Park, + gone: &Cell, + ) -> Result { 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) -> Result { 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>) { + fn run_task(&mut self, task: Task>, gone: &Cell) -> 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, + inner: Option, +} + +impl<'a, T> DropNotGone<'a, T> { + fn new(inner: T, gone: &'a Cell) -> 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() + } +} diff --git a/tokio/tests/thread_pool.rs b/tokio/tests/thread_pool.rs index ef8fcf068..e65f873c3 100644 --- a/tokio/tests/thread_pool.rs +++ b/tokio/tests/thread_pool.rs @@ -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