wip

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

@@ -84,19 +84,17 @@ pub(super) struct Worker {
 
 /// Core data
 struct Core {
+    /// Index holding this worker's remote state
+    index: usize,
+
     /// 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 message passing patterns and
-    /// helps to reduce latency.
-    lifo_slot: Option<Notified>,
-
     /// The worker-local run queue.
     run_queue: queue::Local<Arc<Handle>>,
 
+    lifo: Lifo,
+
     /// True if the worker is currently searching for more work. Searching
     /// involves attempting to steal from other workers.
     is_searching: bool,
@@ -117,6 +115,9 @@ struct Core {
     rand: FastRand,
 }
 
+unsafe impl Send for Core {}
+unsafe impl Sync for Core {}
+
 /// State shared across all workers
 pub(super) struct Shared {
     /// Per-worker remote state. All other workers have access to this and is
@@ -152,6 +153,13 @@ pub(super) struct Shared {
 
 /// 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 message passing patterns and
+    /// helps to reduce latency.
+    lifo_slot: Lifo,
+
     /// Steals tasks from this worker.
     steal: queue::Steal<Arc<Handle>>,
 
@@ -182,6 +190,8 @@ type Task = task::Task<Arc<Handle>>;
 /// A notified task handle
 type Notified = task::Notified<Arc<Handle>>;
 
+type Lifo = task::AtomicCell<Arc<Handle>>;
+
 // Tracks thread-local state
 scoped_thread_local!(static CURRENT: Context);
 
@@ -198,17 +208,24 @@ pub(super) fn create(
     let mut worker_metrics = Vec::with_capacity(size);
 
     // Create the local queues
-    for _ in 0..size {
+    for i in 0..size {
         let (steal, run_queue) = queue::local();
 
         let park = park.clone();
         let unpark = park.unpark();
         let metrics = WorkerMetrics::from_config(&config);
 
+        remotes.push(Remote {
+            lifo_slot: task::AtomicCell::new(),
+            steal,
+            unpark,
+        });
+
         cores.push(Box::new(Core {
+            index: i,
             tick: 0,
-            lifo_slot: None,
             run_queue,
+            lifo: task::AtomicCell::new(),
             is_searching: false,
             is_shutdown: false,
             park: Some(park),
@@ -216,7 +233,6 @@ pub(super) fn create(
             rand: FastRand::new(config.seed_generator.next_seed()),
         }));
 
-        remotes.push(Remote { steal, unpark });
         worker_metrics.push(metrics);
     }
 
@@ -482,7 +498,7 @@ impl Context {
                 core.metrics.end_poll();
 
                 // Check for a task in the LIFO slot
-                let task = match core.lifo_slot.take() {
+                let task = match self.worker.lifo_slot().take_local() {
                     Some(task) => task,
                     None => return Ok(core),
                 };
@@ -601,21 +617,22 @@ impl Core {
     /// Return the next notified task available to this worker.
     fn next_task(&mut self, worker: &Worker) -> Option<Notified> {
         if self.tick % worker.handle.shared.config.global_queue_interval == 0 {
-            worker.inject().pop().or_else(|| self.next_local_task())
+            worker
+                .inject()
+                .pop()
+                .or_else(|| self.next_local_task(&worker.handle))
         } else {
-            self.next_local_task().or_else(|| worker.inject().pop())
+            self.next_local_task(&worker.handle)
+                .or_else(|| worker.inject().pop())
         }
     }
 
-    fn next_local_task(&mut self) -> Option<Notified> {
-        self.lifo_slot.take().or_else(|| self.run_queue.pop())
+    fn next_local_task(&mut self, handle: &Handle) -> Option<Notified> {
+        self.lifo_slot(handle)
+            .take_local()
+            .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<Notified> {
         if !self.transition_to_searching(worker) {
             return None;
@@ -668,7 +685,7 @@ impl Core {
     /// 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.lifo_slot.is_some() || self.run_queue.has_tasks() {
+        if worker.lifo_slot().is_some() || self.run_queue.has_tasks() {
             return false;
         }
 
@@ -696,7 +713,7 @@ impl Core {
     fn transition_from_parked(&mut self, worker: &Worker) -> bool {
         // If a task is in the lifo slot, then we must unpark regardless of
         // being notified
-        if self.lifo_slot.is_some() {
+        if worker.lifo_slot().is_some() {
             // 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"
@@ -741,10 +758,15 @@ impl Core {
         let mut park = self.park.take().expect("park missing");
 
         // Drain the queue
-        while self.next_local_task().is_some() {}
+        while self.next_local_task(handle).is_some() {}
 
         park.shutdown(&handle.driver);
     }
+
+    fn lifo_slot<'a>(&'a self, handle: &'a Handle) -> &'a Lifo {
+        // &handle.shared.remotes[self.index].lifo_slot
+        &self.lifo
+    }
 }
 
 impl Worker {
@@ -752,6 +774,11 @@ impl Worker {
     fn inject(&self) -> &Inject<Arc<Handle>> {
         &self.handle.shared.inject
     }
+
+    /// Returns a reference to the worker's lifo slot
+    fn lifo_slot(&self) -> &Lifo {
+        &self.handle.shared.remotes[self.index].lifo_slot
+    }
 }
 
 // TODO: Move `Handle` impls into handle.rs
@@ -797,29 +824,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 || self.shared.config.disable_lifo_slot {
+        if is_yield || self.shared.config.disable_lifo_slot {
             core.run_queue
                 .push_back(task, &self.shared.inject, &mut core.metrics);
-            true
         } else {
             // Push to the LIFO slot
-            let prev = core.lifo_slot.take();
-            let ret = prev.is_some();
+            let prev = core.lifo_slot(&self).swap_local(task);
 
             if let Some(prev) = prev {
                 core.run_queue
                     .push_back(prev, &self.shared.inject, &mut core.metrics);
             }
-
-            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();
         }
     }

tokio/src/runtime/task/atomic_cell.rs

@@ -0,0 +1,104 @@
+use crate::loom::sync::atomic::AtomicPtr;
+use crate::loom::sync::atomic::Ordering::{AcqRel, Acquire, Release};
+use crate::runtime::task::{Header, Notified};
+
+use std::marker::PhantomData;
+use std::ptr::{self, NonNull};
+
+pub(crate) struct AtomicCell<S: 'static> {
+    // task: AtomicPtr<Header>,
+    task: std::cell::UnsafeCell<Option<Notified<S>>>,
+    _p: PhantomData<S>,
+}
+
+unsafe impl<S: 'static> Send for AtomicCell<S> {}
+unsafe impl<S: 'static> Sync for AtomicCell<S> {}
+
+impl<S> AtomicCell<S> {
+    pub(crate) fn new() -> AtomicCell<S> {
+        AtomicCell {
+            // task: AtomicPtr::default(),
+            task: Default::default(),
+            _p: PhantomData,
+        }
+    }
+
+    /// Should be called from a local context
+    pub(crate) fn is_some(&self) -> bool {
+        // !self.task.load(Acquire).is_null()
+        unsafe { (*self.task.get()).is_some() }
+    }
+
+    pub(crate) fn take_local(&self) -> Option<Notified<S>> {
+        unsafe { (*self.task.get()).take() }
+        /*
+        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(ptr) })
+        */
+    }
+
+    pub(crate) fn swap_local(&self, task: Notified<S>) -> Option<Notified<S>> {
+        std::mem::replace(unsafe { &mut (*self.task.get()) }, Some(task))
+        /*
+        let next = task.into_raw().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, Release, Acquire)
+            .is_ok()
+        {
+            // Safety: we already checked !null above
+            let prev = unsafe { Notified::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<Notified<S>> {
+        let task = self.task.load(Acquire);
+
+        if task.is_null() {
+            return None;
+        }
+
+        // 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(NonNull::new_unchecked(task)) });
+        }
+
+        return None;
+    }
+    */
+}

tokio/src/runtime/task/mod.rs

@@ -183,6 +183,9 @@ mod id;
 pub use id::{id, try_id, Id};
 
 cfg_rt_multi_thread! {
+    mod atomic_cell;
+    pub(super) use atomic_cell::AtomicCell;
+
     mod inject;
     pub(super) use self::inject::Inject;
 }