diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 10de632d1..2c5ed150c 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -111,13 +111,6 @@ pub(super) struct Core { /// injection queue num_seq_local_queue_polls: 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, @@ -217,6 +210,13 @@ 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>, } @@ -246,6 +246,9 @@ type Task = task::Task>; /// A notified task handle type Notified = task::Notified>; +/// Stealable LIFO slot +type Lifo = task::AtomicCell>; + /// Value picked out of thin-air. Running the LIFO slot a handful of times /// seemms sufficient to benefit from locality. More than 3 times probably is /// overweighing. The value can be tuned in the future with data that shows @@ -279,7 +282,6 @@ pub(super) fn create( index: i, tick: 0, num_seq_local_queue_polls: 0, - lifo_slot: None, lifo_enabled: !config.disable_lifo_slot, run_queue, is_searching: false, @@ -290,7 +292,10 @@ pub(super) fn create( rand: FastRand::from_seed(config.seed_generator.next_seed()), })); - remotes.push(Remote { steal }); + remotes.push(Remote { + steal, + lifo_slot: Lifo::new(), + }); worker_metrics.push(metrics); } @@ -693,31 +698,30 @@ impl Worker { // Update the global queue interval, if needed self.tune_global_queue_interval(cx, core); - self.next_remote_task(cx) - .or_else(|| self.next_local_task(core)) - } else { - let maybe_task = self.next_local_task(core); - - if maybe_task.is_some() { - return maybe_task; + if let Some(task) = self.next_remote_task(cx) { + return Some(task); } - - if cx.shared().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( - core.run_queue.remaining_slots(), - core.run_queue.max_capacity() / 2, - ); - - let mut synced = cx.shared().synced.lock(); - self.next_remote_task_batch(cx, &mut synced, core, cap) } + + if let Some(task) = self.next_local_task(cx, core) { + return Some(task); + } + + if cx.shared().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( + core.run_queue.remaining_slots(), + core.run_queue.max_capacity() / 2, + ); + + let mut synced = cx.shared().synced.lock(); + self.next_remote_task_batch(cx, &mut synced, core, cap) } fn next_remote_task(&self, cx: &Context) -> Option { @@ -761,8 +765,11 @@ impl Worker { ret } - fn next_local_task(&self, core: &mut Core) -> Option { - core.lifo_slot.take().or_else(|| core.run_queue.pop()) + fn next_local_task(&self, cx: &Context, core: &mut Core) -> Option { + cx.shared().remotes[core.index] + .lifo_slot + .take_local() + .or_else(|| core.run_queue.pop()) } /// Function responsible for stealing tasks from another worker @@ -847,7 +854,7 @@ impl Worker { }; // Check for a task in the LIFO slot - let task = match core.lifo_slot.take() { + let task = match cx.shared().remotes[core.index].lifo_slot.take_local() { Some(task) => task, None => { self.reset_lifo_enabled(&mut core); @@ -1022,7 +1029,7 @@ impl Worker { f(); } - if self.can_transition_to_parked(&mut core) { + if self.can_transition_to_parked(cx, &mut core) { debug_assert!(!core.is_shutdown); debug_assert!(!core.is_traced); @@ -1142,8 +1149,10 @@ impl Worker { cx.shared().idle.transition_worker_from_searching(core) } - fn can_transition_to_parked(&self, core: &mut Core) -> bool { - core.lifo_slot.is_none() && core.run_queue.is_empty() && !core.is_traced + fn can_transition_to_parked(&self, cx: &Context, core: &mut Core) -> bool { + cx.shared().remotes[core.index].lifo_slot.is_none() + && core.run_queue.is_empty() + && !core.is_traced } /// Signals all tasks to shut down, and waits for them to complete. Must run @@ -1187,7 +1196,7 @@ impl Worker { for mut core in synced.shutdown_cores.drain(..) { // Drain tasks from the local queue - while self.next_local_task(&mut core).is_some() {} + while self.next_local_task(cx, &mut core).is_some() {} } // Shutdown the driver @@ -1260,7 +1269,7 @@ impl Shared { if is_yield { cx.defer.borrow_mut().push(task); } else { - self.schedule_local(core, task); + self.schedule_local(cx, core, task); } } else { // This can happen if either the core was stolen @@ -1277,17 +1286,15 @@ impl Shared { }) } - fn schedule_local(&self, core: &mut Core, task: Notified) { + fn schedule_local(&self, cx: &Context, core: &mut Core, task: Notified) { if core.lifo_enabled { // Push to the LIFO slot - let prev = core.lifo_slot.take(); + 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); } - - core.lifo_slot = Some(task); } else { core.run_queue .push_back_or_overflow(task, self, &mut core.stats); diff --git a/tokio/src/runtime/task/atomic_cell.rs b/tokio/src/runtime/task/atomic_cell.rs new file mode 100644 index 000000000..4ecfe108a --- /dev/null +++ b/tokio/src/runtime/task/atomic_cell.rs @@ -0,0 +1,99 @@ +use crate::loom::sync::atomic::AtomicPtr; +use crate::loom::sync::atomic::Ordering::{AcqRel, Acquire, Release}; +use crate::runtime::task::{Header, Notified, RawTask}; + +use std::marker::PhantomData; +use std::ptr::{self, NonNull}; + +pub(crate) struct AtomicCell { + task: AtomicPtr
, + _p: PhantomData, +} + +impl AtomicCell { + pub(crate) fn new() -> AtomicCell { + AtomicCell { + task: AtomicPtr::default(), + _p: PhantomData, + } + } + + /// Should be called from a local context + pub(crate) fn is_some(&self) -> bool { + !self.is_none() + } + + pub(crate) fn is_none(&self) -> bool { + self.task.load(Acquire).is_null() + } + + pub(crate) fn take_local(&self) -> Option> { + let ptr = self.task.load(Acquire); + + if ptr.is_null() { + return None; + } + + if self + .task + .compare_exchange(ptr, ptr::null_mut(), AcqRel, Acquire) + .is_err() + { + return None; + } + + NonNull::new(ptr).map(|ptr| unsafe { Notified::from_raw(RawTask::from_raw(ptr)) }) + } + + pub(crate) fn swap_local(&self, task: Notified) -> Option> { + let next = task.into_raw().header_ptr().as_ptr(); + let prev = self.task.load(Acquire); + + if prev.is_null() { + // Since this method is only called from the only thread that can + // set the value to !null, it is safe to use a store here. + self.task.store(next, Release); + return None; + } + + if self + .task + .compare_exchange(prev, next, AcqRel, Acquire) + .is_ok() + { + // Safety: we already checked !null above + let prev = + unsafe { Notified::from_raw(RawTask::from_raw(NonNull::new_unchecked(prev))) }; + return Some(prev); + } + + // The compare-exchanged failed, but there is no need to try again since + // this is the only thread that could set the cell to !null. + self.task.store(next, Release); + None + } + + pub(crate) fn take_remote(&self) -> Option> { + let task = self.task.load(Acquire); + + if task.is_null() { + return None; + } + + // std::thread::sleep(std::time::Duration::from_micros(3)); + + // Try to take it once + if self + .task + .compare_exchange(task, ptr::null_mut(), Acquire, Acquire) + .is_ok() + { + // safety: we checked for null above + return Some(unsafe { + Notified::from_raw(RawTask::from_raw(NonNull::new_unchecked(task))) + }); + } + + None + } +} diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 932552fb9..e9fec729a 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -168,6 +168,9 @@ // unstable. This should be removed once `JoinSet` is stabilized. #![cfg_attr(not(tokio_unstable), allow(dead_code))] +mod atomic_cell; +pub(crate) use atomic_cell::AtomicCell; + mod core; use self::core::Cell; use self::core::Header;