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

rt(threaded): move inject queue lock to worker (#5754)

Browse Source 
This commit is a step towards the ongoing effort to unify the mutex in
the multi-threaded scheduler. The Inject queue is split into two
structs. `Shared` holds fields that are concurrently accessed, and
`Synced` holds fields that must be locked to access. The multi-threaded
scheduler is responsible for locking `Synced` and passing it in when
needed.

The commit also splits `inject` into multiple files to help reduce the
amount of code defined in macros.
This commit is contained in:
Carl Lerche
2023-06-02 13:36:06 -07:00
committed by GitHub
parent 1e14ef0093
commit fb4d43017d
15 changed files with 539 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/scheduler/inject.rs
+35 -255
View File
@@ -1,292 +1,72 @@
//! Inject queue used to send wakeups to a work-stealing scheduler
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::{Mutex, MutexGuard};
use crate::loom::sync::Mutex;
use crate::runtime::task;
use std::marker::PhantomData;
use std::sync::atomic::Ordering::{Acquire, Release};
mod pop;
pub(crate) use pop::Pop;
mod shared;
pub(crate) use shared::Shared;
mod synced;
pub(crate) use synced::Synced;
cfg_rt_multi_thread! {
mod rt_multi_thread;
}
cfg_metrics! {
mod metrics;
}
/// 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(crate) 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>,
shared: Shared<T>,
synced: Mutex<Synced>,
}
struct Pointers {
/// True if the queue is closed.
is_closed: bool,
/// Linked-list head.
head: Option<task::RawTask>,
/// Linked-list tail.
tail: Option<task::RawTask>,
}
pub(crate) struct Pop<'a, T: 'static> {
len: usize,
pointers: Option<MutexGuard<'a, Pointers>>,
_p: PhantomData<T>,
}
unsafe impl<T> Send for Inject<T> {}
unsafe impl<T> Sync for Inject<T> {}
impl<T: 'static> Inject<T> {
pub(crate) fn new() -> Inject<T> {
let (shared, synced) = Shared::new();
Inject {
pointers: Mutex::new(Pointers {
is_closed: false,
head: None,
tail: None,
}),
len: AtomicUsize::new(0),
_p: PhantomData,
shared,
synced: Mutex::new(synced),
}
}
pub(crate) fn is_empty(&self) -> bool {
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))))]
#[cfg(tokio_taskdump)]
pub(crate) fn is_closed(&self) -> bool {
self.pointers.lock().is_closed
let synced = self.synced.lock();
self.shared.is_closed(&synced)
}
/// Closes the injection queue, returns `true` if the queue is open when the
/// transition is made.
pub(crate) fn close(&self) -> bool {
let mut p = self.pointers.lock();
if p.is_closed {
return false;
}
p.is_closed = true;
true
}
pub(crate) fn len(&self) -> usize {
self.len.load(Acquire)
let mut synced = self.synced.lock();
self.shared.close(&mut synced)
}
/// Pushes a value into the queue.
///
/// This does nothing if the queue is closed.
pub(crate) fn push(&self, task: task::Notified<T>) {
// Acquire queue lock
let mut p = self.pointers.lock();
if p.is_closed {
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 {
// safety: Holding the Notified for a task guarantees exclusive
// access to the `queue_next` field.
set_next(tail, Some(task));
} else {
p.head = Some(task);
}
p.tail = Some(task);
self.len.store(len + 1, Release);
let mut synced = self.synced.lock();
// safety: passing correct `Synced`
unsafe { self.shared.push(&mut synced, task) }
}
pub(crate) fn pop(&self) -> Option<task::Notified<T>> {
self.pop_n(1).next()
}
pub(crate) fn pop_n(&self, n: usize) -> Pop<'_, T> {
use std::cmp;
// Fast path, if len == 0, then there are no values
if self.is_empty() {
return Pop {
len: 0,
pointers: None,
_p: PhantomData,
};
}
// Lock the queue
let p = self.pointers.lock();
// 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() };
let n = cmp::min(n, len);
// Decrement the count.
self.len.store(len - n, Release);
Pop {
len: n,
pointers: Some(p),
_p: PhantomData,
}
}
}
cfg_rt_multi_thread! {
impl<T: 'static> Inject<T> {
/// Pushes several values into the queue.
#[inline]
pub(crate) fn push_batch<I>(&self, mut iter: I)
where
I: Iterator<Item = task::Notified<T>>,
{
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: task::RawTask,
batch_tail: task::RawTask,
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<T: 'static> Drop for Inject<T> {
fn drop(&mut self) {
if !std::thread::panicking() {
assert!(self.pop().is_none(), "queue not empty");
}
}
}
impl<'a, T: 'static> Iterator for Pop<'a, T> {
type Item = task::Notified<T>;
fn next(&mut self) -> Option<Self::Item> {
if self.len == 0 {
if self.shared.is_empty() {
return None;
}
// `pointers` is always `Some` when `len() > 0`
let pointers = self.pointers.as_mut().unwrap();
let ret = pointers.pop();
debug_assert!(ret.is_some());
self.len -= 1;
if self.len == 0 {
self.pointers = None;
}
ret
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
}
impl<'a, T: 'static> ExactSizeIterator for Pop<'a, T> {
fn len(&self) -> usize {
self.len
}
}
impl<'a, T: 'static> Drop for Pop<'a, T> {
fn drop(&mut self) {
for _ in self.by_ref() {}
}
}
impl Pointers {
fn pop<T: 'static>(&mut self) -> Option<task::Notified<T>> {
let task = self.head?;
self.head = get_next(task);
if self.head.is_none() {
self.tail = None;
}
set_next(task, None);
// safety: a `Notified` is pushed into the queue and now it is popped!
Some(unsafe { task::Notified::from_raw(task) })
}
}
fn get_next(task: task::RawTask) -> Option<task::RawTask> {
unsafe { task.get_queue_next() }
}
fn set_next(task: task::RawTask, val: Option<task::RawTask>) {
unsafe {
task.set_queue_next(val);
let mut synced = self.synced.lock();
// safety: passing correct `Synced`
unsafe { self.shared.pop(&mut synced) }
}
}
tokio/src/runtime/scheduler/inject/metrics.rs
+7
View File
@@ -0,0 +1,7 @@
use super::Inject;
impl<T: 'static> Inject<T> {
pub(crate) fn len(&self) -> usize {
self.shared.len()
}
}
tokio/src/runtime/scheduler/inject/pop.rs
+55
View File
@@ -0,0 +1,55 @@
use super::Synced;
use crate::runtime::task;
use std::marker::PhantomData;
pub(crate) struct Pop<'a, T: 'static> {
len: usize,
synced: &'a mut Synced,
_p: PhantomData<T>,
}
impl<'a, T: 'static> Pop<'a, T> {
pub(super) fn new(len: usize, synced: &'a mut Synced) -> Pop<'a, T> {
Pop {
len,
synced,
_p: PhantomData,
}
}
}
impl<'a, T: 'static> Iterator for Pop<'a, T> {
type Item = task::Notified<T>;
fn next(&mut self) -> Option<Self::Item> {
if self.len == 0 {
return None;
}
let ret = self.synced.pop();
// Should be `Some` when `len > 0`
debug_assert!(ret.is_some());
self.len -= 1;
ret
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
}
impl<'a, T: 'static> ExactSizeIterator for Pop<'a, T> {
fn len(&self) -> usize {
self.len
}
}
impl<'a, T: 'static> Drop for Pop<'a, T> {
fn drop(&mut self) {
for _ in self.by_ref() {}
}
}
tokio/src/runtime/scheduler/inject/rt_multi_thread.rs
+98
View File
@@ -0,0 +1,98 @@
use super::{Shared, Synced};
use crate::runtime::scheduler::Lock;
use crate::runtime::task;
use std::sync::atomic::Ordering::Release;
impl<'a> Lock<Synced> for &'a mut Synced {
type Handle = &'a mut Synced;
fn lock(self) -> Self::Handle {
self
}
}
impl AsMut<Synced> for Synced {
fn as_mut(&mut self) -> &mut Synced {
self
}
}
impl<T: 'static> Shared<T> {
/// Pushes several values into the queue.
///
/// # Safety
///
/// Must be called with the same `Synced` instance returned by `Inject::new`
#[inline]
pub(crate) unsafe fn push_batch<L, I>(&self, shared: L, mut iter: I)
where
L: Lock<Synced>,
I: Iterator<Item = task::Notified<T>>,
{
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.
unsafe { prev.set_queue_next(Some(next)) };
prev = next;
counter += 1;
});
// Now that the tasks are linked together, insert them into the
// linked list.
self.push_batch_inner(shared, 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]
unsafe fn push_batch_inner<L>(
&self,
shared: L,
batch_head: task::RawTask,
batch_tail: task::RawTask,
num: usize,
) where
L: Lock<Synced>,
{
debug_assert!(unsafe { batch_tail.get_queue_next().is_none() });
let mut synced = shared.lock();
let synced = synced.as_mut();
if let Some(tail) = synced.tail {
unsafe {
tail.set_queue_next(Some(batch_head));
}
} else {
synced.head = Some(batch_head);
}
synced.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 = self.len.unsync_load();
self.len.store(len + num, Release);
}
}
tokio/src/runtime/scheduler/inject/shared.rs
+119
View File
@@ -0,0 +1,119 @@
use super::{Pop, Synced};
use crate::loom::sync::atomic::AtomicUsize;
use crate::runtime::task;
use std::marker::PhantomData;
use std::sync::atomic::Ordering::{Acquire, Release};
pub(crate) struct Shared<T: 'static> {
/// Number of pending tasks in the queue. This helps prevent unnecessary
/// locking in the hot path.
pub(super) len: AtomicUsize,
_p: PhantomData<T>,
}
unsafe impl<T> Send for Shared<T> {}
unsafe impl<T> Sync for Shared<T> {}
impl<T: 'static> Shared<T> {
pub(crate) fn new() -> (Shared<T>, Synced) {
let inject = Shared {
len: AtomicUsize::new(0),
_p: PhantomData,
};
let synced = Synced {
is_closed: false,
head: None,
tail: None,
};
(inject, synced)
}
pub(crate) fn is_empty(&self) -> bool {
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, synced: &Synced) -> bool {
synced.is_closed
}
/// Closes the injection queue, returns `true` if the queue is open when the
/// transition is made.
pub(crate) fn close(&self, synced: &mut Synced) -> bool {
if synced.is_closed {
return false;
}
synced.is_closed = true;
true
}
pub(crate) fn len(&self) -> usize {
self.len.load(Acquire)
}
/// Pushes a value into the queue.
///
/// This does nothing if the queue is closed.
///
/// # Safety
///
/// Must be called with the same `Synced` instance returned by `Inject::new`
pub(crate) unsafe fn push(&self, synced: &mut Synced, task: task::Notified<T>) {
if synced.is_closed {
return;
}
// safety: only mutated with the lock held
let len = self.len.unsync_load();
let task = task.into_raw();
// The next pointer should already be null
debug_assert!(unsafe { task.get_queue_next().is_none() });
if let Some(tail) = synced.tail {
// safety: Holding the Notified for a task guarantees exclusive
// access to the `queue_next` field.
unsafe { tail.set_queue_next(Some(task)) };
} else {
synced.head = Some(task);
}
synced.tail = Some(task);
self.len.store(len + 1, Release);
}
/// Pop a value from the queue.
///
/// # Safety
///
/// Must be called with the same `Synced` instance returned by `Inject::new`
pub(crate) unsafe fn pop(&self, synced: &mut Synced) -> Option<task::Notified<T>> {
self.pop_n(synced, 1).next()
}
/// Pop `n` values from the queue
///
/// # Safety
///
/// Must be called with the same `Synced` instance returned by `Inject::new`
pub(crate) unsafe fn pop_n<'a>(&'a self, synced: &'a mut Synced, n: usize) -> Pop<'a, T> {
use std::cmp;
// 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 = self.len.unsync_load();
let n = cmp::min(n, len);
// Decrement the count.
self.len.store(len - n, Release);
Pop::new(n, synced)
}
}
tokio/src/runtime/scheduler/inject/synced.rs
+32
View File
@@ -0,0 +1,32 @@
use crate::runtime::task;
pub(crate) struct Synced {
/// True if the queue is closed.
pub(super) is_closed: bool,
/// Linked-list head.
pub(super) head: Option<task::RawTask>,
/// Linked-list tail.
pub(super) tail: Option<task::RawTask>,
}
unsafe impl Send for Synced {}
unsafe impl Sync for Synced {}
impl Synced {
pub(super) fn pop<T: 'static>(&mut self) -> Option<task::Notified<T>> {
let task = self.head?;
self.head = unsafe { task.get_queue_next() };
if self.head.is_none() {
self.tail = None;
}
unsafe { task.set_queue_next(None) };
// safety: a `Notified` is pushed into the queue and now it is popped!
Some(unsafe { task::Notified::from_raw(task) })
}
}
tokio/src/runtime/scheduler/lock.rs
+6
View File
@@ -0,0 +1,6 @@
/// A lock (mutex) yielding generic data.
pub(crate) trait Lock<T> {
type Handle: AsMut<T>;
fn lock(self) -> Self::Handle;
}
tokio/src/runtime/scheduler/mod.rs
+4 -1
View File
@@ -5,11 +5,14 @@ cfg_rt! {
mod defer;
use defer::Defer;
mod inject;
pub(crate) mod inject;
pub(crate) use inject::Inject;
}
cfg_rt_multi_thread! {
mod lock;
use lock::Lock;
pub(crate) mod multi_thread;
pub(crate) use multi_thread::MultiThread;
}
tokio/src/runtime/scheduler/multi_thread/mod.rs
+4 -2
View File
@@ -6,6 +6,9 @@ use counters::Counters;
mod handle;
pub(crate) use handle::Handle;
mod overflow;
pub(crate) use overflow::Overflow;
mod idle;
use self::idle::Idle;
@@ -18,8 +21,7 @@ pub(crate) use park::{Parker, Unparker};
pub(crate) mod queue;
mod worker;
use worker::Shared;
pub(crate) use worker::{Context, Launch};
pub(crate) use worker::{Context, Launch, Shared};
pub(crate) use worker::block_in_place;
tokio/src/runtime/scheduler/multi_thread/overflow.rs
+26
View File
@@ -0,0 +1,26 @@
use crate::runtime::task;
#[cfg(test)]
use std::cell::RefCell;
pub(crate) trait Overflow<T: 'static> {
fn push(&self, task: task::Notified<T>);
fn push_batch<I>(&self, iter: I)
where
I: Iterator<Item = task::Notified<T>>;
}
#[cfg(test)]
impl<T: 'static> Overflow<T> for RefCell<Vec<task::Notified<T>>> {
fn push(&self, task: task::Notified<T>) {
self.borrow_mut().push(task);
}
fn push_batch<I>(&self, iter: I)
where
I: Iterator<Item = task::Notified<T>>,
{
self.borrow_mut().extend(iter);
}
}
tokio/src/runtime/scheduler/multi_thread/queue.rs
+8 -9
View File
@@ -2,8 +2,7 @@
use crate::loom::cell::UnsafeCell;
use crate::loom::sync::Arc;
use crate::runtime::scheduler::multi_thread::Stats;
use crate::runtime::scheduler::Inject;
use crate::runtime::scheduler::multi_thread::{Overflow, Stats};
use crate::runtime::task;
use std::mem::{self, MaybeUninit};
@@ -183,10 +182,10 @@ impl<T> Local<T> {
/// When the queue overflows, half of the curent contents of the queue is
/// moved to the given Injection queue. This frees up capacity for more
/// tasks to be pushed into the local queue.
pub(crate) fn push_back_or_overflow(
pub(crate) fn push_back_or_overflow<O: Overflow<T>>(
&mut self,
mut task: task::Notified<T>,
inject: &Inject<T>,
overflow: &O,
stats: &mut Stats,
) {
let tail = loop {
@@ -202,12 +201,12 @@ impl<T> Local<T> {
} else if steal != real {
// Concurrently stealing, this will free up capacity, so only
// push the task onto the inject queue
inject.push(task);
overflow.push(task);
return;
} else {
// Push the current task and half of the queue into the
// inject queue.
match self.push_overflow(task, real, tail, inject, stats) {
match self.push_overflow(task, real, tail, overflow, stats) {
Ok(_) => return,
// Lost the race, try again
Err(v) => {
@@ -248,12 +247,12 @@ impl<T> Local<T> {
/// workers "missed" some of the tasks during a steal, they will get
/// another opportunity.
#[inline(never)]
fn push_overflow(
fn push_overflow<O: Overflow<T>>(
&mut self,
task: task::Notified<T>,
head: UnsignedShort,
tail: UnsignedShort,
inject: &Inject<T>,
overflow: &O,
stats: &mut Stats,
) -> Result<(), task::Notified<T>> {
/// How many elements are we taking from the local queue.
@@ -336,7 +335,7 @@ impl<T> Local<T> {
head: head as UnsignedLong,
i: 0,
};
inject.push_batch(batch_iter.chain(std::iter::once(task)));
overflow.push_batch(batch_iter.chain(std::iter::once(task)));
// Add 1 to factor in the task currently being scheduled.
stats.incr_overflow_count();
tokio/src/runtime/scheduler/multi_thread/worker.rs
+94 -18
View File
@@ -60,9 +60,9 @@ use crate::loom::sync::{Arc, Mutex};
use crate::runtime;
use crate::runtime::context;
use crate::runtime::scheduler::multi_thread::{
idle, queue, Counters, Handle, Idle, Parker, Stats, Unparker,
idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, Unparker,
};
use crate::runtime::scheduler::{Defer, Inject};
use crate::runtime::scheduler::{inject, Defer, Lock};
use crate::runtime::task::OwnedTasks;
use crate::runtime::{
blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics,
@@ -129,7 +129,7 @@ struct Core {
}
/// State shared across all workers
pub(super) struct Shared {
pub(crate) struct Shared {
/// Per-worker remote state. All other workers have access to this and is
/// how they communicate between each other.
remotes: Box<[Remote]>,
@@ -137,7 +137,7 @@ pub(super) struct Shared {
/// Global task queue used for:
/// 1. Submit work to the scheduler while **not** currently on a worker thread.
/// 2. Submit work to the scheduler when a worker run queue is saturated
inject: Inject<Arc<Handle>>,
inject: inject::Shared<Arc<Handle>>,
/// Coordinates idle workers
idle: Idle,
@@ -172,7 +172,11 @@ pub(super) struct Shared {
/// Data synchronized by the scheduler mutex
pub(super) struct Synced {
/// Synchronized state for `Idle`.
pub(super) idle: idle::Synced,
/// Synchronized state for `Inject`.
inject: inject::Synced,
}
/// Used to communicate with a worker from other threads.
@@ -256,14 +260,18 @@ pub(super) fn create(
}
let (idle, idle_synced) = Idle::new(size);
let (inject, inject_synced) = inject::Shared::new();
let handle = Arc::new(Handle {
shared: Shared {
remotes: remotes.into_boxed_slice(),
inject: Inject::new(),
inject,
idle,
owned: OwnedTasks::new(),
synced: Mutex::new(Synced { idle: idle_synced }),
synced: Mutex::new(Synced {
idle: idle_synced,
inject: inject_synced,
}),
shutdown_cores: Mutex::new(vec![]),
config,
scheduler_metrics: SchedulerMetrics::new(),
@@ -561,7 +569,7 @@ impl Context {
// the back of the queue and return.
core.run_queue.push_back_or_overflow(
task,
self.worker.inject(),
&*self.worker.handle,
&mut core.stats,
);
// If we hit this point, the LIFO slot should be enabled.
@@ -710,7 +718,10 @@ impl Core {
// Update the global queue interval, if needed
self.tune_global_queue_interval(worker);
worker.inject().pop().or_else(|| self.next_local_task())
worker
.handle
.next_remote_task()
.or_else(|| self.next_local_task())
} else {
let maybe_task = self.next_local_task();
@@ -718,6 +729,10 @@ impl Core {
return maybe_task;
}
if worker.inject().is_empty() {
return None;
}
// Other threads can only **remove** tasks from the current worker's
// `run_queue`. So, we can be confident that by the time we call
// `run_queue.push_back` below, there will be *at least* `cap`
@@ -735,7 +750,9 @@ impl Core {
cap,
);
let mut tasks = worker.inject().pop_n(n);
let mut synced = worker.handle.shared.synced.lock();
// safety: passing in the correct `inject::Synced`.
let mut tasks = unsafe { worker.inject().pop_n(&mut synced.inject, n) };
// Pop the first task to return immedietly
let ret = tasks.next();
@@ -783,7 +800,7 @@ impl Core {
}
// Fallback on checking the global queue
worker.handle.shared.inject.pop()
worker.handle.next_remote_task()
}
fn transition_to_searching(&mut self, worker: &Worker) -> bool {
@@ -870,7 +887,8 @@ impl Core {
if !self.is_shutdown {
// Check if the scheduler has been shutdown
self.is_shutdown = worker.inject().is_closed();
let synced = worker.handle.shared.synced.lock();
self.is_shutdown = worker.inject().is_closed(&synced.inject);
}
}
@@ -911,7 +929,7 @@ impl Core {
impl Worker {
/// Returns a reference to the scheduler's injection queue.
fn inject(&self) -> &Inject<Arc<Handle>> {
fn inject(&self) -> &inject::Shared<Arc<Handle>> {
&self.handle.shared.inject
}
}
@@ -946,8 +964,7 @@ impl Handle {
}
// Otherwise, use the inject queue.
self.shared.inject.push(task);
self.shared.scheduler_metrics.inc_remote_schedule_count();
self.push_remote_task(task);
self.notify_parked_remote();
})
}
@@ -961,7 +978,7 @@ impl Handle {
// flexibility and the task may go to the front of the queue.
let should_notify = if is_yield || !core.lifo_enabled {
core.run_queue
.push_back_or_overflow(task, &self.shared.inject, &mut core.stats);
.push_back_or_overflow(task, self, &mut core.stats);
true
} else {
// Push to the LIFO slot
@@ -970,7 +987,7 @@ impl Handle {
if let Some(prev) = prev {
core.run_queue
.push_back_or_overflow(prev, &self.shared.inject, &mut core.stats);
.push_back_or_overflow(prev, self, &mut core.stats);
}
core.lifo_slot = Some(task);
@@ -986,8 +1003,32 @@ impl Handle {
}
}
fn next_remote_task(&self) -> Option<Notified> {
if self.shared.inject.is_empty() {
return None;
}
let mut synced = self.shared.synced.lock();
// safety: passing in correct `idle::Synced`
unsafe { self.shared.inject.pop(&mut synced.inject) }
}
fn push_remote_task(&self, task: Notified) {
self.shared.scheduler_metrics.inc_remote_schedule_count();
let mut synced = self.shared.synced.lock();
// safety: passing in correct `idle::Synced`
unsafe {
self.shared.inject.push(&mut synced.inject, task);
}
}
pub(super) fn close(&self) {
if self.shared.inject.close() {
if self
.shared
.inject
.close(&mut self.shared.synced.lock().inject)
{
self.notify_all();
}
}
@@ -1055,7 +1096,7 @@ impl Handle {
// Drain the injection queue
//
// We already shut down every task, so we can simply drop the tasks.
while let Some(task) = self.shared.inject.pop() {
while let Some(task) = self.next_remote_task() {
drop(task);
}
}
@@ -1065,6 +1106,41 @@ impl Handle {
}
}
impl Overflow<Arc<Handle>> for Handle {
fn push(&self, task: task::Notified<Arc<Handle>>) {
self.push_remote_task(task);
}
fn push_batch<I>(&self, iter: I)
where
I: Iterator<Item = task::Notified<Arc<Handle>>>,
{
unsafe {
self.shared.inject.push_batch(self, iter);
}
}
}
pub(crate) struct InjectGuard<'a> {
lock: crate::loom::sync::MutexGuard<'a, Synced>,
}
impl<'a> AsMut<inject::Synced> for InjectGuard<'a> {
fn as_mut(&mut self) -> &mut inject::Synced {
&mut self.lock.inject
}
}
impl<'a> Lock<inject::Synced> for &'a Handle {
type Handle = InjectGuard<'a>;
fn lock(self) -> Self::Handle {
InjectGuard {
lock: self.shared.synced.lock(),
}
}
}
#[track_caller]
fn with_current<R>(f: impl FnOnce(Option<&Context>) -> R) -> R {
use scheduler::Context::MultiThread;
tokio/src/runtime/tests/inject.rs
+32 -16
View File
@@ -1,38 +1,54 @@
use crate::runtime::scheduler::Inject;
use crate::runtime::scheduler::inject;
#[test]
fn push_and_pop() {
let inject = Inject::new();
const N: usize = 2;
for _ in 0..10 {
let (inject, mut synced) = inject::Shared::new();
for i in 0..N {
assert_eq!(inject.len(), i);
let (task, _) = super::unowned(async {});
inject.push(task);
unsafe { inject.push(&mut synced, task) };
}
for _ in 0..10 {
assert!(inject.pop().is_some());
for i in 0..N {
assert_eq!(inject.len(), N - i);
assert!(unsafe { inject.pop(&mut synced) }.is_some());
}
assert!(inject.pop().is_none());
println!("--------------");
assert!(unsafe { inject.pop(&mut synced) }.is_none());
}
#[test]
fn push_batch_and_pop() {
let inject = Inject::new();
let (inject, mut inject_synced) = inject::Shared::new();
inject.push_batch((0..10).map(|_| super::unowned(async {}).0));
unsafe {
inject.push_batch(
&mut inject_synced,
(0..10).map(|_| super::unowned(async {}).0),
);
assert_eq!(5, inject.pop_n(5).count());
assert_eq!(5, inject.pop_n(5).count());
assert_eq!(0, inject.pop_n(5).count());
assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count());
assert_eq!(5, inject.pop_n(&mut inject_synced, 5).count());
assert_eq!(0, inject.pop_n(&mut inject_synced, 5).count());
}
}
#[test]
fn pop_n_drains_on_drop() {
let inject = Inject::new();
let (inject, mut inject_synced) = inject::Shared::new();
inject.push_batch((0..10).map(|_| super::unowned(async {}).0));
let _ = inject.pop_n(10);
unsafe {
inject.push_batch(
&mut inject_synced,
(0..10).map(|_| super::unowned(async {}).0),
);
let _ = inject.pop_n(&mut inject_synced, 10);
assert_eq!(inject.len(), 0);
assert_eq!(inject.len(), 0);
}
}
tokio/src/runtime/tests/loom_queue.rs
+8 -15
View File
@@ -1,8 +1,8 @@
use crate::runtime::scheduler::multi_thread::{queue, Stats};
use crate::runtime::scheduler::Inject;
use crate::runtime::tests::NoopSchedule;
use loom::thread;
use std::cell::RefCell;
fn new_stats() -> Stats {
Stats::new(&crate::runtime::WorkerMetrics::new())
@@ -12,7 +12,7 @@ fn new_stats() -> Stats {
fn basic() {
loom::model(|| {
let (steal, mut local) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
let mut stats = new_stats();
let th = thread::spawn(move || {
@@ -54,9 +54,7 @@ fn basic() {
}
}
while inject.pop().is_some() {
n += 1;
}
n += inject.borrow_mut().drain(..).count();
n += th.join().unwrap();
@@ -68,7 +66,7 @@ fn basic() {
fn steal_overflow() {
loom::model(|| {
let (steal, mut local) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
let mut stats = new_stats();
let th = thread::spawn(move || {
@@ -108,9 +106,7 @@ fn steal_overflow() {
n += 1;
}
while inject.pop().is_some() {
n += 1;
}
n += inject.borrow_mut().drain(..).count();
assert_eq!(7, n);
});
@@ -139,7 +135,7 @@ fn multi_stealer() {
loom::model(|| {
let (steal, mut local) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
let mut stats = new_stats();
// Push work
@@ -161,9 +157,7 @@ fn multi_stealer() {
n += 1;
}
while inject.pop().is_some() {
n += 1;
}
n += inject.borrow_mut().drain(..).count();
n += th1.join().unwrap();
n += th2.join().unwrap();
@@ -178,7 +172,7 @@ fn chained_steal() {
let mut stats = new_stats();
let (s1, mut l1) = queue::local();
let (s2, mut l2) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
// Load up some tasks
for _ in 0..4 {
@@ -207,6 +201,5 @@ fn chained_steal() {
while l1.pop().is_some() {}
while l2.pop().is_some() {}
while inject.pop().is_some() {}
});
}
tokio/src/runtime/tests/queue.rs
+11 -19
View File
@@ -1,7 +1,7 @@
use crate::runtime::scheduler::multi_thread::{queue, Stats};
use crate::runtime::scheduler::Inject;
use crate::runtime::task::{self, Schedule, Task};
use std::cell::RefCell;
use std::thread;
use std::time::Duration;
@@ -29,7 +29,7 @@ fn new_stats() -> Stats {
#[test]
fn fits_256_one_at_a_time() {
let (_, mut local) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
let mut stats = new_stats();
for _ in 0..256 {
@@ -41,7 +41,7 @@ fn fits_256_one_at_a_time() {
assert_metrics!(stats, overflow_count == 0);
}
assert!(inject.pop().is_none());
assert!(inject.borrow_mut().pop().is_none());
while local.pop().is_some() {}
}
@@ -87,7 +87,7 @@ fn fits_256_all_in_chunks() {
#[test]
fn overflow() {
let (_, mut local) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
let mut stats = new_stats();
for _ in 0..257 {
@@ -101,9 +101,7 @@ fn overflow() {
let mut n = 0;
while inject.pop().is_some() {
n += 1;
}
n += inject.borrow_mut().drain(..).count();
while local.pop().is_some() {
n += 1;
@@ -118,7 +116,7 @@ fn steal_batch() {
let (steal1, mut local1) = queue::local();
let (_, mut local2) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
for _ in 0..4 {
let (task, _) = super::unowned(async {});
@@ -164,7 +162,7 @@ fn stress1() {
for _ in 0..NUM_ITER {
let (steal, mut local) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
let th = thread::spawn(move || {
let mut stats = new_stats();
@@ -207,9 +205,7 @@ fn stress1() {
}
}
while inject.pop().is_some() {
n += 1;
}
n += inject.borrow_mut().drain(..).count();
n += th.join().unwrap();
@@ -227,7 +223,7 @@ fn stress2() {
for _ in 0..NUM_ITER {
let (steal, mut local) = queue::local();
let inject = Inject::new();
let inject = RefCell::new(vec![]);
let th = thread::spawn(move || {
let mut stats = new_stats();
@@ -259,9 +255,7 @@ fn stress2() {
num_pop += 1;
}
while inject.pop().is_some() {
num_pop += 1;
}
num_pop += inject.borrow_mut().drain(..).count();
}
num_pop += th.join().unwrap();
@@ -270,9 +264,7 @@ fn stress2() {
num_pop += 1;
}
while inject.pop().is_some() {
num_pop += 1;
}
num_pop += inject.borrow_mut().drain(..).count();
assert_eq!(num_pop, NUM_TASKS);
}