runtime: steal tasks from the LIFO slot (#7431)

Author: hawkw

tokio/src/runtime/builder.rs

@@ -1257,22 +1257,21 @@ impl Builder {
         /// scheduled task being polled first.
         ///
         /// To implement this heuristic, each worker thread has a slot which
-        /// holds the task that should be polled next. However, this slot cannot
-        /// be stolen by other worker threads, which can result in lower total
-        /// throughput when tasks tend to have longer poll times.
+        /// holds the task that should be polled next. In earlier versions of
+        /// Tokio, this slot could not be stolen by other worker threads, which
+        /// can result in lower total throughput when tasks tend to have longer
+        /// poll times.
         ///
         /// This configuration option will disable this heuristic resulting in
-        /// all scheduled tasks being pushed into the worker-local queue, which
-        /// is stealable.
-        ///
-        /// Consider trying this option when the task "scheduled" time is high
-        /// but the runtime is underutilized. Use [tokio-rs/tokio-metrics] to
-        /// collect this data.
+        /// all scheduled tasks being pushed into the worker-local queue. This
+        /// was intended as a workaround for the LIFO slot not being stealable.
+        /// As of Tokio 1.51, tasks can be stolen from the LIFO slot. In a
+        /// future version, this option may be deprecated.
         ///
         /// # Unstable
         ///
-        /// This configuration option is considered a workaround for the LIFO
-        /// slot not being stealable. When the slot becomes stealable, we will
+        /// This configuration option was considered a workaround for the LIFO
+        /// slot not being stealable. Since this is no longer the case, we will
         /// revisit whether or not this option is necessary. See
         /// issue [tokio-rs/tokio#4941].
         ///

tokio/src/runtime/config.rs

@@ -34,11 +34,13 @@ pub(crate) struct Config {
 
     /// The multi-threaded scheduler includes a per-worker LIFO slot used to
     /// store the last scheduled task. This can improve certain usage patterns,
-    /// especially message passing between tasks. However, this LIFO slot is not
-    /// currently stealable.
+    /// especially message passing between tasks.
     ///
-    /// Eventually, the LIFO slot **will** become stealable, however as a
-    /// stop-gap, this unstable option lets users disable the LIFO task.
+    /// In Tokio versions before 1.51, tasks in the LIFO slot could not be
+    /// stolen, which could cause issues in applications with long poll times.
+    /// As a stop-gap, this unstable option lets users disable the LIFO task.
+    /// Now that the LIFO slot is stealable, we may remove this option in a
+    /// future version.
     pub(crate) disable_lifo_slot: bool,
 
     /// Random number generator seed to configure runtimes to act in a

tokio/src/runtime/mod.rs

@@ -367,8 +367,8 @@
 //! three times in a row, it is temporarily disabled until the worker thread has
 //! scheduled a task that didn't come from the lifo slot. The lifo slot can be
 //! disabled using the [`disable_lifo_slot`] setting. The lifo slot is separate
-//! from the local queue, so other worker threads cannot steal the task in the
-//! lifo slot.
+//! from the local queue, and is stolen from by other worker threads only when
+//! a worker's local queue has been drained.
 //!
 //! When a task is woken from a thread that is not a worker thread, then the
 //! task is placed in the global queue.

tokio/src/runtime/scheduler/multi_thread/queue.rs

@@ -52,6 +52,13 @@ pub(crate) struct Inner<T: 'static> {
     /// Only updated by producer thread but read by many threads.
     tail: AtomicUnsignedShort,
 
+    /// 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: task::AtomicNotified<T>,
+
     /// Elements
     buffer: Box<[UnsafeCell<MaybeUninit<task::Notified<T>>>; LOCAL_QUEUE_CAPACITY]>,
 }
@@ -92,6 +99,7 @@ pub(crate) fn local<T: 'static>() -> (Steal<T>, Local<T>) {
     let inner = Arc::new(Inner {
         head: AtomicUnsignedLong::new(0),
         tail: AtomicUnsignedShort::new(0),
+        lifo: task::AtomicNotified::empty(),
         buffer: make_fixed_size(buffer.into_boxed_slice()),
     });
 
@@ -108,9 +116,10 @@ impl<T> Local<T> {
     /// Returns the number of entries in the queue
     pub(crate) fn len(&self) -> usize {
         let (_, head) = unpack(self.inner.head.load(Acquire));
+        let lifo = self.inner.lifo.is_some() as usize;
         // safety: this is the **only** thread that updates this cell.
         let tail = unsafe { self.inner.tail.unsync_load() };
-        len(head, tail)
+        len(head, tail) + lifo
     }
 
     /// How many tasks can be pushed into the queue
@@ -388,14 +397,28 @@ impl<T> Local<T> {
 
         Some(self.inner.buffer[idx].with(|ptr| unsafe { ptr::read(ptr).assume_init() }))
     }
+
+    /// Pushes a task to the LIFO slot, returning the task previously in the
+    /// LIFO slot (if there was one).
+    pub(crate) fn push_lifo(&self, task: task::Notified<T>) -> Option<task::Notified<T>> {
+        self.inner.lifo.swap(Some(task))
+    }
+
+    /// Pops the task currently held in the LIFO slot, if there is one;
+    /// otherwise, returns `None`.
+    pub(crate) fn pop_lifo(&self) -> Option<task::Notified<T>> {
+        // LIFO-suction!
+        self.inner.lifo.take()
+    }
 }
 
 impl<T> Steal<T> {
     /// Returns the number of entries in the queue
     pub(crate) fn len(&self) -> usize {
         let (_, head) = unpack(self.0.head.load(Acquire));
         let tail = self.0.tail.load(Acquire);
-        len(head, tail)
+        let lifo = self.0.lifo.is_some() as usize;
+        len(head, tail) + lifo
     }
 
     /// Return true if the queue is empty,
@@ -430,8 +453,14 @@ impl<T> Steal<T> {
         let mut n = self.steal_into2(dst, dst_tail);
 
         if n == 0 {
-            // No tasks were stolen
-            return None;
+            // If no tasks were stolen, let's see if there's one in the LIFO
+            // slot.
+            let lifo = self.0.lifo.take();
+            if lifo.is_some() {
+                dst_stats.incr_steal_count(1);
+                dst_stats.incr_steal_operations();
+            }
+            return lifo;
         }
 
         dst_stats.incr_steal_count(n as u16);
@@ -569,6 +598,7 @@ impl<T> Drop for Local<T> {
     fn drop(&mut self) {
         if !std::thread::panicking() {
             assert!(self.pop().is_none(), "queue not empty");
+            assert!(self.pop_lifo().is_none(), "LIFO slot not empty");
         }
     }
 }

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -112,13 +112,6 @@ 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<Notified>,
-
     /// When `true`, locally scheduled tasks go to the LIFO slot. When `false`,
     /// they go to the back of the `run_queue`.
     lifo_enabled: bool,
@@ -280,7 +273,6 @@ pub(super) fn create(
 
         cores.push(Box::new(Core {
             tick: 0,
-            lifo_slot: None,
             lifo_enabled: !config.disable_lifo_slot,
             run_queue,
             #[cfg(all(tokio_unstable, feature = "time"))]
@@ -440,7 +432,7 @@ where
         // If we heavily call `spawn_blocking`, there might be no available thread to
         // run this core. Except for the task in the lifo_slot, all tasks can be
         // stolen, so we move the task out of the lifo_slot to the run_queue.
-        if let Some(task) = core.lifo_slot.take() {
+        if let Some(task) = core.run_queue.pop_lifo() {
             core.run_queue
                 .push_back_or_overflow(task, &*cx.worker.handle, &mut core.stats);
         }
@@ -670,7 +662,7 @@ impl Context {
                 };
 
                 // Check for a task in the LIFO slot
-                let task = match core.lifo_slot.take() {
+                let task = match core.run_queue.pop_lifo() {
                     Some(task) => task,
                     None => {
                         self.reset_lifo_enabled(&mut core);
@@ -1079,7 +1071,7 @@ impl Core {
     }
 
     fn next_local_task(&mut self) -> Option<Notified> {
-        self.lifo_slot.take().or_else(|| self.run_queue.pop())
+        self.run_queue.pop_lifo().or_else(|| self.run_queue.pop())
     }
 
     /// Function responsible for stealing tasks from another worker
@@ -1135,7 +1127,7 @@ impl Core {
     }
 
     fn has_tasks(&self) -> bool {
-        self.lifo_slot.is_some() || self.run_queue.has_tasks()
+        self.run_queue.has_tasks()
     }
 
     fn should_notify_others(&self) -> bool {
@@ -1144,7 +1136,7 @@ impl Core {
         if self.is_searching {
             return false;
         }
-        self.lifo_slot.is_some() as usize + self.run_queue.len() > 1
+        self.run_queue.len() > 1
     }
 
     /// Prepares the worker state for parking.
@@ -1306,29 +1298,23 @@ impl Handle {
         // 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 {
+        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 {
+            if let Some(prev) = core.run_queue.push_lifo(task) {
+                // There was a previous task in the LIFO slot which needs
+                // to be pushed to the back of the run queue.
                 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() {
+        if core.park.is_some() {
             self.notify_parked_local();
         }
     }

tokio/src/runtime/task/atomic_notified.rs

@@ -0,0 +1,58 @@
+use crate::loom::sync::atomic::AtomicPtr;
+use crate::runtime::task::{Header, Notified, RawTask};
+
+use std::marker::PhantomData;
+use std::ptr;
+use std::ptr::NonNull;
+use std::sync::atomic::Ordering::SeqCst;
+
+/// An atomic cell which can contain a pointer to a [`Notified`] task.
+///
+/// This is similar to the `crate::util::AtomicCell` type, but specialized to
+/// hold a task pointer --- this type "remembers" the task's scheduler generic
+/// when a task is stored in the cell, so that the pointer can be turned back
+/// into a [`Notified`] task with the correct generic type when it is retrieved.
+pub(crate) struct AtomicNotified<S: 'static> {
+    task: AtomicPtr<Header>,
+    _scheduler: PhantomData<S>,
+}
+
+impl<S: 'static> AtomicNotified<S> {
+    pub(crate) fn empty() -> Self {
+        Self {
+            task: AtomicPtr::new(ptr::null_mut()),
+            _scheduler: PhantomData,
+        }
+    }
+
+    pub(crate) fn swap(&self, task: Option<Notified<S>>) -> Option<Notified<S>> {
+        let new = task
+            .map(|t| t.into_raw().header_ptr().as_ptr())
+            .unwrap_or_else(ptr::null_mut);
+        let old = self.task.swap(new, SeqCst);
+        NonNull::new(old).map(|ptr| unsafe {
+            // Safety: since we only allow tasks with the same scheduler type to
+            // be placed in this cell, we know that the pointed task's scheduler
+            // type matches the type parameter S.
+            Notified::from_raw(RawTask::from_raw(ptr))
+        })
+    }
+
+    pub(crate) fn take(&self) -> Option<Notified<S>> {
+        self.swap(None)
+    }
+
+    pub(crate) fn is_some(&self) -> bool {
+        !self.task.load(SeqCst).is_null()
+    }
+}
+
+unsafe impl<S: Send> Send for AtomicNotified<S> {}
+unsafe impl<S: Send> Sync for AtomicNotified<S> {}
+
+impl<S> Drop for AtomicNotified<S> {
+    fn drop(&mut self) {
+        // Ensure the task reference is dropped if this cell is dropped.
+        let _ = self.take();
+    }
+}

tokio/src/runtime/task/mod.rs

@@ -209,6 +209,11 @@ pub(crate) use self::raw::RawTask;
 mod state;
 use self::state::State;
 
+#[cfg(feature = "rt-multi-thread")]
+mod atomic_notified;
+#[cfg(feature = "rt-multi-thread")]
+pub(crate) use self::atomic_notified::AtomicNotified;
+
 mod waker;
 
 pub(crate) use self::spawn_location::SpawnLocation;