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

rt: relax worker notification

Browse Source
This commit is contained in:
Carl Lerche
2023-05-16 17:30:12 -07:00
parent 94767c30d1
commit 2a1e36ac44
2 changed files with 201 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/scheduler/multi_thread/idle.rs
+111 -148
View File
@@ -1,18 +1,25 @@
//! Coordinates idling workers
use crate::loom::sync::atomic::AtomicUsize;
use crate::loom::sync::atomic::{AtomicBool, AtomicUsize};
use crate::loom::sync::Mutex;
use std::fmt;
use std::sync::atomic::Ordering::{self, SeqCst};
use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
pub(super) struct Idle {
/// Tracks both the number of searching workers and the number of unparked
/// workers.
///
/// Used as a fast-path to avoid acquiring the lock when needed.
state: AtomicUsize,
/// Number of searching workers
num_searching: AtomicUsize,
/// Number of sleeping workers
num_sleeping: AtomicUsize,
/// Used to catch false-negatives when notifying workers.
needs_searching: AtomicBool,
// /// Tracks both the number of searching workers and the number of unparked
// /// workers.
// ///
// /// Used as a fast-path to avoid acquiring the lock when needed.
// state: AtomicUsize,
/// Sleeping workers
sleepers: Mutex<Vec<usize>>,
@@ -20,19 +27,15 @@ pub(super) struct Idle {
num_workers: usize,
}
const UNPARK_SHIFT: usize = 16;
const UNPARK_MASK: usize = !SEARCH_MASK;
const SEARCH_MASK: usize = (1 << UNPARK_SHIFT) - 1;
#[derive(Copy, Clone)]
struct State(usize);
impl Idle {
pub(super) fn new(num_workers: usize) -> Idle {
let init = State::new(num_workers);
Idle {
state: AtomicUsize::new(init.into()),
num_searching: AtomicUsize::new(0),
num_sleeping: AtomicUsize::new(0),
needs_searching: AtomicBool::new(false),
sleepers: Mutex::new(Vec::with_capacity(num_workers)),
num_workers,
}
@@ -40,73 +43,122 @@ impl Idle {
/// If there are no workers actively searching, returns the index of a
/// worker currently sleeping.
pub(super) fn worker_to_notify(&self) -> Option<usize> {
// If at least one worker is spinning, work being notified will
// eventually be found. A searching thread will find **some** work and
// notify another worker, eventually leading to our work being found.
pub(super) fn worker_to_notify_local(&self) -> Option<usize> {
// Because this is only called from a worker thread, we can be more
// relaxed here. If we get false-negatives, the caller will eventually
// process the work.
//
// For this to happen, this load must happen before the thread
// transitioning `num_searching` to zero. Acquire / Release does not
// provide sufficient guarantees, so this load is done with `SeqCst` and
// will pair with the `fetch_sub(1)` when transitioning out of
// searching.
if !self.notify_should_wakeup() {
// Before attempting the CAS (which is expensive), do a load which is cheap.
if self.num_searching.load(Acquire) != 0 {
return None;
}
// Acquire the lock
let mut sleepers = self.sleepers.lock();
// Check again, now that the lock is acquired
if !self.notify_should_wakeup() {
if self
.num_searching
.compare_exchange(0, 1, Acquire, Acquire)
.is_err()
{
return None;
}
// A worker should be woken up, atomically increment the number of
// searching workers as well as the number of unparked workers.
State::unpark_one(&self.state, 1);
// We will notify a worker.
let mut sleepers = self.sleepers.lock();
// Get the worker to unpark
let ret = sleepers.pop();
debug_assert!(ret.is_some());
if let Some(ret) = sleepers.pop() {
self.num_sleeping
.store(self.num_sleeping.load(Acquire) - 1, Release);
return Some(ret);
}
ret
self.num_searching.fetch_sub(1, Release);
None
}
/// Returns `true` if the worker needs to do a final check for submitted
/// work.
pub(super) fn transition_worker_to_parked(&self, worker: usize, is_searching: bool) -> bool {
pub(super) fn worker_to_notify_remote(&self) -> Option<usize> {
// Because this function is called from *outside* the runtime, we need
// to be more aggressive with our synchronization. We need to create a
// barrier between pushing a task into the queue (done right before
// calling this method) and ensuring there is at least one spinning
// thread. A load is not sufficient, we must also write to create the
// release relationship.
if self.num_searching.fetch_add(0, AcqRel) != 0 {
return None;
}
// We just created the release ordering and also noticed there are no
// searchers. Try incrementing the number of searches.
if self
.num_searching
.compare_exchange(0, 1, AcqRel, Acquire)
.is_err()
{
// A worker started searching, because of the ordering set by
// `fetch_add` we don't need to do anything else
return None;
}
// We will notify a worker.
let mut sleepers = self.sleepers.lock();
if let Some(ret) = sleepers.pop() {
self.num_sleeping
.store(self.num_sleeping.load(Acquire) - 1, Release);
return Some(ret);
}
self.num_searching.fetch_sub(1, Release);
// We failed to find a worker to wake, we need to make sure we don't lose this wake.
self.needs_searching.store(true, Release);
None
}
/// Returns `false` if the worker must transition back to searching
pub(super) fn transition_worker_to_parked(&self, worker: usize) -> bool {
// Acquire the lock
let mut sleepers = self.sleepers.lock();
// Decrement the number of unparked threads
let ret = State::dec_num_unparked(&self.state, is_searching);
// Track the sleeping worker
sleepers.push(worker);
ret
}
pub(super) fn transition_worker_to_searching(&self) -> bool {
let state = State::load(&self.state, SeqCst);
if 2 * state.num_searching() >= self.num_workers {
if self.needs_searching.load(Acquire) {
return false;
}
// It is possible for this routine to allow more than 50% of the workers
// to search. That is OK. Limiting searchers is only an optimization to
// prevent too much contention.
State::inc_num_searching(&self.state, SeqCst);
// Track the sleeping worker
sleepers.push(worker);
self.num_sleeping
.store(self.num_sleeping.load(Acquire) + 1, Release);
true
}
/// Returns `true` if the worker has become searching
pub(super) fn try_transition_worker_to_searching(&self) -> bool {
let num_searching = self.num_searching.load(Acquire);
let num_sleeping = self.num_sleeping.load(Acquire);
if 2 * num_searching >= self.num_workers - num_sleeping {
return false;
}
self.transition_worker_to_searching();
true
}
pub(super) fn transition_worker_to_searching(&self) {
// Because we are about to become a searching worker, we can
// optimistically clear the need searching flag.
self.needs_searching.store(false, Release);
self.num_searching.fetch_add(1, AcqRel);
}
/// A lightweight transition from searching -> running.
///
/// Returns `true` if this is the final searching worker. The caller
/// **must** notify a new worker.
pub(super) fn transition_worker_from_searching(&self) -> bool {
State::dec_num_searching(&self.state)
let prev = self.num_searching.fetch_sub(1, AcqRel);
debug_assert!(prev > 0);
prev == 1
}
/// Unpark a specific worker. This happens if tasks are submitted from
@@ -118,11 +170,10 @@ impl Idle {
for index in 0..sleepers.len() {
if sleepers[index] == worker_id {
self.num_sleeping
.store(self.num_sleeping.load(Acquire) - 1, Release);
sleepers.swap_remove(index);
// Update the state accordingly while the lock is held.
State::unpark_one(&self.state, 0);
return true;
}
}
@@ -135,92 +186,4 @@ impl Idle {
let sleepers = self.sleepers.lock();
sleepers.contains(&worker_id)
}
fn notify_should_wakeup(&self) -> bool {
let state = State(self.state.fetch_add(0, SeqCst));
state.num_searching() == 0 && state.num_unparked() < self.num_workers
}
}
impl State {
fn new(num_workers: usize) -> State {
// All workers start in the unparked state
let ret = State(num_workers << UNPARK_SHIFT);
debug_assert_eq!(num_workers, ret.num_unparked());
debug_assert_eq!(0, ret.num_searching());
ret
}
fn load(cell: &AtomicUsize, ordering: Ordering) -> State {
State(cell.load(ordering))
}
fn unpark_one(cell: &AtomicUsize, num_searching: usize) {
cell.fetch_add(num_searching | (1 << UNPARK_SHIFT), SeqCst);
}
fn inc_num_searching(cell: &AtomicUsize, ordering: Ordering) {
cell.fetch_add(1, ordering);
}
/// Returns `true` if this is the final searching worker
fn dec_num_searching(cell: &AtomicUsize) -> bool {
let state = State(cell.fetch_sub(1, SeqCst));
state.num_searching() == 1
}
/// Track a sleeping worker
///
/// Returns `true` if this is the final searching worker.
fn dec_num_unparked(cell: &AtomicUsize, is_searching: bool) -> bool {
let mut dec = 1 << UNPARK_SHIFT;
if is_searching {
dec += 1;
}
let prev = State(cell.fetch_sub(dec, SeqCst));
is_searching && prev.num_searching() == 1
}
/// Number of workers currently searching
fn num_searching(self) -> usize {
self.0 & SEARCH_MASK
}
/// Number of workers currently unparked
fn num_unparked(self) -> usize {
(self.0 & UNPARK_MASK) >> UNPARK_SHIFT
}
}
impl From<usize> for State {
fn from(src: usize) -> State {
State(src)
}
}
impl From<State> for usize {
fn from(src: State) -> usize {
src.0
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("worker::State")
.field("num_unparked", &self.num_unparked())
.field("num_searching", &self.num_searching())
.finish()
}
}
#[test]
fn test_state() {
assert_eq!(0, UNPARK_MASK & SEARCH_MASK);
assert_eq!(0, !(UNPARK_MASK | SEARCH_MASK));
let state = State::new(10);
assert_eq!(10, state.num_unparked());
assert_eq!(0, state.num_searching());
}
tokio/src/runtime/scheduler/multi_thread/worker.rs
+90 -66
View File
@@ -437,18 +437,21 @@ impl Context {
continue;
}
// There is no more **local** work to process, try to steal work
// from other workers.
if let Some(task) = core.steal_work(&self.worker) {
core = self.run_task(task, core)?;
} else {
// Wait for work
core = if did_defer_tasks() {
self.park_timeout(core, Some(Duration::from_millis(0)))
} else {
self.park(core)
};
if core.transition_to_searching(&self.worker) {
// There is no more **local** work to process, try to steal work
// from other workers.
if let Some(task) = core.steal_work(&self.worker) {
core = self.run_task(task, core)?;
continue;
}
}
// Wait for work
core = if did_defer_tasks() {
self.park_timeout(core, Some(Duration::from_millis(0)))
} else {
self.park(core)
};
}
core.pre_shutdown(&self.worker);
@@ -544,23 +547,29 @@ impl Context {
/// Also, we rely on the workstealing algorithm to spread the tasks amongst workers
/// after all the IOs get dispatched
fn park(&self, mut core: Box<Core>) -> Box<Core> {
// First, try to transition to the parked state. If this doesn't
// succeed, then we abort the parking process. This transition can fail
// if we detect that there *may* have been a lost wakeup. See the
// transition fn in `idle.rs` for more details.
if !core.transition_to_parked(&self.worker) {
return core;
}
if let Some(f) = &self.worker.handle.shared.config.before_park {
f();
}
if core.transition_to_parked(&self.worker) {
while !core.is_shutdown {
core.metrics.about_to_park();
core = self.park_timeout(core, None);
core.metrics.returned_from_park();
// Run regularly scheduled maintenance
core.maintenance(&self.worker);
if core.transition_from_parked(&self.worker) {
break;
}
while !core.is_shutdown {
if core.transition_from_parked_if_notified(&self.worker) {
break;
}
core.metrics.about_to_park();
core = self.park_timeout(core, None);
core.metrics.returned_from_park();
// Run regularly scheduled maintenance
core.maintenance(&self.worker);
}
if let Some(f) = &self.worker.handle.shared.config.after_unpark {
@@ -594,7 +603,8 @@ impl Context {
// If there are tasks available to steal, but this worker is not
// looking for tasks to steal, notify another worker.
if !core.is_searching && core.run_queue.is_stealable() {
self.worker.handle.notify_parked();
// TODO: is this correct?
self.worker.handle.notify_parked_local();
}
core
@@ -634,10 +644,6 @@ impl Core {
fn steal_work(&mut self, worker: &Worker) -> Option<Notified> {
const ATTEMPTS: usize = 4;
if !self.transition_to_searching(worker) {
return None;
}
// Number of remotes
let num = worker.handle.shared.remotes.len();
@@ -682,55 +688,72 @@ impl Core {
fn transition_to_searching(&mut self, worker: &Worker) -> bool {
if !self.is_searching {
self.is_searching = worker.handle.shared.idle.transition_worker_to_searching();
self.is_searching = worker
.handle
.shared
.idle
.try_transition_worker_to_searching();
}
self.is_searching
}
/// Called right before running a task. In this case, if this is the last
/// searching worker, we need to wake up another worker as there might be
/// other queued tasks.
fn transition_from_searching(&mut self, worker: &Worker) {
if !self.is_searching {
return;
}
self.is_searching = false;
worker.handle.transition_worker_from_searching();
if worker.handle.shared.idle.transition_worker_from_searching() {
// We are the final searching worker, so notify a new one
worker.handle.notify_parked_local();
}
}
/// Prepares the worker state for parking.
///
/// Returns true if the transition happened, false if there is work to do first.
fn transition_to_parked(&mut self, worker: &Worker) -> bool {
// Workers should not park if they have work to do
if worker.lifo_slot().is_some() || self.run_queue.has_tasks() {
return false;
if self.is_searching {
self.is_searching = false;
if worker.handle.shared.idle.transition_worker_from_searching() {
if worker.handle.work_to_steal() {
// Transition *back* to searching
self.is_searching = true;
worker.handle.shared.idle.transition_worker_to_searching();
return false;
}
}
}
// When the final worker transitions **out** of searching to parked, it
// must check all the queues one last time in case work materialized
// between the last work scan and transitioning out of searching.
let is_last_searcher = worker
// When the final worker transitions **out** of searching to parked, the
// transition might fail. In this case, it becomes a searching worker.
let res = worker
.handle
.shared
.idle
.transition_worker_to_parked(worker.index, self.is_searching);
.transition_worker_to_parked(worker.index);
// The worker is no longer searching. Setting this is the local cache
// only.
self.is_searching = false;
if is_last_searcher {
worker.handle.notify_if_work_pending();
if res {
true
} else {
self.is_searching = true;
worker.handle.shared.idle.transition_worker_to_searching();
false
}
true
}
/// Returns `true` if the transition happened.
fn transition_from_parked(&mut self, worker: &Worker) -> bool {
fn transition_from_parked_if_notified(&mut self, worker: &Worker) -> bool {
// If a task is in the lifo slot, then we must unpark regardless of
// being notified
if worker.lifo_slot().is_some() {
if worker.lifo_slot().is_some() || self.run_queue.has_tasks() {
// When a worker wakes, it should only transition to the "searching"
// state when the wake originates from another worker *or* a new task
// is pushed. We do *not* want the worker to transition to "searching"
@@ -829,7 +852,8 @@ impl Handle {
// Otherwise, use the inject queue.
self.shared.inject.push(task);
self.shared.scheduler_metrics.inc_remote_schedule_count();
self.notify_parked();
self.notify_parked_remote();
})
}
@@ -857,7 +881,7 @@ impl Handle {
// scheduling is from a resource driver. As notifications often come in
// batches, the notification is delayed until the park is complete.
if core.park.is_some() {
self.notify_parked();
self.notify_parked_local();
}
}
@@ -867,8 +891,19 @@ impl Handle {
}
}
fn notify_parked(&self) {
if let Some(index) = self.shared.idle.worker_to_notify() {
/// A weaker notification. It is possible for there to be false-negatives (a
/// worker should be woken, but isn't). This is very rare and considered OK
/// as it is called from a worker thread which will, eventually, process the
/// work.
fn notify_parked_local(&self) {
if let Some(index) = self.shared.idle.worker_to_notify_local() {
self.shared.remotes[index].unpark.unpark(&self.driver);
}
}
/// A stronger notify that ensures a worker will always wake up if needed.
fn notify_parked_remote(&self) {
if let Some(index) = self.shared.idle.worker_to_notify_remote() {
self.shared.remotes[index].unpark.unpark(&self.driver);
}
}
@@ -879,25 +914,14 @@ impl Handle {
}
}
fn notify_if_work_pending(&self) {
fn work_to_steal(&self) -> bool {
for remote in &self.shared.remotes[..] {
if !remote.steal.is_empty() {
self.notify_parked();
return;
return true;
}
}
if !self.shared.inject.is_empty() {
self.notify_parked();
}
}
fn transition_worker_from_searching(&self) {
if self.shared.idle.transition_worker_from_searching() {
// We are the final searching worker. Because work was found, we
// need to notify another worker.
self.notify_parked();
}
!self.shared.inject.is_empty()
}
/// Signals that a worker has observed the shutdown signal and has replaced