From 8a2279f2abad1f82fe7c3f95b67144da322a4a4a Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Mon, 26 Jun 2023 20:18:40 +0000 Subject: [PATCH] wip --- tokio/src/runtime/builder.rs | 10 +- tokio/src/runtime/scheduler/mod.rs | 7 - .../scheduler/multi_thread/counters.rs | 114 +- .../runtime/scheduler/multi_thread/handle.rs | 5 +- .../runtime/scheduler/multi_thread/idle.rs | 537 ++---- .../src/runtime/scheduler/multi_thread/mod.rs | 18 +- .../runtime/scheduler/multi_thread/queue.rs | 23 +- .../runtime/scheduler/multi_thread/stats.rs | 65 +- .../runtime/scheduler/multi_thread/worker.rs | 1591 +++++++---------- tokio/src/util/mod.rs | 5 + 10 files changed, 877 insertions(+), 1498 deletions(-) diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index ef636e7d6..d2e10b004 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -1198,7 +1198,7 @@ cfg_rt_multi_thread! { fn build_threaded_runtime(&mut self) -> io::Result { use crate::loom::sys::num_cpus; use crate::runtime::{Config, runtime::Scheduler}; - use crate::runtime::scheduler::MultiThread; + use crate::runtime::scheduler::{self, MultiThread}; let core_threads = self.worker_threads.unwrap_or_else(num_cpus); @@ -1213,7 +1213,7 @@ cfg_rt_multi_thread! { let seed_generator_1 = self.seed_generator.next_generator(); let seed_generator_2 = self.seed_generator.next_generator(); - let (scheduler, handle) = MultiThread::new( + let (scheduler, handle, launch) = MultiThread::new( core_threads, driver, driver_handle, @@ -1232,6 +1232,12 @@ cfg_rt_multi_thread! { }, ); + let handle = Handle { inner: scheduler::Handle::MultiThread(handle) }; + + // Spawn the thread pool workers + let _enter = handle.enter(); + launch.launch(); + Ok(Runtime::from_parts(Scheduler::MultiThread(scheduler), handle, blocking_pool)) } diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs index e0e1ede65..af0e56ad7 100644 --- a/tokio/src/runtime/scheduler/mod.rs +++ b/tokio/src/runtime/scheduler/mod.rs @@ -160,13 +160,6 @@ cfg_rt! { } cfg_rt_multi_thread! { - pub(crate) fn expect_multi_thread(&self) -> &Arc { - match self { - Handle::MultiThread(handle) => handle, - _ => panic!("not a `MultiThread` handle"), - } - } - cfg_unstable! { pub(crate) fn expect_multi_thread_alt(&self) -> &Arc { match self { diff --git a/tokio/src/runtime/scheduler/multi_thread/counters.rs b/tokio/src/runtime/scheduler/multi_thread/counters.rs index edda0d46d..50bcc1198 100644 --- a/tokio/src/runtime/scheduler/multi_thread/counters.rs +++ b/tokio/src/runtime/scheduler/multi_thread/counters.rs @@ -5,63 +5,24 @@ mod imp { static NUM_MAINTENANCE: AtomicUsize = AtomicUsize::new(0); static NUM_NOTIFY_LOCAL: AtomicUsize = AtomicUsize::new(0); - static NUM_NOTIFY_REMOTE: AtomicUsize = AtomicUsize::new(0); static NUM_UNPARKS_LOCAL: AtomicUsize = AtomicUsize::new(0); - static NUM_UNPARKS_REMOTE: AtomicUsize = AtomicUsize::new(0); static NUM_LIFO_SCHEDULES: AtomicUsize = AtomicUsize::new(0); static NUM_LIFO_CAPPED: AtomicUsize = AtomicUsize::new(0); - static NUM_STEALS: AtomicUsize = AtomicUsize::new(0); - static NUM_OVERFLOW: AtomicUsize = AtomicUsize::new(0); - static NUM_PARK: AtomicUsize = AtomicUsize::new(0); - static NUM_POLLS: AtomicUsize = AtomicUsize::new(0); - static NUM_LIFO_POLLS: AtomicUsize = AtomicUsize::new(0); - static NUM_REMOTE_BATCH: AtomicUsize = AtomicUsize::new(0); - static NUM_GLOBAL_QUEUE_INTERVAL: AtomicUsize = AtomicUsize::new(0); - static NUM_NO_AVAIL_CORE: AtomicUsize = AtomicUsize::new(0); - static NUM_RELAY_SEARCH: AtomicUsize = AtomicUsize::new(0); - static NUM_SPIN_STALL: AtomicUsize = AtomicUsize::new(0); - static NUM_NO_LOCAL_WORK: AtomicUsize = AtomicUsize::new(0); impl Drop for super::Counters { fn drop(&mut self) { let notifies_local = NUM_NOTIFY_LOCAL.load(Relaxed); - let notifies_remote = NUM_NOTIFY_REMOTE.load(Relaxed); let unparks_local = NUM_UNPARKS_LOCAL.load(Relaxed); - let unparks_remote = NUM_UNPARKS_REMOTE.load(Relaxed); let maintenance = NUM_MAINTENANCE.load(Relaxed); let lifo_scheds = NUM_LIFO_SCHEDULES.load(Relaxed); let lifo_capped = NUM_LIFO_CAPPED.load(Relaxed); - let num_steals = NUM_STEALS.load(Relaxed); - let num_overflow = NUM_OVERFLOW.load(Relaxed); - let num_park = NUM_PARK.load(Relaxed); - let num_polls = NUM_POLLS.load(Relaxed); - let num_lifo_polls = NUM_LIFO_POLLS.load(Relaxed); - let num_remote_batch = NUM_REMOTE_BATCH.load(Relaxed); - let num_global_queue_interval = NUM_GLOBAL_QUEUE_INTERVAL.load(Relaxed); - let num_no_avail_core = NUM_NO_AVAIL_CORE.load(Relaxed); - let num_relay_search = NUM_RELAY_SEARCH.load(Relaxed); - let num_spin_stall = NUM_SPIN_STALL.load(Relaxed); - let num_no_local_work = NUM_NO_LOCAL_WORK.load(Relaxed); println!("---"); - println!("notifies (remote): {}", notifies_remote); - println!(" notifies (local): {}", notifies_local); - println!(" unparks (local): {}", unparks_local); - println!(" unparks (remote): {}", unparks_remote); - println!(" notify, no core: {}", num_no_avail_core); - println!(" maintenance: {}", maintenance); - println!(" LIFO schedules: {}", lifo_scheds); - println!(" LIFO capped: {}", lifo_capped); - println!(" steals: {}", num_steals); - println!(" queue overflows: {}", num_overflow); - println!(" parks: {}", num_park); - println!(" polls: {}", num_polls); - println!(" polls (LIFO): {}", num_lifo_polls); - println!("remote task batch: {}", num_remote_batch); - println!("global Q interval: {}", num_global_queue_interval); - println!(" relay search: {}", num_relay_search); - println!(" spin stall: {}", num_spin_stall); - println!(" no local work: {}", num_no_local_work); + println!("notifies (local): {}", notifies_local); + println!(" unparks (local): {}", unparks_local); + println!(" maintenance: {}", maintenance); + println!(" LIFO schedules: {}", lifo_scheds); + println!(" LIFO capped: {}", lifo_capped); } } @@ -69,18 +30,10 @@ mod imp { NUM_NOTIFY_LOCAL.fetch_add(1, Relaxed); } - pub(crate) fn inc_num_notify_remote() { - NUM_NOTIFY_REMOTE.fetch_add(1, Relaxed); - } - pub(crate) fn inc_num_unparks_local() { NUM_UNPARKS_LOCAL.fetch_add(1, Relaxed); } - pub(crate) fn inc_num_unparks_remote() { - NUM_UNPARKS_REMOTE.fetch_add(1, Relaxed); - } - pub(crate) fn inc_num_maintenance() { NUM_MAINTENANCE.fetch_add(1, Relaxed); } @@ -92,72 +45,15 @@ mod imp { pub(crate) fn inc_lifo_capped() { NUM_LIFO_CAPPED.fetch_add(1, Relaxed); } - - pub(crate) fn inc_num_steals() { - NUM_STEALS.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_overflows() { - NUM_OVERFLOW.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_parks() { - NUM_PARK.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_polls() { - NUM_POLLS.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_lifo_polls() { - NUM_LIFO_POLLS.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_remote_batch() { - NUM_REMOTE_BATCH.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_global_queue_interval() { - NUM_GLOBAL_QUEUE_INTERVAL.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_notify_no_core() { - NUM_NO_AVAIL_CORE.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_relay_search() { - NUM_RELAY_SEARCH.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_spin_stall() { - NUM_SPIN_STALL.fetch_add(1, Relaxed); - } - - pub(crate) fn inc_num_no_local_work() { - NUM_NO_LOCAL_WORK.fetch_add(1, Relaxed); - } } #[cfg(not(tokio_internal_mt_counters))] mod imp { pub(crate) fn inc_num_inc_notify_local() {} - pub(crate) fn inc_num_notify_remote() {} pub(crate) fn inc_num_unparks_local() {} - pub(crate) fn inc_num_unparks_remote() {} pub(crate) fn inc_num_maintenance() {} pub(crate) fn inc_lifo_schedules() {} pub(crate) fn inc_lifo_capped() {} - pub(crate) fn inc_num_steals() {} - pub(crate) fn inc_num_overflows() {} - pub(crate) fn inc_num_parks() {} - pub(crate) fn inc_num_polls() {} - pub(crate) fn inc_num_lifo_polls() {} - pub(crate) fn inc_num_remote_batch() {} - pub(crate) fn inc_global_queue_interval() {} - pub(crate) fn inc_notify_no_core() {} - pub(crate) fn inc_num_relay_search() {} - pub(crate) fn inc_num_spin_stall() {} - pub(crate) fn inc_num_no_local_work() {} } #[derive(Debug)] diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 781433d95..98e476585 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -43,8 +43,7 @@ impl Handle { } pub(crate) fn shutdown(&self) { - self.shared.close(); - self.driver.unpark(); + self.close(); } pub(super) fn bind_new_task(me: &Arc, future: T, id: task::Id) -> JoinHandle @@ -55,7 +54,7 @@ impl Handle { let (handle, notified) = me.shared.owned.bind(future, me.clone(), id); if let Some(notified) = notified { - me.shared.schedule_task(notified, false); + me.schedule_task(notified, false); } handle diff --git a/tokio/src/runtime/scheduler/multi_thread/idle.rs b/tokio/src/runtime/scheduler/multi_thread/idle.rs index bd8ffa399..834bc2b66 100644 --- a/tokio/src/runtime/scheduler/multi_thread/idle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/idle.rs @@ -1,425 +1,240 @@ //! Coordinates idling workers -use crate::loom::sync::atomic::{AtomicBool, AtomicUsize}; -use crate::loom::sync::MutexGuard; -use crate::runtime::scheduler::multi_thread::{worker, Core, Shared}; +use crate::loom::sync::atomic::AtomicUsize; +use crate::runtime::scheduler::multi_thread::Shared; -use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; +use std::fmt; +use std::sync::atomic::Ordering::{self, SeqCst}; pub(super) struct Idle { - /// Number of searching cores - num_searching: AtomicUsize, + /// 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 idle cores - num_idle: AtomicUsize, - - /// Map of idle cores - idle_map: IdleMap, - - /// Used to catch false-negatives when waking workers - needs_searching: AtomicBool, - - /// Total number of cores - num_cores: usize, -} - -pub(super) struct IdleMap { - chunks: Vec, -} - -pub(super) struct Snapshot { - chunks: Vec, + /// Total number of workers. + num_workers: usize, } /// Data synchronized by the scheduler mutex pub(super) struct Synced { - /// Worker IDs that are currently sleeping + /// Sleeping workers sleepers: Vec, - - /// Cores available for workers - available_cores: Vec>, } +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(cores: Vec>, num_workers: usize) -> (Idle, Synced) { + pub(super) fn new(num_workers: usize) -> (Idle, Synced) { + let init = State::new(num_workers); + let idle = Idle { - num_searching: AtomicUsize::new(0), - num_idle: AtomicUsize::new(cores.len()), - idle_map: IdleMap::new(&cores), - needs_searching: AtomicBool::new(false), - num_cores: cores.len(), + state: AtomicUsize::new(init.into()), + num_workers, }; let synced = Synced { sleepers: Vec::with_capacity(num_workers), - available_cores: cores, }; (idle, synced) } - pub(super) fn num_idle(&self, synced: &Synced) -> usize { - debug_assert_eq!(synced.available_cores.len(), self.num_idle.load(Acquire)); - synced.available_cores.len() - } - - pub(super) fn num_searching(&self) -> usize { - self.num_searching.load(Acquire) - } - - pub(super) fn snapshot(&self, snapshot: &mut Snapshot) { - snapshot.update(&self.idle_map) - } - - /// Try to acquire an available core - pub(super) fn try_acquire_available_core(&self, synced: &mut Synced) -> Option> { - let ret = synced.available_cores.pop(); - - if let Some(core) = &ret { - // Decrement the number of idle cores - let num_idle = self.num_idle.load(Acquire) - 1; - debug_assert_eq!(num_idle, synced.available_cores.len()); - self.num_idle.store(num_idle, Release); - - self.idle_map.unset(core.index); - debug_assert!(self.idle_map.matches(&synced.available_cores)); + /// If there are no workers actively searching, returns the index of a + /// worker currently sleeping. + pub(super) fn worker_to_notify(&self, shared: &Shared) -> 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. + // + // 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() { + return None; } + // Acquire the lock + let mut lock = shared.synced.lock(); + + // Check again, now that the lock is acquired + if !self.notify_should_wakeup() { + 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); + + // Get the worker to unpark + let ret = lock.idle.sleepers.pop(); + debug_assert!(ret.is_some()); + ret } - /// We need at least one searching worker - pub(super) fn notify_local(&self, shared: &Shared) { - if self.num_searching.load(Acquire) != 0 { - // There already is a searching worker. Note, that this could be a - // false positive. However, because this method is called **from** a - // worker, we know that there is at least one worker currently - // awake, so the scheduler won't deadlock. - return; - } - - if self.num_idle.load(Acquire) == 0 { - self.needs_searching.store(true, Release); - return; - } - - // There aren't any searching workers. Try to initialize one - if self - .num_searching - .compare_exchange(0, 1, AcqRel, Acquire) - .is_err() - { - // Failing the compare_exchange means another thread concurrently - // launched a searching worker. - return; - } - - super::counters::inc_num_unparks_local(); - - // Acquire the lock - let synced = shared.synced.lock(); - self.notify_synced(synced, shared); - } - - /// Notifies a single worker - pub(super) fn notify_remote(&self, synced: MutexGuard<'_, worker::Synced>, shared: &Shared) { - if synced.idle.sleepers.is_empty() { - self.needs_searching.store(true, Release); - return; - } - - // We need to establish a stronger barrier than with `notify_local` - if self - .num_searching - .compare_exchange(0, 1, AcqRel, Acquire) - .is_err() - { - return; - } - - self.notify_synced(synced, shared); - } - - /// Notify a worker while synced - fn notify_synced(&self, mut synced: MutexGuard<'_, worker::Synced>, shared: &Shared) { - // Find a sleeping worker - if let Some(worker) = synced.idle.sleepers.pop() { - // Find an available core - if let Some(mut core) = synced.idle.available_cores.pop() { - debug_assert!(!core.is_searching); - core.is_searching = true; - - self.idle_map.unset(core.index); - debug_assert!(self.idle_map.matches(&synced.idle.available_cores)); - - // Assign the core to the worker - synced.assigned_cores[worker] = Some(core); - - let num_idle = synced.idle.available_cores.len(); - debug_assert_eq!(num_idle, self.num_idle.load(Acquire) - 1); - - // Update the number of sleeping workers - self.num_idle.store(num_idle, Release); - - // Drop the lock before notifying the condvar. - drop(synced); - - super::counters::inc_num_unparks_remote(); - - // Notify the worker - shared.condvars[worker].notify_one(); - return; - } else { - synced.idle.sleepers.push(worker); - } - } - - super::counters::inc_notify_no_core(); - - // Set the `needs_searching` flag, this happens *while* the lock is held. - self.needs_searching.store(true, Release); - self.num_searching.fetch_sub(1, Release); - - // Explicit mutex guard drop to show that holding the guard to this - // point is significant. `needs_searching` and `num_searching` must be - // updated in the critical section. - drop(synced); - } - - pub(super) fn notify_mult( + /// Returns `true` if the worker needs to do a final check for submitted + /// work. + pub(super) fn transition_worker_to_parked( &self, - synced: &mut worker::Synced, - workers: &mut Vec, - num: usize, - ) { - debug_assert!(workers.is_empty()); - - for _ in 0..num { - if let Some(worker) = synced.idle.sleepers.pop() { - if let Some(core) = synced.idle.available_cores.pop() { - debug_assert!(!core.is_searching); - - self.idle_map.unset(core.index); - - synced.assigned_cores[worker] = Some(core); - - workers.push(worker); - - continue; - } else { - synced.idle.sleepers.push(worker); - } - } - - break; - } - - if !workers.is_empty() { - debug_assert!(self.idle_map.matches(&synced.idle.available_cores)); - let num_idle = synced.idle.available_cores.len(); - self.num_idle.store(num_idle, Release); - } else { - debug_assert_eq!( - synced.idle.available_cores.len(), - self.num_idle.load(Acquire) - ); - self.needs_searching.store(true, Release); - } - } - - pub(super) fn shutdown(&self, synced: &mut worker::Synced, shared: &Shared) { - // Wake every sleeping worker and assign a core to it. There may not be - // enough sleeping workers for all cores, but other workers will - // eventually find the cores and shut them down. - while !synced.idle.sleepers.is_empty() && !synced.idle.available_cores.is_empty() { - let worker = synced.idle.sleepers.pop().unwrap(); - let core = synced.idle.available_cores.pop().unwrap(); - - self.idle_map.unset(core.index); - - synced.assigned_cores[worker] = Some(core); - shared.condvars[worker].notify_one(); - - self.num_idle - .store(synced.idle.available_cores.len(), Release); - } - - debug_assert!(self.idle_map.matches(&synced.idle.available_cores)); - - // Wake up any other workers - while let Some(index) = synced.idle.sleepers.pop() { - shared.condvars[index].notify_one(); - } - } - - /// The worker releases the given core, making it available to other workers - /// that are waiting. - pub(super) fn release_core(&self, synced: &mut worker::Synced, core: Box) { - // The core should not be searching at this point - debug_assert!(!core.is_searching); - - // Check that this isn't the final worker to go idle *and* - // `needs_searching` is set. - debug_assert!(!self.needs_searching.load(Acquire) || num_active_workers(&synced.idle) > 1); - - let num_idle = synced.idle.available_cores.len(); - debug_assert_eq!(num_idle, self.num_idle.load(Acquire)); - - self.idle_map.set(core.index); - - // Store the core in the list of available cores - synced.idle.available_cores.push(core); - - debug_assert!(self.idle_map.matches(&synced.idle.available_cores)); - - // Update `num_idle` - self.num_idle.store(num_idle + 1, Release); - } - - pub(super) fn transition_worker_to_parked(&self, synced: &mut worker::Synced, index: usize) { - // Store the worker index in the list of sleepers - synced.idle.sleepers.push(index); - - // The worker's assigned core slot should be empty - debug_assert!(synced.assigned_cores[index].is_none()); - } - - pub(super) fn try_transition_worker_to_searching(&self, core: &mut Core) { - debug_assert!(!core.is_searching); - - let num_searching = self.num_searching.load(Acquire); - let num_idle = self.num_idle.load(Acquire); - - if 2 * num_searching >= self.num_cores - num_idle { - return; - } - - self.transition_worker_to_searching(core); - } - - /// Needs to happen while synchronized in order to avoid races - pub(super) fn transition_worker_to_searching_if_needed( - &self, - _synced: &mut Synced, - core: &mut Core, + shared: &Shared, + worker: usize, + is_searching: bool, ) -> bool { - if self.needs_searching.load(Acquire) { - // Needs to be called while holding the lock - self.transition_worker_to_searching(core); - true - } else { - false - } + // Acquire the lock + let mut lock = shared.synced.lock(); + + // Decrement the number of unparked threads + let ret = State::dec_num_unparked(&self.state, is_searching); + + // Track the sleeping worker + lock.idle.sleepers.push(worker); + + ret } - fn transition_worker_to_searching(&self, core: &mut Core) { - core.is_searching = true; - self.num_searching.fetch_add(1, AcqRel); - self.needs_searching.store(false, Release); + pub(super) fn transition_worker_to_searching(&self) -> bool { + let state = State::load(&self.state, SeqCst); + if 2 * state.num_searching() >= self.num_workers { + 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); + true } /// 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, core: &mut Core) -> bool { - debug_assert!(core.is_searching); - core.is_searching = false; - - let prev = self.num_searching.fetch_sub(1, AcqRel); - debug_assert!(prev > 0); - - prev == 1 - } -} - -const BITS: usize = usize::BITS as usize; -const BIT_MASK: usize = (usize::BITS - 1) as usize; - -impl IdleMap { - fn new(cores: &[Box]) -> IdleMap { - let ret = IdleMap::new_n(num_chunks(cores.len())); - ret.set_all(cores); - - ret + pub(super) fn transition_worker_from_searching(&self) -> bool { + State::dec_num_searching(&self.state) } - fn new_n(n: usize) -> IdleMap { - let chunks = (0..n).map(|_| AtomicUsize::new(0)).collect(); - IdleMap { chunks } - } + /// Unpark a specific worker. This happens if tasks are submitted from + /// within the worker's park routine. + /// + /// Returns `true` if the worker was parked before calling the method. + pub(super) fn unpark_worker_by_id(&self, shared: &Shared, worker_id: usize) -> bool { + let mut lock = shared.synced.lock(); + let sleepers = &mut lock.idle.sleepers; - fn set(&self, index: usize) { - let (chunk, mask) = index_to_mask(index); - let prev = self.chunks[chunk].load(Acquire); - let next = prev | mask; - self.chunks[chunk].store(next, Release); - } + for index in 0..sleepers.len() { + if sleepers[index] == worker_id { + sleepers.swap_remove(index); - fn set_all(&self, cores: &[Box]) { - for core in cores { - self.set(core.index); - } - } + // Update the state accordingly while the lock is held. + State::unpark_one(&self.state, 0); - fn unset(&self, index: usize) { - let (chunk, mask) = index_to_mask(index); - let prev = self.chunks[chunk].load(Acquire); - let next = prev & !mask; - self.chunks[chunk].store(next, Release); - } - - fn matches(&self, idle_cores: &[Box]) -> bool { - let expect = IdleMap::new_n(self.chunks.len()); - expect.set_all(idle_cores); - - for (i, chunk) in expect.chunks.iter().enumerate() { - if chunk.load(Acquire) != self.chunks[i].load(Acquire) { - return false; + return true; } } - true + false + } + + /// Returns `true` if `worker_id` is contained in the sleep set. + pub(super) fn is_parked(&self, shared: &Shared, worker_id: usize) -> bool { + let lock = shared.synced.lock(); + lock.idle.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 Snapshot { - pub(crate) fn new(idle: &Idle) -> Snapshot { - let chunks = vec![0; idle.idle_map.chunks.len()]; - let mut ret = Snapshot { chunks }; - ret.update(&idle.idle_map); +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 update(&mut self, idle_map: &IdleMap) { - for i in 0..self.chunks.len() { - self.chunks[i] = idle_map.chunks[i].load(Acquire); + 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 } - pub(super) fn is_idle(&self, index: usize) -> bool { - let (chunk, mask) = index_to_mask(index); - debug_assert!( - chunk < self.chunks.len(), - "index={}; chunks={}", - index, - self.chunks.len() - ); - self.chunks[chunk] & mask == mask + /// 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 } } -fn num_chunks(max_cores: usize) -> usize { - (max_cores / BITS) + 1 +impl From for State { + fn from(src: usize) -> State { + State(src) + } } -fn index_to_mask(index: usize) -> (usize, usize) { - let mask = 1 << (index & BIT_MASK); - let chunk = index / BITS; - - (chunk, mask) +impl From for usize { + fn from(src: State) -> usize { + src.0 + } } -fn num_active_workers(synced: &Synced) -> usize { - synced.available_cores.capacity() - synced.available_cores.len() +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/mod.rs b/tokio/src/runtime/scheduler/multi_thread/mod.rs index 11be279ca..d85a0ae0a 100644 --- a/tokio/src/runtime/scheduler/multi_thread/mod.rs +++ b/tokio/src/runtime/scheduler/multi_thread/mod.rs @@ -15,11 +15,13 @@ use self::idle::Idle; mod stats; pub(crate) use stats::Stats; +mod park; +pub(crate) use park::{Parker, Unparker}; + pub(crate) mod queue; mod worker; -use worker::Core; -pub(crate) use worker::{Context, Shared}; +pub(crate) use worker::{Context, Launch, Shared}; cfg_taskdump! { mod trace; @@ -35,8 +37,9 @@ cfg_not_taskdump! { pub(crate) use worker::block_in_place; +use crate::loom::sync::Arc; use crate::runtime::{ - self, blocking, + blocking, driver::{self, Driver}, scheduler, Config, }; @@ -58,17 +61,18 @@ impl MultiThread { blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, - ) -> (MultiThread, runtime::Handle) { - let handle = worker::create( + ) -> (MultiThread, Arc, Launch) { + let parker = Parker::new(driver); + let (handle, launch) = worker::create( size, - driver, + parker, driver_handle, blocking_spawner, seed_generator, config, ); - (MultiThread, handle) + (MultiThread, handle, launch) } /// Blocks the current thread waiting for the future to complete. diff --git a/tokio/src/runtime/scheduler/multi_thread/queue.rs b/tokio/src/runtime/scheduler/multi_thread/queue.rs index a07d76c96..dd66fa2dd 100644 --- a/tokio/src/runtime/scheduler/multi_thread/queue.rs +++ b/tokio/src/runtime/scheduler/multi_thread/queue.rs @@ -33,7 +33,6 @@ pub(crate) struct Local { /// Consumer handle. May be used from many threads. pub(crate) struct Steal(Arc>); -#[repr(align(128))] pub(crate) struct Inner { /// Concurrently updated by many threads. /// @@ -106,6 +105,11 @@ pub(crate) fn local() -> (Steal, Local) { } impl Local { + /// Returns the number of entries in the queue + pub(crate) fn len(&self) -> usize { + self.inner.len() as usize + } + /// How many tasks can be pushed into the queue pub(crate) fn remaining_slots(&self) -> usize { self.inner.remaining_slots() @@ -115,9 +119,12 @@ impl Local { LOCAL_QUEUE_CAPACITY } - /// Returns `true` if there are no entries in the queue - pub(crate) fn is_empty(&self) -> bool { - self.inner.is_empty() + /// Returns false if there are any entries in the queue + /// + /// Separate to is_stealable so that refactors of is_stealable to "protect" + /// some tasks from stealing won't affect this + pub(crate) fn has_tasks(&self) -> bool { + !self.inner.is_empty() } /// Pushes a batch of tasks to the back of the queue. All tasks must fit in @@ -192,13 +199,11 @@ impl Local { // There is capacity for the task break tail; } else if steal != real { - super::counters::inc_num_overflows(); // Concurrently stealing, this will free up capacity, so only // push the task onto the inject queue overflow.push(task); return; } else { - super::counters::inc_num_overflows(); // Push the current task and half of the queue into the // inject queue. match self.push_overflow(task, real, tail, overflow, stats) { @@ -381,6 +386,10 @@ impl Local { } impl Steal { + pub(crate) fn is_empty(&self) -> bool { + self.0.is_empty() + } + /// Steals half the tasks from self and place them into `dst`. pub(crate) fn steal_into( &self, @@ -411,8 +420,6 @@ impl Steal { return None; } - super::counters::inc_num_steals(); - dst_stats.incr_steal_count(n as u16); dst_stats.incr_steal_operations(); diff --git a/tokio/src/runtime/scheduler/multi_thread/stats.rs b/tokio/src/runtime/scheduler/multi_thread/stats.rs index 57657bb03..f01daaa1b 100644 --- a/tokio/src/runtime/scheduler/multi_thread/stats.rs +++ b/tokio/src/runtime/scheduler/multi_thread/stats.rs @@ -10,16 +10,6 @@ pub(crate) struct Stats { /// user. batch: MetricsBatch, - /// Exponentially-weighted moving average of time spent polling scheduled a - /// task. - /// - /// Tracked in nanoseconds, stored as a f64 since that is what we use with - /// the EWMA calculations - task_poll_time_ewma: f64, -} - -/// Transient state -pub(crate) struct Ephemeral { /// Instant at which work last resumed (continued after park). /// /// This duplicates the value stored in `MetricsBatch`. We will unify @@ -29,20 +19,12 @@ pub(crate) struct Ephemeral { /// Number of tasks polled in the batch of scheduled tasks tasks_polled_in_batch: usize, - /// Used to ensure calls to start / stop batch are paired - #[cfg(debug_assertions)] - batch_started: bool, -} - -impl Ephemeral { - pub(crate) fn new() -> Ephemeral { - Ephemeral { - processing_scheduled_tasks_started_at: Instant::now(), - tasks_polled_in_batch: 0, - #[cfg(debug_assertions)] - batch_started: false, - } - } + /// Exponentially-weighted moving average of time spent polling scheduled a + /// task. + /// + /// Tracked in nanoseconds, stored as a f64 since that is what we use with + /// the EWMA calculations + task_poll_time_ewma: f64, } /// How to weigh each individual poll time, value is plucked from thin air. @@ -58,9 +40,6 @@ const MAX_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 127; const TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL: u32 = 61; impl Stats { - pub(crate) const DEFAULT_GLOBAL_QUEUE_INTERVAL: u32 = - TARGET_TASKS_POLLED_PER_GLOBAL_QUEUE_INTERVAL; - pub(crate) fn new(worker_metrics: &WorkerMetrics) -> Stats { // Seed the value with what we hope to see. let task_poll_time_ewma = @@ -68,6 +47,8 @@ impl Stats { Stats { batch: MetricsBatch::new(worker_metrics), + processing_scheduled_tasks_started_at: Instant::now(), + tasks_polled_in_batch: 0, task_poll_time_ewma, } } @@ -104,36 +85,24 @@ impl Stats { self.batch.inc_local_schedule_count(); } - pub(crate) fn start_processing_scheduled_tasks(&mut self, ephemeral: &mut Ephemeral) { + pub(crate) fn start_processing_scheduled_tasks(&mut self) { self.batch.start_processing_scheduled_tasks(); - #[cfg(debug_assertions)] - { - debug_assert!(!ephemeral.batch_started); - ephemeral.batch_started = true; - } - - ephemeral.processing_scheduled_tasks_started_at = Instant::now(); - ephemeral.tasks_polled_in_batch = 0; + self.processing_scheduled_tasks_started_at = Instant::now(); + self.tasks_polled_in_batch = 0; } - pub(crate) fn end_processing_scheduled_tasks(&mut self, ephemeral: &mut Ephemeral) { + pub(crate) fn end_processing_scheduled_tasks(&mut self) { self.batch.end_processing_scheduled_tasks(); - #[cfg(debug_assertions)] - { - debug_assert!(ephemeral.batch_started); - ephemeral.batch_started = false; - } - // Update the EWMA task poll time - if ephemeral.tasks_polled_in_batch > 0 { + if self.tasks_polled_in_batch > 0 { let now = Instant::now(); // If we "overflow" this conversion, we have bigger problems than // slightly off stats. - let elapsed = (now - ephemeral.processing_scheduled_tasks_started_at).as_nanos() as f64; - let num_polls = ephemeral.tasks_polled_in_batch as f64; + let elapsed = (now - self.processing_scheduled_tasks_started_at).as_nanos() as f64; + let num_polls = self.tasks_polled_in_batch as f64; // Calculate the mean poll duration for a single task in the batch let mean_poll_duration = elapsed / num_polls; @@ -147,10 +116,10 @@ impl Stats { } } - pub(crate) fn start_poll(&mut self, ephemeral: &mut Ephemeral) { + pub(crate) fn start_poll(&mut self) { self.batch.start_poll(); - ephemeral.tasks_polled_in_batch += 1; + self.tasks_polled_in_batch += 1; } pub(crate) fn end_poll(&mut self) { diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 4209a9f36..7fc335f51 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -56,22 +56,21 @@ //! the inject queue indefinitely. This would be a ref-count cycle and a memory //! leak. -use crate::loom::sync::{Arc, Condvar, Mutex, MutexGuard}; +use crate::loom::sync::{Arc, Mutex}; use crate::runtime; use crate::runtime::context; use crate::runtime::scheduler::multi_thread::{ - idle, queue, stats, Counters, Handle, Idle, Overflow, Stats, TraceStatus, + idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker, }; -use crate::runtime::scheduler::{self, inject, Lock}; +use crate::runtime::scheduler::{inject, Defer, Lock}; use crate::runtime::task::OwnedTasks; use crate::runtime::{ - blocking, coop, driver, task, Config, Driver, SchedulerMetrics, WorkerMetrics, + blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics, }; use crate::util::atomic_cell::AtomicCell; use crate::util::rand::{FastRand, RngSeedGenerator}; -use std::cell::{Cell, RefCell}; -use std::cmp; +use std::cell::RefCell; use std::task::Waker; use std::time::Duration; @@ -88,61 +87,66 @@ cfg_not_taskdump! { } /// A scheduler worker -/// -/// Data is stack-allocated and never migrates threads pub(super) struct Worker { - /// Used to schedule bookkeeping tasks every so often. - tick: u32, + /// Reference to scheduler's handle + handle: Arc, - /// True if the scheduler is being shutdown - pub(super) is_shutdown: bool, + /// Index holding this worker's remote state + index: usize, - /// True if the scheduler is being traced - is_traced: bool, - - /// Counter used to track when to poll from the local queue vs. the - /// injection queue - num_seq_local_queue_polls: u32, - - /// How often to check the global queue - global_queue_interval: u32, - - /// Used to collect a list of workers to notify - workers_to_notify: Vec, - - /// Snapshot of idle core list. This helps speedup stealing - idle_snapshot: idle::Snapshot, - - stats: stats::Ephemeral, + /// Used to hand-off a worker's core to another thread. + core: AtomicCell, } /// Core data -/// -/// Data is heap-allocated and migrates threads. -#[repr(align(128))] -pub(super) struct Core { - /// Index holding this core's remote/shared state. - pub(super) index: usize, +struct Core { + /// Used to schedule bookkeeping tasks every so often. + tick: u32, + /// When a task is scheduled from a worker, it is stored in this slot. The + /// worker will check this slot for a task **before** checking the run + /// queue. This effectively results in the **last** scheduled task to be run + /// next (LIFO). This is an optimization for improving locality which + /// benefits message passing patterns and helps to reduce latency. lifo_slot: Option, + /// When `true`, locally scheduled tasks go to the LIFO slot. When `false`, + /// they go to the back of the `run_queue`. + lifo_enabled: bool, + /// The worker-local run queue. run_queue: queue::Local>, /// True if the worker is currently searching for more work. Searching /// involves attempting to steal from other workers. - pub(super) is_searching: bool, + is_searching: bool, + + /// True if the scheduler is being shutdown + is_shutdown: bool, + + /// True if the scheduler is being traced + is_traced: bool, + + /// Parker + /// + /// Stored in an `Option` as the parker is added / removed to make the + /// borrow checker happy. + park: Option, /// Per-worker runtime stats stats: Stats, + /// How often to check the global queue + global_queue_interval: u32, + /// Fast random number generator. rand: FastRand, } /// State shared across all workers pub(crate) struct Shared { - /// Per-core remote state. + /// Per-worker remote state. All other workers have access to this and is + /// how they communicate between each other. remotes: Box<[Remote]>, /// Global task queue used for: @@ -159,13 +163,12 @@ pub(crate) struct Shared { /// Data synchronized by the scheduler mutex pub(super) synced: Mutex, - /// Power's Tokio's I/O, timers, etc... the responsibility of polling the - /// driver is shared across workers. - driver: AtomicCell, - - /// Condition variables used to unblock worker threads. Each worker thread - /// has its own condvar it waits on. - pub(super) condvars: Vec, + /// Cores that have observed the shutdown signal + /// + /// The core is **not** placed back in the worker to avoid it from being + /// stolen by a thread that was spawned as part of `block_in_place`. + #[allow(clippy::vec_box)] // we're moving an already-boxed value + shutdown_cores: Mutex>>, /// The number of cores that have observed the trace signal. pub(super) trace_status: TraceStatus, @@ -187,16 +190,6 @@ pub(crate) struct Shared { /// Data synchronized by the scheduler mutex pub(crate) struct Synced { - /// When worker is notified, it is assigned a core. The core is placed here - /// until the worker wakes up to take it. - pub(super) assigned_cores: Vec>>, - - /// Cores that have observed the shutdown signal - /// - /// The core is **not** placed back in the worker to avoid it from being - /// stolen by a thread that was spawned as part of `block_in_place`. - shutdown_cores: Vec>, - /// Synchronized state for `Idle`. pub(super) idle: idle::Synced, @@ -206,44 +199,33 @@ pub(crate) struct Synced { /// Used to communicate with a worker from other threads. struct Remote { - /// When a task is scheduled from a worker, it is stored in this slot. The - /// worker will check this slot for a task **before** checking the run - /// queue. This effectively results in the **last** scheduled task to be run - /// next (LIFO). This is an optimization for improving locality which - /// benefits message passing patterns and helps to reduce latency. - // lifo_slot: Lifo, - /// Steals tasks from this worker. pub(super) steal: queue::Steal>, + + /// Unparks the associated worker thread + unpark: Unparker, } /// Thread-local context pub(crate) struct Context { - // Current scheduler's handle - handle: Arc, - - /// Worker index - index: usize, - - /// True when the LIFO slot is enabled - lifo_enabled: Cell, + /// Worker + worker: Arc, /// Core data core: RefCell>>, - /// Used to pass cores to other threads when `block_in_place` is called - handoff_core: Arc>, - /// Tasks to wake after resource drivers are polled. This is mostly to /// handle yielded tasks. - pub(crate) defer: RefCell>, + pub(crate) defer: Defer, } +/// Starts the workers +pub(crate) struct Launch(Vec>); + /// Running a task may consume the core. If the core is still available when /// running the task completes, it is returned. Otherwise, the worker will need /// to stop processing. type RunResult = Result, ()>; -type NextTaskResult = Result<(Option, Box), ()>; /// A task handle type Task = task::Task>; @@ -258,49 +240,48 @@ type Notified = task::Notified>; const MAX_LIFO_POLLS_PER_TICK: usize = 3; pub(super) fn create( - num_cores: usize, - driver: Driver, + size: usize, + park: Parker, driver_handle: driver::Handle, blocking_spawner: blocking::Spawner, seed_generator: RngSeedGenerator, config: Config, -) -> runtime::Handle { - // Allocate num_cores + 1 workers so that one worker can handle the I/O - // driver, if needed. - let num_workers = num_cores + 1; - let mut cores = Vec::with_capacity(num_cores); - let mut remotes = Vec::with_capacity(num_cores); - // Worker metrics are actually core based - let mut worker_metrics = Vec::with_capacity(num_cores); +) -> (Arc, Launch) { + let mut cores = Vec::with_capacity(size); + let mut remotes = Vec::with_capacity(size); + let mut worker_metrics = Vec::with_capacity(size); // Create the local queues - for i in 0..num_cores { + for _ in 0..size { let (steal, run_queue) = queue::local(); + let park = park.clone(); + let unpark = park.unpark(); let metrics = WorkerMetrics::from_config(&config); let stats = Stats::new(&metrics); cores.push(Box::new(Core { - index: i, + tick: 0, lifo_slot: None, + lifo_enabled: !config.disable_lifo_slot, run_queue, is_searching: false, + is_shutdown: false, + is_traced: false, + park: Some(park), + global_queue_interval: stats.tuned_global_queue_interval(&config), stats, rand: FastRand::from_seed(config.seed_generator.next_seed()), })); - remotes.push(Remote { - steal, - // lifo_slot: Lifo::new(), - }); + remotes.push(Remote { steal, unpark }); worker_metrics.push(metrics); } - // Allocate num-cores + 1 workers, so one worker can handle the I/O driver, - // if needed. - let (idle, idle_synced) = Idle::new(cores, num_workers); + let (idle, idle_synced) = Idle::new(size); let (inject, inject_synced) = inject::Shared::new(); + let remotes_len = remotes.len(); let handle = Arc::new(Handle { shared: Shared { remotes: remotes.into_boxed_slice(), @@ -308,14 +289,11 @@ pub(super) fn create( idle, owned: OwnedTasks::new(), synced: Mutex::new(Synced { - assigned_cores: (0..num_workers).map(|_| None).collect(), - shutdown_cores: Vec::with_capacity(num_cores), idle: idle_synced, inject: inject_synced, }), - driver: AtomicCell::new(Some(Box::new(driver))), - condvars: (0..num_workers).map(|_| Condvar::new()).collect(), - trace_status: TraceStatus::new(num_cores), + shutdown_cores: Mutex::new(vec![]), + trace_status: TraceStatus::new(remotes_len), config, scheduler_metrics: SchedulerMetrics::new(), worker_metrics: worker_metrics.into_boxed_slice(), @@ -326,22 +304,17 @@ pub(super) fn create( seed_generator, }); - let rt_handle = runtime::Handle { - inner: scheduler::Handle::MultiThread(handle), - }; + let mut launch = Launch(vec![]); - // Eagerly start worker threads - for index in 0..num_workers { - let handle = rt_handle.inner.expect_multi_thread(); - let h2 = handle.clone(); - let handoff_core = Arc::new(AtomicCell::new(None)); - - handle - .blocking_spawner - .spawn_blocking(&rt_handle, move || run(index, h2, handoff_core, false)); + for (index, core) in cores.drain(..).enumerate() { + launch.0.push(Arc::new(Worker { + handle: handle.clone(), + index, + core: AtomicCell::new(Some(core)), + })); } - rt_handle + (handle, launch) } #[track_caller] @@ -356,7 +329,7 @@ where fn drop(&mut self) { with_current(|maybe_cx| { if let Some(cx) = maybe_cx { - let core = cx.handoff_core.take(); + let core = cx.worker.core.take(); let mut cx_core = cx.core.borrow_mut(); assert!(cx_core.is_none()); *cx_core = core; @@ -421,21 +394,22 @@ where None => return Ok(()), }; + // The parker should be set here + assert!(core.park.is_some()); + // In order to block, the core must be sent to another thread for // execution. // // First, move the core back into the worker's shared core slot. - cx.handoff_core.set(core); + cx.worker.core.set(core); // Next, clone the worker handle and send it to a new thread for // processing. // // Once the blocking task is done executing, we will attempt to // steal the core back. - let index = cx.index; - let handle = cx.handle.clone(); - let handoff_core = cx.handoff_core.clone(); - runtime::spawn_blocking(move || run(index, handle, handoff_core, true)); + let worker = cx.worker.clone(); + runtime::spawn_blocking(move || run(worker)); Ok(()) }); @@ -454,12 +428,15 @@ where } } -fn run( - index: usize, - handle: Arc, - handoff_core: Arc>, - blocking_in_place: bool, -) { +impl Launch { + pub(crate) fn launch(mut self) { + for worker in self.0.drain(..) { + runtime::spawn_blocking(move || run(worker)); + } + } +} + +fn run(worker: Arc) { struct AbortOnPanic; impl Drop for AbortOnPanic { @@ -476,444 +453,113 @@ fn run( #[cfg(debug_assertions)] let _abort_on_panic = AbortOnPanic; - let num_workers = handle.shared.condvars.len(); - - let mut worker = Worker { - tick: 0, - num_seq_local_queue_polls: 0, - global_queue_interval: Stats::DEFAULT_GLOBAL_QUEUE_INTERVAL, - is_shutdown: false, - is_traced: false, - workers_to_notify: Vec::with_capacity(num_workers - 1), - idle_snapshot: idle::Snapshot::new(&handle.shared.idle), - stats: stats::Ephemeral::new(), + // Acquire a core. If this fails, then another thread is running this + // worker and there is nothing further to do. + let core = match worker.core.take() { + Some(core) => core, + None => return, }; - let sched_handle = scheduler::Handle::MultiThread(handle.clone()); + let handle = scheduler::Handle::MultiThread(worker.handle.clone()); - crate::runtime::context::enter_runtime(&sched_handle, true, |_| { + crate::runtime::context::enter_runtime(&handle, true, |_| { // Set the worker context. let cx = scheduler::Context::MultiThread(Context { - index, - lifo_enabled: Cell::new(!handle.shared.config.disable_lifo_slot), - handle, + worker, core: RefCell::new(None), - handoff_core, - defer: RefCell::new(Vec::with_capacity(64)), + defer: Defer::new(), }); context::set_scheduler(&cx, || { let cx = cx.expect_multi_thread(); - // Run the worker - let res = worker.run(&cx, blocking_in_place); - // `err` here signifies the core was lost, this is an expected end - // state for a worker. - debug_assert!(res.is_err()); + // This should always be an error. It only returns a `Result` to support + // using `?` to short circuit. + assert!(cx.run(core).is_err()); // Check if there are any deferred tasks to notify. This can happen when // the worker core is lost due to `block_in_place()` being called from // within the task. - if !cx.defer.borrow().is_empty() { - worker.schedule_deferred_without_core(&cx, &mut cx.shared().synced.lock()); - } + cx.defer.wake(); }); }); } -macro_rules! try_task { - ($e:expr) => {{ - let (task, core) = $e?; - if task.is_some() { - return Ok((task, core)); - } - core - }}; -} +impl Context { + fn run(&self, mut core: Box) -> RunResult { + // Reset `lifo_enabled` here in case the core was previously stolen from + // a task that had the LIFO slot disabled. + self.reset_lifo_enabled(&mut core); -macro_rules! try_task_new_batch { - ($w:expr, $e:expr) => {{ - let (task, mut core) = $e?; - if task.is_some() { - core.stats.start_processing_scheduled_tasks(&mut $w.stats); - return Ok((task, core)); - } - core - }}; -} + // Start as "processing" tasks as polling tasks from the local queue + // will be one of the first things we do. + core.stats.start_processing_scheduled_tasks(); -impl Worker { - fn run(&mut self, cx: &Context, blocking_in_place: bool) -> RunResult { - let (maybe_task, mut core) = { - if blocking_in_place { - if let Some(core) = cx.handoff_core.take() { - (None, core) - } else { - // Just shutdown - return Err(()); - } - } else { - let mut synced = cx.shared().synced.lock(); + while !core.is_shutdown { + self.assert_lifo_enabled_is_correct(&core); - // First try to acquire an available core - if let Some(core) = self.try_acquire_available_core(cx, &mut synced) { - // Try to poll a task from the global queue - let maybe_task = self.next_remote_task_synced(cx, &mut synced); - (maybe_task, core) - } else { - // block the thread to wait for a core to be assinged to us - self.wait_for_core(cx, synced)? - } + if core.is_traced { + core = self.worker.handle.trace_core(core); } - }; - core.stats.start_processing_scheduled_tasks(&mut self.stats); + // Increment the tick + core.tick(); - if let Some(task) = maybe_task { - core = self.run_task(cx, core, task)?; - } + // Run maintenance, if needed + core = self.maintenance(core); - while !self.is_shutdown { - let (maybe_task, c) = self.next_task(cx, core)?; - core = c; + // First, check work available to the current worker. + if let Some(task) = core.next_task(&self.worker) { + core = self.run_task(task, core)?; + continue; + } - if let Some(task) = maybe_task { - core = self.run_task(cx, core, task)?; + // We consumed all work in the queues and will start searching for work. + core.stats.end_processing_scheduled_tasks(); + + // There is no more **local** work to process, try to steal work + // from other workers. + if let Some(task) = core.steal_work(&self.worker) { + // Found work, switch back to processing + core.stats.start_processing_scheduled_tasks(); + core = self.run_task(task, core)?; } else { - // The only reason to get `None` from `next_task` is we have - // entered the shutdown phase. - assert!(self.is_shutdown); - break; + // Wait for work + core = if !self.defer.is_empty() { + self.park_timeout(core, Some(Duration::from_millis(0))) + } else { + self.park(core) + }; } } - self.pre_shutdown(cx, &mut core); + core.pre_shutdown(&self.worker); // Signal shutdown - self.shutdown_core(cx, core); - - // It is possible that tasks wake others during drop, so we need to - // clear the defer list. - self.shutdown_clear_defer(cx); - + self.worker.handle.shutdown_core(core); Err(()) } - // Try to acquire an available core, but do not block the thread - fn try_acquire_available_core( - &mut self, - cx: &Context, - synced: &mut Synced, - ) -> Option> { - if let Some(mut core) = cx - .shared() - .idle - .try_acquire_available_core(&mut synced.idle) - { - self.reset_acquired_core(cx, synced, &mut core); - Some(core) - } else { - None - } - } - - // Block the current thread, waiting for an available core - fn wait_for_core( - &mut self, - cx: &Context, - mut synced: MutexGuard<'_, Synced>, - ) -> NextTaskResult { - cx.shared() - .idle - .transition_worker_to_parked(&mut synced, cx.index); - - // Wait until a core is available, then exit the loop. - let mut core = loop { - if let Some(core) = synced.assigned_cores[cx.index].take() { - break core; - } - - // If shutting down, abort - if cx.shared().inject.is_closed(&synced.inject) { - self.shutdown_clear_defer(cx); - return Err(()); - } - - synced = cx.shared().condvars[cx.index].wait(synced).unwrap(); - }; - - self.reset_acquired_core(cx, &mut synced, &mut core); - - if self.is_shutdown { - // Currently shutting down, don't do any more work - return Ok((None, core)); - } - - let n = core.run_queue.max_capacity() / 2; - let maybe_task = self.next_remote_task_batch_synced(cx, &mut synced, &mut core, n); - - Ok((maybe_task, core)) - } - - /// Ensure core's state is set correctly for the worker to start using. - fn reset_acquired_core(&mut self, cx: &Context, synced: &mut Synced, core: &mut Core) { - self.global_queue_interval = core.stats.tuned_global_queue_interval(&cx.shared().config); - debug_assert!(self.global_queue_interval > 1); - - // Reset `lifo_enabled` here in case the core was previously stolen from - // a task that had the LIFO slot disabled. - self.reset_lifo_enabled(cx); - - // At this point, the local queue should be empty - debug_assert!(core.run_queue.is_empty()); - - // Update shutdown state while locked - self.update_global_flags(cx, synced); - } - - /// Finds the next task to run, this could be from a queue or stealing. If - /// none are available, the thread sleeps and tries again. - fn next_task(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { - self.assert_lifo_enabled_is_correct(cx); - - if self.is_traced { - core = cx.handle.trace_core(core); - } - - // Increment the tick - self.tick = self.tick.wrapping_add(1); - - // Runs maintenance every so often. When maintenance is run, the - // driver is checked, which may result in a task being found. - core = try_task!(self.maybe_maintenance(&cx, core)); - - // Check the LIFO slot, local run queue, and the injection queue for - // a notified task. - core = try_task!(self.next_notified_task(cx, core)); - - // We consumed all work in the queues and will start searching for work. - core.stats.end_processing_scheduled_tasks(&mut self.stats); - - super::counters::inc_num_no_local_work(); - - if !cx.defer.borrow().is_empty() { - // We are deferring tasks, so poll the resource driver and schedule - // the deferred tasks. - try_task_new_batch!(self, self.park_yield(cx, core)); - - panic!("what happened to the deferred tasks? 🤔"); - } - - while !self.is_shutdown { - // Search for more work, this involves trying to poll the resource - // driver, steal from other workers, and check the global queue - // again. - core = try_task_new_batch!(self, self.search_for_work(cx, core)); - - debug_assert!(cx.defer.borrow().is_empty()); - core = try_task_new_batch!(self, self.park(cx, core)); - } - - // Shutting down, drop any deferred tasks - self.shutdown_clear_defer(cx); - - Ok((None, core)) - } - - fn next_notified_task(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { - self.num_seq_local_queue_polls += 1; - - if self.num_seq_local_queue_polls % self.global_queue_interval == 0 { - super::counters::inc_global_queue_interval(); - - self.num_seq_local_queue_polls = 0; - - // Update the global queue interval, if needed - self.tune_global_queue_interval(cx, &mut core); - - if let Some(task) = self.next_remote_task(cx) { - return Ok((Some(task), core)); - } - } - - if let Some(task) = self.next_local_task(&mut core) { - return Ok((Some(task), core)); - } - - self.next_remote_task_batch(cx, core) - } - - fn next_remote_task(&self, cx: &Context) -> Option { - if cx.shared().inject.is_empty() { - return None; - } - - let mut synced = cx.shared().synced.lock(); - self.next_remote_task_synced(cx, &mut synced) - } - - fn next_remote_task_synced(&self, cx: &Context, synced: &mut Synced) -> Option { - // safety: we only have access to a valid `Synced` in this file. - unsafe { cx.shared().inject.pop(&mut synced.inject) } - } - - fn next_remote_task_batch(&self, cx: &Context, mut core: Box) -> NextTaskResult { - if cx.shared().inject.is_empty() { - return Ok((None, core)); - } - - // 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` - // available slots in the queue. - let cap = usize::min( - core.run_queue.remaining_slots(), - core.run_queue.max_capacity() / 2, - ); - - let mut synced = cx.shared().synced.lock(); - let maybe_task = self.next_remote_task_batch_synced(cx, &mut synced, &mut core, cap); - Ok((maybe_task, core)) - } - - fn next_remote_task_batch_synced( - &self, - cx: &Context, - synced: &mut Synced, - core: &mut Core, - max: usize, - ) -> Option { - super::counters::inc_num_remote_batch(); - - // The worker is currently idle, pull a batch of work from the - // injection queue. We don't want to pull *all* the work so other - // workers can also get some. - let n = if core.is_searching { - cx.shared().inject.len() / cx.shared().idle.num_searching() + 1 - } else { - cx.shared().inject.len() / cx.shared().remotes.len() + 1 - }; - - let n = usize::min(n, max); - - // safety: passing in the correct `inject::Synced`. - let mut tasks = unsafe { cx.shared().inject.pop_n(&mut synced.inject, n) }; - - // Pop the first task to return immedietly - let ret = tasks.next(); - - // Push the rest of the on the run queue - core.run_queue.push_back(tasks); - - ret - } - - fn next_local_task(&self, core: &mut Core) -> Option { - self.next_lifo_task(core).or_else(|| core.run_queue.pop()) - } - - fn next_lifo_task(&self, core: &mut Core) -> Option { - core.lifo_slot.take() - } - - /// Function responsible for stealing tasks from another worker - /// - /// Note: Only if less than half the workers are searching for tasks to steal - /// a new worker will actually try to steal. The idea is to make sure not all - /// workers will be trying to steal at the same time. - fn search_for_work(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { - #[cfg(not(loom))] - const ROUNDS: usize = 1; - - #[cfg(loom)] - const ROUNDS: usize = 1; - - debug_assert!(core.lifo_slot.is_none()); - debug_assert!(core.run_queue.is_empty()); - - if !self.transition_to_searching(cx, &mut core) { - return Ok((None, core)); - } - - // core = try_task!(self, self.poll_driver(cx, core)); - - // Get a snapshot of which workers are idle - cx.shared().idle.snapshot(&mut self.idle_snapshot); - - let num = cx.shared().remotes.len(); - - for i in 0..ROUNDS { - // Start from a random worker - let start = core.rand.fastrand_n(num as u32) as usize; - - if let Some(task) = self.steal_one_round(cx, &mut core, start) { - return Ok((Some(task), core)); - } - - core = try_task!(self.next_remote_task_batch(cx, core)); - - if i > 0 { - super::counters::inc_num_spin_stall(); - std::thread::sleep(std::time::Duration::from_micros(i as u64)); - } - } - - Ok((None, core)) - } - - fn steal_one_round(&self, cx: &Context, core: &mut Core, start: usize) -> Option { - let num = cx.shared().remotes.len(); - - for i in 0..num { - let i = (start + i) % num; - - // Don't steal from ourself! We know we don't have work. - if i == core.index { - continue; - } - - // If the core is currently idle, then there is nothing to steal. - if self.idle_snapshot.is_idle(i) { - continue; - } - - let target = &cx.shared().remotes[i]; - - if let Some(task) = target - .steal - .steal_into(&mut core.run_queue, &mut core.stats) - { - return Some(task); - } - } - - None - } - - fn run_task(&mut self, cx: &Context, mut core: Box, task: Notified) -> RunResult { - let task = cx.shared().owned.assert_owner(task); + fn run_task(&self, task: Notified, mut core: Box) -> RunResult { + let task = self.worker.handle.shared.owned.assert_owner(task); // Make sure the worker is not in the **searching** state. This enables // another idle worker to try to steal work. - if self.transition_from_searching(cx, &mut core) { - super::counters::inc_num_relay_search(); - cx.shared().notify_parked_local(); - } + core.transition_from_searching(&self.worker); - self.assert_lifo_enabled_is_correct(cx); + self.assert_lifo_enabled_is_correct(&core); // Measure the poll start time. Note that we may end up polling other // tasks under this measurement. In this case, the tasks came from the // LIFO slot and are considered part of the current task for scheduling // purposes. These tasks inherent the "parent"'s limits. - core.stats.start_poll(&mut self.stats); + core.stats.start_poll(); // Make the core available to the runtime context - *cx.core.borrow_mut() = Some(core); + *self.core.borrow_mut() = Some(core); // Run the task coop::budget(|| { - super::counters::inc_num_polls(); task.run(); let mut lifo_polls = 0; @@ -922,7 +568,7 @@ impl Worker { loop { // Check if we still have the core. If not, the core was stolen // by another worker. - let mut core = match cx.core.borrow_mut().take() { + let mut core = match self.core.borrow_mut().take() { Some(core) => core, None => { // In this case, we cannot call `reset_lifo_enabled()` @@ -933,10 +579,10 @@ impl Worker { }; // Check for a task in the LIFO slot - let task = match self.next_lifo_task(&mut core) { + let task = match core.lifo_slot.take() { Some(task) => task, None => { - self.reset_lifo_enabled(cx); + self.reset_lifo_enabled(&mut core); core.stats.end_poll(); return Ok(core); } @@ -947,11 +593,14 @@ impl Worker { // Not enough budget left to run the LIFO task, push it to // the back of the queue and return. - core.run_queue - .push_back_or_overflow(task, cx.shared(), &mut core.stats); + core.run_queue.push_back_or_overflow( + task, + &*self.worker.handle, + &mut core.stats, + ); // If we hit this point, the LIFO slot should be enabled. // There is no need to reset it. - debug_assert!(cx.lifo_enabled.get()); + debug_assert!(core.lifo_enabled); return Ok(core); } @@ -967,356 +616,349 @@ impl Worker { // repeatedly schedule the other. To mitigate this, we limit the // number of times the LIFO slot is prioritized. if lifo_polls >= MAX_LIFO_POLLS_PER_TICK { - cx.lifo_enabled.set(false); + core.lifo_enabled = false; super::counters::inc_lifo_capped(); } // Run the LIFO task, then loop - *cx.core.borrow_mut() = Some(core); - let task = cx.shared().owned.assert_owner(task); - super::counters::inc_num_lifo_polls(); + *self.core.borrow_mut() = Some(core); + let task = self.worker.handle.shared.owned.assert_owner(task); task.run(); } }) } - fn schedule_deferred_with_core<'a>( - &mut self, - cx: &'a Context, - mut core: Box, - synced: impl FnOnce() -> MutexGuard<'a, Synced>, - ) -> NextTaskResult { - let mut defer = cx.defer.borrow_mut(); - - // Grab a task to run next - let task = defer.pop(); - - if task.is_none() { - return Ok((None, core)); - } - - if !defer.is_empty() { - let mut synced = synced(); - - // Number of tasks we want to try to spread across idle workers - let num_fanout = cmp::min(defer.len(), cx.shared().idle.num_idle(&synced.idle)); - - if num_fanout > 0 { - cx.shared() - .push_remote_task_batch_synced(&mut synced, defer.drain(..num_fanout)); - - cx.shared() - .idle - .notify_mult(&mut synced, &mut self.workers_to_notify, num_fanout); - } - - // Do not run the task while holding the lock... - drop(synced); - } - - // Notify any workers - for worker in self.workers_to_notify.drain(..) { - cx.shared().condvars[worker].notify_one() - } - - if !defer.is_empty() { - // Push the rest of the tasks on the local queue - for task in defer.drain(..) { - core.run_queue - .push_back_or_overflow(task, cx.shared(), &mut core.stats); - } - - cx.shared().notify_parked_local(); - } - - Ok((task, core)) + fn reset_lifo_enabled(&self, core: &mut Core) { + core.lifo_enabled = !self.worker.handle.shared.config.disable_lifo_slot; } - fn schedule_deferred_without_core<'a>(&mut self, cx: &Context, synced: &mut Synced) { - let mut defer = cx.defer.borrow_mut(); - let num = defer.len(); - - if num > 0 { - // Push all tasks to the injection queue - cx.shared() - .push_remote_task_batch_synced(synced, defer.drain(..)); - - debug_assert!(self.workers_to_notify.is_empty()); - - // Notify workers - cx.shared() - .idle - .notify_mult(synced, &mut self.workers_to_notify, num); - - // Notify any workers - for worker in self.workers_to_notify.drain(..) { - cx.shared().condvars[worker].notify_one() - } - } + fn assert_lifo_enabled_is_correct(&self, core: &Core) { + debug_assert_eq!( + core.lifo_enabled, + !self.worker.handle.shared.config.disable_lifo_slot + ); } - fn maybe_maintenance(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { - if self.tick % cx.shared().config.event_interval == 0 { + fn maintenance(&self, mut core: Box) -> Box { + if core.tick % self.worker.handle.shared.config.event_interval == 0 { super::counters::inc_num_maintenance(); - core.stats.end_processing_scheduled_tasks(&mut self.stats); + core.stats.end_processing_scheduled_tasks(); + + // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... + // to run without actually putting the thread to sleep. + core = self.park_timeout(core, Some(Duration::from_millis(0))); // Run regularly scheduled maintenance - core = try_task_new_batch!(self, self.park_yield(cx, core)); + core.maintenance(&self.worker); - core.stats.start_processing_scheduled_tasks(&mut self.stats); + core.stats.start_processing_scheduled_tasks(); } - Ok((None, core)) + core } - fn flush_metrics(&self, cx: &Context, core: &mut Core) { - core.stats.submit(&cx.shared().worker_metrics[core.index]); - } - - fn update_global_flags(&mut self, cx: &Context, synced: &mut Synced) { - if !self.is_shutdown { - self.is_shutdown = cx.shared().inject.is_closed(&synced.inject); - } - - if !self.is_traced { - self.is_traced = cx.shared().trace_status.trace_requested(); - } - } - - fn park_yield(&mut self, cx: &Context, core: Box) -> NextTaskResult { - // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... - // to run without actually putting the thread to sleep. - if let Some(mut driver) = cx.shared().driver.take() { - driver.park_timeout(&cx.handle.driver, Duration::from_millis(0)); - - cx.shared().driver.set(driver); - } - - // If there are more I/O events, schedule them. - let (maybe_task, mut core) = - self.schedule_deferred_with_core(cx, core, || cx.shared().synced.lock())?; - - self.flush_metrics(cx, &mut core); - self.update_global_flags(cx, &mut cx.shared().synced.lock()); - - Ok((maybe_task, core)) - } - - /* - fn poll_driver(&mut self, cx: &Context, core: Box) -> NextTaskResult { - // Call `park` with a 0 timeout. This enables the I/O driver, timer, ... - // to run without actually putting the thread to sleep. - if let Some(mut driver) = cx.shared().driver.take() { - driver.park_timeout(&cx.handle.driver, Duration::from_millis(0)); - - cx.shared().driver.set(driver); - - // If there are more I/O events, schedule them. - self.schedule_deferred_with_core(cx, core, || cx.shared().synced.lock()) - } else { - Ok((None, core)) - } - } - */ - - fn park(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { - if let Some(f) = &cx.shared().config.before_park { + /// Parks the worker thread while waiting for tasks to execute. + /// + /// This function checks if indeed there's no more work left to be done before parking. + /// Also important to notice that, before parking, the worker thread will try to take + /// ownership of the Driver (IO/Time) and dispatch any events that might have fired. + /// Whenever a worker thread executes the Driver loop, all waken tasks are scheduled + /// in its own local queue until the queue saturates (ntasks > LOCAL_QUEUE_CAPACITY). + /// When the local queue is saturated, the overflow tasks are added to the injection queue + /// from where other workers can pick them up. + /// 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 { + if let Some(f) = &self.worker.handle.shared.config.before_park { f(); } - if self.can_transition_to_parked(&mut core) { - debug_assert!(!self.is_shutdown); - debug_assert!(!self.is_traced); + if core.transition_to_parked(&self.worker) { + while !core.is_shutdown && !core.is_traced { + core.stats.about_to_park(); + core = self.park_timeout(core, None); - core = try_task!(self.do_park(cx, core)); - } + // Run regularly scheduled maintenance + core.maintenance(&self.worker); - if let Some(f) = &cx.shared().config.after_unpark { - f(); - } - - Ok((None, core)) - } - - fn do_park(&mut self, cx: &Context, mut core: Box) -> NextTaskResult { - let was_searching = core.is_searching; - - // Before we park, if we are searching, we need to transition away from searching - if self.transition_from_searching(cx, &mut core) { - cx.shared().idle.snapshot(&mut self.idle_snapshot); - // We were the last searching worker, we need to do one last check - if let Some(task) = self.steal_one_round(cx, &mut core, 0) { - cx.shared().notify_parked_local(); - - return Ok((Some(task), core)); + if core.transition_from_parked(&self.worker) { + break; + } } } - // Acquire the lock - let mut synced = cx.shared().synced.lock(); + if let Some(f) = &self.worker.handle.shared.config.after_unpark { + f(); + } + core + } - // Try one last time to get tasks - let n = core.run_queue.max_capacity() / 2; - if let Some(task) = self.next_remote_task_batch_synced(cx, &mut synced, &mut core, n) { - return Ok((Some(task), core)); + fn park_timeout(&self, mut core: Box, duration: Option) -> Box { + self.assert_lifo_enabled_is_correct(&core); + + // Take the parker out of core + let mut park = core.park.take().expect("park missing"); + + // Store `core` in context + *self.core.borrow_mut() = Some(core); + + // Park thread + if let Some(timeout) = duration { + park.park_timeout(&self.worker.handle.driver, timeout); + } else { + park.park(&self.worker.handle.driver); } - if !was_searching { - if cx - .shared() - .idle - .transition_worker_to_searching_if_needed(&mut synced.idle, &mut core) + self.defer.wake(); + + // Remove `core` from context + core = self.core.borrow_mut().take().expect("core missing"); + + // Place `park` back in `core` + core.park = Some(park); + + if core.should_notify_others() { + self.worker.handle.notify_parked_local(); + } + + core + } + + pub(crate) fn defer(&self, waker: &Waker) { + self.defer.defer(waker); + } +} + +impl Core { + /// Increment the tick + fn tick(&mut self) { + self.tick = self.tick.wrapping_add(1); + } + + /// Return the next notified task available to this worker. + fn next_task(&mut self, worker: &Worker) -> Option { + if self.tick % self.global_queue_interval == 0 { + // Update the global queue interval, if needed + self.tune_global_queue_interval(worker); + + worker + .handle + .next_remote_task() + .or_else(|| self.next_local_task()) + } else { + let maybe_task = self.next_local_task(); + + if maybe_task.is_some() { + 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` + // available slots in the queue. + let cap = usize::min( + self.run_queue.remaining_slots(), + self.run_queue.max_capacity() / 2, + ); + + // The worker is currently idle, pull a batch of work from the + // injection queue. We don't want to pull *all* the work so other + // workers can also get some. + let n = usize::min( + worker.inject().len() / worker.handle.shared.remotes.len() + 1, + cap, + ); + + 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(); + + // Push the rest of the on the run queue + self.run_queue.push_back(tasks); + + ret + } + } + + fn next_local_task(&mut self) -> Option { + self.lifo_slot.take().or_else(|| self.run_queue.pop()) + } + + /// Function responsible for stealing tasks from another worker + /// + /// Note: Only if less than half the workers are searching for tasks to steal + /// a new worker will actually try to steal. The idea is to make sure not all + /// workers will be trying to steal at the same time. + fn steal_work(&mut self, worker: &Worker) -> Option { + if !self.transition_to_searching(worker) { + return None; + } + + let num = worker.handle.shared.remotes.len(); + // Start from a random worker + let start = self.rand.fastrand_n(num as u32) as usize; + + for i in 0..num { + let i = (start + i) % num; + + // Don't steal from ourself! We know we don't have work. + if i == worker.index { + continue; + } + + let target = &worker.handle.shared.remotes[i]; + if let Some(task) = target + .steal + .steal_into(&mut self.run_queue, &mut self.stats) { - // Skip parking, go back to searching - return Ok((None, core)); + return Some(task); } } - super::counters::inc_num_parks(); - core.stats.about_to_park(); - // Flush metrics to the runtime metrics aggregator - self.flush_metrics(cx, &mut core); - - // If the runtime is shutdown, skip parking - self.update_global_flags(cx, &mut synced); - - if self.is_shutdown { - return Ok((None, core)); - } - - // Core being returned must not be in the searching state - debug_assert!(!core.is_searching); - - // Release the core - cx.shared().idle.release_core(&mut synced, core); - - if let Some(mut driver) = cx.shared().driver.take() { - // Drop the lock before parking on the driver - drop(synced); - - // Wait for driver events - driver.park(&cx.handle.driver); - - synced = cx.shared().synced.lock(); - - // Put the driver back - cx.shared().driver.set(driver); - - if cx.shared().inject.is_closed(&mut synced.inject) { - self.shutdown_clear_defer(cx); - self.shutdown_finalize(cx, synced); - return Err(()); - } - - // Try to acquire an available core to schedule I/O events - if let Some(core) = self.try_acquire_available_core(cx, &mut synced) { - // This may result in a task being run - self.schedule_deferred_with_core(cx, core, move || synced) - } else { - // Schedule any deferred tasks - self.schedule_deferred_without_core(cx, &mut synced); - - // Wait for a core. - self.wait_for_core(cx, synced) - } - } else { - // Wait for a core to be assigned to us - self.wait_for_core(cx, synced) - } + // Fallback on checking the global queue + worker.handle.next_remote_task() } - fn transition_to_searching(&self, cx: &Context, core: &mut Core) -> bool { - if !core.is_searching { - cx.shared().idle.try_transition_worker_to_searching(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(); } - core.is_searching + self.is_searching } - /// Returns `true` if another worker must be notified - fn transition_from_searching(&self, cx: &Context, core: &mut Core) -> bool { - if !core.is_searching { + fn transition_from_searching(&mut self, worker: &Worker) { + if !self.is_searching { + return; + } + + self.is_searching = false; + worker.handle.transition_worker_from_searching(); + } + + fn has_tasks(&self) -> bool { + self.lifo_slot.is_some() || self.run_queue.has_tasks() + } + + fn should_notify_others(&self) -> bool { + // If there are tasks available to steal, but this worker is not + // looking for tasks to steal, notify another worker. + if self.is_searching { + return false; + } + self.lifo_slot.is_some() as usize + self.run_queue.len() > 1 + } + + /// 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 self.has_tasks() || self.is_traced { return false; } - cx.shared().idle.transition_worker_from_searching(core) + // 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.handle.shared.idle.transition_worker_to_parked( + &worker.handle.shared, + worker.index, + self.is_searching, + ); + + // 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(); + } + + true } - fn can_transition_to_parked(&self, core: &mut Core) -> bool { - !self.has_tasks(core) && !self.is_shutdown && !self.is_traced + /// Returns `true` if the transition happened. + fn transition_from_parked(&mut self, worker: &Worker) -> bool { + // If a task is in the lifo slot/run queue, then we must unpark regardless of + // being notified + if self.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" + // when it wakes when the I/O driver receives new events. + self.is_searching = !worker + .handle + .shared + .idle + .unpark_worker_by_id(&worker.handle.shared, worker.index); + return true; + } + + if worker + .handle + .shared + .idle + .is_parked(&worker.handle.shared, worker.index) + { + return false; + } + + // When unparked, the worker is in the searching state. + self.is_searching = true; + true } - fn has_tasks(&self, core: &Core) -> bool { - core.lifo_slot.is_some() || !core.run_queue.is_empty() + /// Runs maintenance work such as checking the pool's state. + fn maintenance(&mut self, worker: &Worker) { + self.stats + .submit(&worker.handle.shared.worker_metrics[worker.index]); + + if !self.is_shutdown { + // Check if the scheduler has been shutdown + let synced = worker.handle.shared.synced.lock(); + self.is_shutdown = worker.inject().is_closed(&synced.inject); + } + + if !self.is_traced { + // Check if the worker should be tracing. + self.is_traced = worker.handle.shared.trace_status.trace_requested(); + } } /// Signals all tasks to shut down, and waits for them to complete. Must run /// before we enter the single-threaded phase of shutdown processing. - fn pre_shutdown(&self, cx: &Context, core: &mut Core) { + fn pre_shutdown(&mut self, worker: &Worker) { // Signal to all tasks to shut down. - cx.shared().owned.close_and_shutdown_all(); + worker.handle.shared.owned.close_and_shutdown_all(); - core.stats.submit(&cx.shared().worker_metrics[core.index]); + self.stats + .submit(&worker.handle.shared.worker_metrics[worker.index]); } - /// Signals that a worker has observed the shutdown signal and has replaced - /// its core back into its handle. - /// - /// If all workers have reached this point, the final cleanup is performed. - fn shutdown_core(&self, cx: &Context, core: Box) { - let mut synced = cx.shared().synced.lock(); - synced.shutdown_cores.push(core); + /// Shuts down the core. + fn shutdown(&mut self, handle: &Handle) { + // Take the core + let mut park = self.park.take().expect("park missing"); - self.shutdown_finalize(cx, synced); + // Drain the queue + while self.next_local_task().is_some() {} + + park.shutdown(&handle.driver); } - fn shutdown_finalize(&self, cx: &Context, mut synced: MutexGuard<'_, Synced>) { - // Wait for all cores - if synced.shutdown_cores.len() != cx.shared().remotes.len() { - return; - } - - let mut driver = match cx.shared().driver.take() { - Some(driver) => driver, - None => return, - }; - - debug_assert!(cx.shared().owned.is_empty()); - - for mut core in synced.shutdown_cores.drain(..) { - // Drain tasks from the local queue - while self.next_local_task(&mut core).is_some() {} - } - - // Shutdown the driver - driver.shutdown(&cx.handle.driver); - - // Drain the injection queue - // - // We already shut down every task, so we can simply drop the tasks. We - // cannot call `next_remote_task()` because we already hold the lock. - // - // safety: passing in correct `idle::Synced` - while let Some(task) = self.next_remote_task_synced(cx, &mut synced) { - drop(task); - } - } - - fn reset_lifo_enabled(&self, cx: &Context) { - cx.lifo_enabled - .set(!cx.handle.shared.config.disable_lifo_slot); - } - - fn assert_lifo_enabled_is_correct(&self, cx: &Context) { - debug_assert_eq!( - cx.lifo_enabled.get(), - !cx.handle.shared.config.disable_lifo_slot - ); - } - - fn tune_global_queue_interval(&mut self, cx: &Context, core: &mut Core) { - let next = core.stats.tuned_global_queue_interval(&cx.shared().config); + fn tune_global_queue_interval(&mut self, worker: &Worker) { + let next = self + .stats + .tuned_global_queue_interval(&worker.handle.shared.config); debug_assert!(next > 1); @@ -1325,166 +967,199 @@ impl Worker { self.global_queue_interval = next; } } +} - fn shutdown_clear_defer(&self, cx: &Context) { - let mut defer = cx.defer.borrow_mut(); - - for task in defer.drain(..) { - drop(task); - } - } -} - -impl Context { - pub(crate) fn defer(&self, waker: &Waker) { - // TODO: refactor defer across all runtimes - waker.wake_by_ref(); - } - - fn shared(&self) -> &Shared { - &self.handle.shared - } -} - -impl Shared { - pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { - use std::ptr; - - with_current(|maybe_cx| { - if let Some(cx) = maybe_cx { - // Make sure the task is part of the **current** scheduler. - if ptr::eq(self, &cx.handle.shared) { - // And the current thread still holds a core - if let Some(core) = cx.core.borrow_mut().as_mut() { - if is_yield { - cx.defer.borrow_mut().push(task); - } else { - self.schedule_local(cx, core, task); - } - } else { - // This can happen if either the core was stolen - // (`block_in_place`) or the notification happens from - // the driver. - cx.defer.borrow_mut().push(task); - } - return; - } - } - - // Otherwise, use the inject queue. - self.schedule_remote(task); - }) - } - - fn schedule_local(&self, cx: &Context, core: &mut Core, task: Notified) { - core.stats.inc_local_schedule_count(); - - if cx.lifo_enabled.get() { - // Push to the LIFO slot - let prev = std::mem::replace(&mut core.lifo_slot, Some(task)); - // let prev = cx.shared().remotes[core.index].lifo_slot.swap_local(task); - - if let Some(prev) = prev { - core.run_queue - .push_back_or_overflow(prev, self, &mut core.stats); - } else { - return; - } - } else { - core.run_queue - .push_back_or_overflow(task, self, &mut core.stats); - } - - self.notify_parked_local(); - } - - fn notify_parked_local(&self) { - super::counters::inc_num_inc_notify_local(); - self.idle.notify_local(self); - } - - fn schedule_remote(&self, task: Notified) { - super::counters::inc_num_notify_remote(); - self.scheduler_metrics.inc_remote_schedule_count(); - - let mut synced = self.synced.lock(); - // Push the task in the - self.push_remote_task(&mut synced, task); - - // Notify a worker. The mutex is passed in and will be released as part - // of the method call. - self.idle.notify_remote(synced, self); - } - - pub(super) fn close(&self) { - let mut synced = self.synced.lock(); - - if self.inject.close(&mut synced.inject) { - // Set the shutdown flag on all available cores - self.idle.shutdown(&mut synced, self); - } - } - - fn push_remote_task(&self, synced: &mut Synced, task: Notified) { - // safety: passing in correct `idle::Synced` - unsafe { - self.inject.push(&mut synced.inject, task); - } - } - - fn push_remote_task_batch(&self, iter: I) - where - I: Iterator>>, - { - unsafe { - self.inject.push_batch(self, iter); - } - } - - fn push_remote_task_batch_synced(&self, synced: &mut Synced, iter: I) - where - I: Iterator>>, - { - unsafe { - self.inject.push_batch(&mut synced.inject, iter); - } - } -} - -impl Overflow> for Shared { - fn push(&self, task: task::Notified>) { - self.push_remote_task(&mut self.synced.lock(), task); - } - - fn push_batch(&self, iter: I) - where - I: Iterator>>, - { - self.push_remote_task_batch(iter) - } -} - -impl<'a> Lock for &'a Shared { - type Handle = InjectGuard<'a>; - - fn lock(self) -> Self::Handle { - InjectGuard { - lock: self.synced.lock(), - } +impl Worker { + /// Returns a reference to the scheduler's injection queue. + fn inject(&self) -> &inject::Shared> { + &self.handle.shared.inject } } +// TODO: Move `Handle` impls into handle.rs impl task::Schedule for Arc { fn release(&self, task: &Task) -> Option { self.shared.owned.remove(task) } fn schedule(&self, task: Notified) { - self.shared.schedule_task(task, false); + self.schedule_task(task, false); } fn yield_now(&self, task: Notified) { - self.shared.schedule_task(task, true); + self.schedule_task(task, true); + } +} + +impl Handle { + pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { + with_current(|maybe_cx| { + if let Some(cx) = maybe_cx { + // Make sure the task is part of the **current** scheduler. + if self.ptr_eq(&cx.worker.handle) { + // And the current thread still holds a core + if let Some(core) = cx.core.borrow_mut().as_mut() { + self.schedule_local(core, task, is_yield); + return; + } + } + } + + // Otherwise, use the inject queue. + self.push_remote_task(task); + self.notify_parked_remote(); + }) + } + + fn schedule_local(&self, core: &mut Core, task: Notified, is_yield: bool) { + core.stats.inc_local_schedule_count(); + + // Spawning from the worker thread. If scheduling a "yield" then the + // task must always be pushed to the back of the queue, enabling other + // tasks to be executed. If **not** a yield, then there is more + // 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, &mut core.stats); + true + } else { + // Push to the LIFO slot + let prev = core.lifo_slot.take(); + let ret = prev.is_some(); + + if let Some(prev) = prev { + core.run_queue + .push_back_or_overflow(prev, self, &mut core.stats); + } + + core.lifo_slot = Some(task); + + ret + }; + + // Only notify if not currently parked. If `park` is `None`, then the + // scheduling is from a resource driver. As notifications often come in + // batches, the notification is delayed until the park is complete. + if should_notify && core.park.is_some() { + self.notify_parked_local(); + } + } + + fn next_remote_task(&self) -> Option { + 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(&mut self.shared.synced.lock().inject) + { + self.notify_all(); + } + } + + fn notify_parked_local(&self) { + super::counters::inc_num_inc_notify_local(); + + if let Some(index) = self.shared.idle.worker_to_notify(&self.shared) { + super::counters::inc_num_unparks_local(); + self.shared.remotes[index].unpark.unpark(&self.driver); + } + } + + fn notify_parked_remote(&self) { + if let Some(index) = self.shared.idle.worker_to_notify(&self.shared) { + self.shared.remotes[index].unpark.unpark(&self.driver); + } + } + + pub(super) fn notify_all(&self) { + for remote in &self.shared.remotes[..] { + remote.unpark.unpark(&self.driver); + } + } + + fn notify_if_work_pending(&self) { + for remote in &self.shared.remotes[..] { + if !remote.steal.is_empty() { + self.notify_parked_local(); + return; + } + } + + if !self.shared.inject.is_empty() { + self.notify_parked_local(); + } + } + + 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_local(); + } + } + + /// Signals that a worker has observed the shutdown signal and has replaced + /// its core back into its handle. + /// + /// If all workers have reached this point, the final cleanup is performed. + fn shutdown_core(&self, core: Box) { + let mut cores = self.shared.shutdown_cores.lock(); + cores.push(core); + + if cores.len() != self.shared.remotes.len() { + return; + } + + debug_assert!(self.shared.owned.is_empty()); + + for mut core in cores.drain(..) { + core.shutdown(self); + } + + // Drain the injection queue + // + // We already shut down every task, so we can simply drop the tasks. + while let Some(task) = self.next_remote_task() { + drop(task); + } + } + + fn ptr_eq(&self, other: &Handle) -> bool { + std::ptr::eq(self, other) + } +} + +impl Overflow> for Handle { + fn push(&self, task: task::Notified>) { + self.push_remote_task(task); + } + + fn push_batch(&self, iter: I) + where + I: Iterator>>, + { + unsafe { + self.shared.inject.push_batch(self, iter); + } } } @@ -1498,6 +1173,16 @@ impl<'a> AsMut for InjectGuard<'a> { } } +impl<'a> Lock for &'a Handle { + type Handle = InjectGuard<'a>; + + fn lock(self) -> Self::Handle { + InjectGuard { + lock: self.shared.synced.lock(), + } + } +} + #[track_caller] fn with_current(f: impl FnOnce(Option<&Context>) -> R) -> R { use scheduler::Context::MultiThread; diff --git a/tokio/src/util/mod.rs b/tokio/src/util/mod.rs index ba0f6dd74..6a7d4b103 100644 --- a/tokio/src/util/mod.rs +++ b/tokio/src/util/mod.rs @@ -68,6 +68,11 @@ cfg_rt! { pub(crate) use rc_cell::RcCell; } +cfg_rt_multi_thread! { + mod try_lock; + pub(crate) use try_lock::TryLock; +} + pub(crate) mod trace; pub(crate) mod error;