diff --git a/benches/Cargo.toml b/benches/Cargo.toml index 25beac754..fdbf50c3c 100644 --- a/benches/Cargo.toml +++ b/benches/Cargo.toml @@ -40,6 +40,11 @@ name = "sync_watch" path = "sync_watch.rs" harness = false +[[bench]] +name = "rt_current_thread" +path = "rt_current_thread.rs" +harness = false + [[bench]] name = "rt_multi_threaded" path = "rt_multi_threaded.rs" diff --git a/benches/rt_current_thread.rs b/benches/rt_current_thread.rs new file mode 100644 index 000000000..dc832193b --- /dev/null +++ b/benches/rt_current_thread.rs @@ -0,0 +1,83 @@ +//! Benchmark implementation details of the threaded 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 bencher::{benchmark_group, benchmark_main, Bencher}; + +const NUM_SPAWN: usize = 1_000; + +fn spawn_many_local(b: &mut Bencher) { + let rt = rt(); + let mut handles = Vec::with_capacity(NUM_SPAWN); + + b.iter(|| { + rt.block_on(async { + for _ in 0..NUM_SPAWN { + handles.push(tokio::spawn(async move {})); + } + + for handle in handles.drain(..) { + handle.await.unwrap(); + } + }); + }); +} + +fn spawn_many_remote_idle(b: &mut Bencher) { + let rt = rt(); + let rt_handle = rt.handle(); + let mut handles = Vec::with_capacity(NUM_SPAWN); + + b.iter(|| { + for _ in 0..NUM_SPAWN { + handles.push(rt_handle.spawn(async {})); + } + + rt.block_on(async { + for handle in handles.drain(..) { + handle.await.unwrap(); + } + }); + }); +} + +fn spawn_many_remote_busy(b: &mut Bencher) { + let rt = rt(); + let rt_handle = rt.handle(); + let mut handles = Vec::with_capacity(NUM_SPAWN); + + rt.spawn(async { + fn iter() { + tokio::spawn(async { iter() }); + } + + iter() + }); + + b.iter(|| { + for _ in 0..NUM_SPAWN { + handles.push(rt_handle.spawn(async {})); + } + + rt.block_on(async { + for handle in handles.drain(..) { + handle.await.unwrap(); + } + }); + }); +} + +fn rt() -> Runtime { + runtime::Builder::new_current_thread().build().unwrap() +} + +benchmark_group!( + scheduler, + spawn_many_local, + spawn_many_remote_idle, + spawn_many_remote_busy +); + +benchmark_main!(scheduler); diff --git a/tokio/src/runtime/scheduler/current_thread.rs b/tokio/src/runtime/scheduler/current_thread.rs index aa9bacf03..8b0dfaaf3 100644 --- a/tokio/src/runtime/scheduler/current_thread.rs +++ b/tokio/src/runtime/scheduler/current_thread.rs @@ -1,8 +1,8 @@ use crate::future::poll_fn; use crate::loom::sync::atomic::AtomicBool; -use crate::loom::sync::{Arc, Mutex}; +use crate::loom::sync::Arc; use crate::runtime::driver::{self, Driver}; -use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task}; +use crate::runtime::task::{self, Inject, JoinHandle, OwnedTasks, Schedule, Task}; use crate::runtime::{blocking, context, scheduler, Config}; use crate::runtime::{MetricsBatch, SchedulerMetrics, WorkerMetrics}; use crate::sync::notify::Notify; @@ -66,8 +66,8 @@ struct Core { /// Scheduler state shared between threads. struct Shared { - /// Remote run queue. None if the `Runtime` has been dropped. - queue: Mutex>>, + /// Remote run queue + inject: Inject>, /// Collection of all active tasks spawned onto this executor. owned: OwnedTasks>, @@ -115,7 +115,7 @@ impl CurrentThread { let handle = Arc::new(Handle { shared: Shared { - queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))), + inject: Inject::new(), owned: OwnedTasks::new(), woken: AtomicBool::new(false), config, @@ -217,15 +217,12 @@ impl CurrentThread { drop(task); } - // Drain remote queue and set it to None - let remote_queue = handle.shared.queue.lock().take(); + // Close the injection queue + handle.shared.inject.close(); - // Using `Option::take` to replace the shared queue with `None`. - // We already shut down every task, so we just need to drop the task. - if let Some(remote_queue) = remote_queue { - for task in remote_queue { - drop(task); - } + // Drain remote queue + while let Some(task) = handle.shared.inject.pop() { + drop(task); } assert!(handle.shared.owned.is_empty()); @@ -259,9 +256,12 @@ impl Core { fn next_task(&mut self, handle: &Handle) -> Option { if self.tick % handle.shared.config.global_queue_interval == 0 { - handle.pop().or_else(|| self.next_local_task(handle)) + handle + .next_remote_task() + .or_else(|| self.next_local_task(handle)) } else { - self.next_local_task(handle).or_else(|| handle.pop()) + self.next_local_task(handle) + .or_else(|| handle.next_remote_task()) } } @@ -440,14 +440,11 @@ impl Handle { }; let local = &mut core.tasks; - let mut injection = self.shared.queue.lock(); - let injection = if let Some(injection) = injection.as_mut() { - injection - } else { + if self.shared.inject.is_closed() { return; - }; + } - traces = trace_current_thread(&self.shared.owned, local, injection) + traces = trace_current_thread(&self.shared.owned, local, &self.shared.inject) .into_iter() .map(dump::Task::new) .collect(); @@ -461,11 +458,8 @@ impl Handle { dump::Dump::new(traces) } - fn pop(&self) -> Option { - match self.shared.queue.lock().as_mut() { - Some(queue) => queue.pop_front(), - None => None, - } + fn next_remote_task(&self) -> Option { + self.shared.inject.pop() } fn waker_ref(me: &Arc) -> WakerRef<'_> { @@ -488,14 +482,7 @@ cfg_metrics! { } pub(crate) fn injection_queue_depth(&self) -> usize { - // TODO: avoid having to lock. The multi-threaded injection queue - // could probably be used here. - self.shared - .queue - .lock() - .as_ref() - .map(|queue| queue.len()) - .unwrap_or(0) + self.shared.inject.len() } pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { @@ -549,14 +536,9 @@ impl Schedule for Arc { // Track that a task was scheduled from **outside** of the runtime. self.shared.scheduler_metrics.inc_remote_schedule_count(); - // If the queue is None, then the runtime has shut down. We - // don't need to do anything with the notification in that case. - let mut guard = self.shared.queue.lock(); - if let Some(queue) = guard.as_mut() { - queue.push_back(task); - drop(guard); - self.driver.unpark(); - } + // Schedule the task + self.shared.inject.push(task); + self.driver.unpark(); } }); } diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index bcccc6988..640ed42bb 100644 --- a/tokio/src/runtime/task/core.rs +++ b/tokio/src/runtime/task/core.rs @@ -278,15 +278,11 @@ impl Core { } } -cfg_rt_multi_thread! { - impl Header { - pub(super) unsafe fn set_next(&self, next: Option>) { - self.queue_next.with_mut(|ptr| *ptr = next); - } - } -} - impl Header { + pub(super) unsafe fn set_next(&self, next: Option>) { + self.queue_next.with_mut(|ptr| *ptr = next); + } + // safety: The caller must guarantee exclusive access to this field, and // must ensure that the id is either 0 or the id of the OwnedTasks // containing this task. diff --git a/tokio/src/runtime/task/inject.rs b/tokio/src/runtime/task/inject.rs index 1585e13a0..eb17ee644 100644 --- a/tokio/src/runtime/task/inject.rs +++ b/tokio/src/runtime/task/inject.rs @@ -52,6 +52,12 @@ impl Inject { self.len() == 0 } + // Kind of annoying to have to include the cfg here + #[cfg(any(tokio_taskdump, all(feature = "rt-multi-thread", not(tokio_wasi))))] + pub(crate) fn is_closed(&self) -> bool { + self.pointers.lock().is_closed + } + /// Closes the injection queue, returns `true` if the queue is open when the /// transition is made. pub(crate) fn close(&self) -> bool { @@ -65,10 +71,6 @@ impl Inject { true } - pub(crate) fn is_closed(&self) -> bool { - self.pointers.lock().is_closed - } - pub(crate) fn len(&self) -> usize { self.len.load(Acquire) } @@ -104,71 +106,6 @@ impl Inject { self.len.store(len + 1, Release); } - /// Pushes several values into the queue. - #[inline] - pub(crate) fn push_batch(&self, mut iter: I) - where - I: Iterator>, - { - let first = match iter.next() { - Some(first) => first.into_raw(), - None => return, - }; - - // Link up all the tasks. - let mut prev = first; - let mut counter = 1; - - // We are going to be called with an `std::iter::Chain`, and that - // iterator overrides `for_each` to something that is easier for the - // compiler to optimize than a loop. - iter.for_each(|next| { - let next = next.into_raw(); - - // safety: Holding the Notified for a task guarantees exclusive - // access to the `queue_next` field. - set_next(prev, Some(next)); - prev = next; - counter += 1; - }); - - // Now that the tasks are linked together, insert them into the - // linked list. - self.push_batch_inner(first, prev, counter); - } - - /// Inserts several tasks that have been linked together into the queue. - /// - /// The provided head and tail may be be the same task. In this case, a - /// single task is inserted. - #[inline] - fn push_batch_inner( - &self, - batch_head: NonNull, - batch_tail: NonNull, - num: usize, - ) { - debug_assert!(get_next(batch_tail).is_none()); - - let mut p = self.pointers.lock(); - - 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(crate) fn pop(&self) -> Option> { // Fast path, if len == 0, then there are no values if self.is_empty() { @@ -201,6 +138,75 @@ impl Inject { } } +cfg_rt_multi_thread! { + impl Inject { + /// Pushes several values into the queue. + #[inline] + pub(crate) fn push_batch(&self, mut iter: I) + where + I: Iterator>, + { + let first = match iter.next() { + Some(first) => first.into_raw(), + None => return, + }; + + // Link up all the tasks. + let mut prev = first; + let mut counter = 1; + + // We are going to be called with an `std::iter::Chain`, and that + // iterator overrides `for_each` to something that is easier for the + // compiler to optimize than a loop. + iter.for_each(|next| { + let next = next.into_raw(); + + // safety: Holding the Notified for a task guarantees exclusive + // access to the `queue_next` field. + set_next(prev, Some(next)); + prev = next; + counter += 1; + }); + + // Now that the tasks are linked together, insert them into the + // linked list. + self.push_batch_inner(first, prev, counter); + } + + /// Inserts several tasks that have been linked together into the queue. + /// + /// The provided head and tail may be be the same task. In this case, a + /// single task is inserted. + #[inline] + fn push_batch_inner( + &self, + batch_head: NonNull, + batch_tail: NonNull, + num: usize, + ) { + debug_assert!(get_next(batch_tail).is_none()); + + let mut p = self.pointers.lock(); + + 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); + } + } +} + impl Drop for Inject { fn drop(&mut self) { if !std::thread::panicking() { diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 35809aca0..8da59cb06 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -182,10 +182,8 @@ mod id; #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] pub use id::{id, try_id, Id}; -cfg_rt_multi_thread! { - mod inject; - pub(super) use self::inject::Inject; -} +mod inject; +pub(super) use self::inject::Inject; #[cfg(feature = "rt")] mod abort; @@ -370,25 +368,23 @@ impl Notified { } } -cfg_rt_multi_thread! { - impl Notified { - unsafe fn from_raw(ptr: NonNull
) -> Notified { - Notified(Task::from_raw(ptr)) - } +impl Notified { + unsafe fn from_raw(ptr: NonNull
) -> Notified { + Notified(Task::from_raw(ptr)) } +} - impl Task { - fn into_raw(self) -> NonNull
{ - let ret = self.raw.header_ptr(); - mem::forget(self); - ret - } +impl Task { + fn into_raw(self) -> NonNull
{ + let ret = self.raw.header_ptr(); + mem::forget(self); + ret } +} - impl Notified { - fn into_raw(self) -> NonNull
{ - self.0.into_raw() - } +impl Notified { + fn into_raw(self) -> NonNull
{ + self.0.into_raw() } } diff --git a/tokio/src/runtime/task/trace/mod.rs b/tokio/src/runtime/task/trace/mod.rs index 5a2152755..fb1909c35 100644 --- a/tokio/src/runtime/task/trace/mod.rs +++ b/tokio/src/runtime/task/trace/mod.rs @@ -1,5 +1,6 @@ use crate::loom::sync::Arc; use crate::runtime::scheduler::current_thread; +use crate::runtime::task::Inject; use backtrace::BacktraceFrame; use std::cell::Cell; use std::collections::VecDeque; @@ -236,11 +237,14 @@ impl Future for Root { pub(in crate::runtime) fn trace_current_thread( owned: &OwnedTasks>, local: &mut VecDeque>>, - injection: &mut VecDeque>>, + injection: &Inject>, ) -> Vec { // clear the local and injection queues local.clear(); - injection.clear(); + + while let Some(task) = injection.pop() { + drop(task); + } // notify each task let mut tasks = vec![]; diff --git a/tokio/tests/rt_common.rs b/tokio/tests/rt_common.rs index 039e27e9b..9c6add047 100644 --- a/tokio/tests/rt_common.rs +++ b/tokio/tests/rt_common.rs @@ -1317,4 +1317,39 @@ rt_test! { } }); } + + #[test] + #[cfg(not(target_os="wasi"))] + fn shutdown_concurrent_spawn() { + const NUM_TASKS: usize = 10_000; + for _ in 0..5 { + let (tx, rx) = std::sync::mpsc::channel(); + let rt = rt(); + + let mut txs = vec![]; + + for _ in 0..NUM_TASKS { + let (tx, rx) = tokio::sync::oneshot::channel(); + txs.push(tx); + rt.spawn(async move { + rx.await.unwrap(); + }); + } + + // Prime the tasks + rt.block_on(async { tokio::task::yield_now().await }); + + let th = std::thread::spawn(move || { + tx.send(()).unwrap(); + for tx in txs.drain(..) { + let _ = tx.send(()); + } + }); + + rx.recv().unwrap(); + drop(rt); + + th.join().unwrap(); + } + } }