From 2a1e36ac446e0e2a60198b156b889a071225fe43 Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Tue, 16 May 2023 14:38:34 -0700 Subject: [PATCH] rt: relax worker notification --- .../runtime/scheduler/multi_thread/idle.rs | 259 ++++++++---------- .../runtime/scheduler/multi_thread/worker.rs | 156 ++++++----- 2 files changed, 201 insertions(+), 214 deletions(-) diff --git a/tokio/src/runtime/scheduler/multi_thread/idle.rs b/tokio/src/runtime/scheduler/multi_thread/idle.rs index a57bf6a0b..d0b313cfd 100644 --- a/tokio/src/runtime/scheduler/multi_thread/idle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/idle.rs @@ -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>, @@ -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 { - // 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 { + // 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 { + // 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 for State { - fn from(src: usize) -> State { - State(src) - } -} - -impl From 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()); } diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 93f4dc99e..24ccd284e 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -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) -> Box { + // 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 { 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