runtime: steal tasks from the LIFO slot

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,
@@ -429,23 +452,37 @@ impl<T> Steal<T> {
         // tasks in `dst`.
         let mut n = self.steal_into2(dst, dst_tail);
 
-        if n == 0 {
+        // Ooh, there's another task just sitting there? Grab that one, too!
+        let lifo = self.0.lifo.take();
+
+        if n == 0 && lifo.is_none() {
             // No tasks were stolen
             return None;
         }
 
-        dst_stats.incr_steal_count(n as u16);
+        // If we also grabbed the task from the LIFO slot, include that in the
+        // steal count as well.
+        dst_stats.incr_steal_count(n as u16 + lifo.is_some() as u16);
         dst_stats.incr_steal_operations();
 
-        // We are returning a task here
-        n -= 1;
+        let ret = if let Some(lifo) = lifo {
+            // If we took the task from the LIFO slot, just return it as the
+            // next task to run, rather than messing around with tasks from the
+            // queue.
+            lifo
+        } else {
+            // We are returning a task from the queue here.
+            n -= 1;
 
-        let ret_pos = dst_tail.wrapping_add(n);
-        let ret_idx = ret_pos as usize & MASK;
+            // The LIFO slot was was empty, so take the last task we squirted
+            // into `dst` instead.
+            let ret_pos = dst_tail.wrapping_add(n);
+            let ret_idx = ret_pos as usize & MASK;
 
-        // safety: the value was written as part of `steal_into2` and not
-        // exposed to stealers, so no other thread can access it.
-        let ret = dst.inner.buffer[ret_idx].with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) });
+            // safety: the value was written as part of `steal_into2` and not
+            // exposed to stealers, so no other thread can access it.
+            dst.inner.buffer[ret_idx].with(|ptr| unsafe { ptr::read((*ptr).as_ptr()) })
+        };
 
         if n == 0 {
             // The `dst` queue is empty, but a single task was stolen

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

@@ -101,13 +101,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,
@@ -257,7 +250,6 @@ pub(super) fn create(
 
         cores.push(Box::new(Core {
             tick: 0,
-            lifo_slot: None,
             lifo_enabled: !config.disable_lifo_slot,
             run_queue,
             is_searching: false,
@@ -405,7 +397,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);
         }
@@ -620,7 +612,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);
@@ -858,7 +850,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
@@ -914,7 +906,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 {
@@ -923,7 +915,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.
@@ -1084,29 +1076,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::{AcqRel, Acquire};
+
+/// 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, AcqRel);
+        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(Acquire).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

@@ -214,6 +214,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;