rt: move Inject to runtime::scheduler (#5748)

Previously, `Inject` was defined in `runtime::task`. This was because it
used some internal fns as part of the intrusive linked-list
implementation.

In the future, we want to remove the mutex from Inject and move it to
the scheduler proper (to reduce mutex ops). To set this up, this commit
moves `Inject` to `runtime::scheduler`. To make this work, we have to
`pub(crate)` `task::RawTask` and use that as the interface to access the
next / previous pointers.

tokio/src/runtime/scheduler/current_thread.rs

@@ -2,8 +2,8 @@ use crate::future::poll_fn;
 use crate::loom::sync::atomic::AtomicBool;
 use crate::loom::sync::Arc;
 use crate::runtime::driver::{self, Driver};
-use crate::runtime::scheduler::{self, Defer};
-use crate::runtime::task::{self, Inject, JoinHandle, OwnedTasks, Schedule, Task};
+use crate::runtime::scheduler::{self, Defer, Inject};
+use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task};
 use crate::runtime::{blocking, context, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics};
 use crate::sync::notify::Notify;
 use crate::util::atomic_cell::AtomicCell;

tokio/src/runtime/scheduler/inject.rs

@@ -5,7 +5,6 @@ use crate::loom::sync::{Mutex, MutexGuard};
 use crate::runtime::task;
 
 use std::marker::PhantomData;
-use std::ptr::NonNull;
 use std::sync::atomic::Ordering::{Acquire, Release};
 
 /// Growable, MPMC queue used to inject new tasks into the scheduler and as an
@@ -26,10 +25,10 @@ struct Pointers {
     is_closed: bool,
 
     /// Linked-list head.
-    head: Option<NonNull<task::Header>>,
+    head: Option<task::RawTask>,
 
     /// Linked-list tail.
-    tail: Option<NonNull<task::Header>>,
+    tail: Option<task::RawTask>,
 }
 
 pub(crate) struct Pop<'a, T: 'static> {
@@ -190,8 +189,8 @@ cfg_rt_multi_thread! {
         #[inline]
         fn push_batch_inner(
             &self,
-            batch_head: NonNull<task::Header>,
-            batch_tail: NonNull<task::Header>,
+            batch_head: task::RawTask,
+            batch_tail: task::RawTask,
             num: usize,
         ) {
             debug_assert!(get_next(batch_tail).is_none());
@@ -282,12 +281,12 @@ impl Pointers {
     }
 }
 
-fn get_next(header: NonNull<task::Header>) -> Option<NonNull<task::Header>> {
-    unsafe { header.as_ref().queue_next.with(|ptr| *ptr) }
+fn get_next(task: task::RawTask) -> Option<task::RawTask> {
+    unsafe { task.get_queue_next() }
 }
 
-fn set_next(header: NonNull<task::Header>, val: Option<NonNull<task::Header>>) {
+fn set_next(task: task::RawTask, val: Option<task::RawTask>) {
     unsafe {
-        header.as_ref().set_next(val);
+        task.set_queue_next(val);
     }
 }

tokio/src/runtime/scheduler/mod.rs

@@ -4,6 +4,9 @@ cfg_rt! {
 
     mod defer;
     use defer::Defer;
+
+    mod inject;
+    pub(crate) use inject::Inject;
 }
 
 cfg_rt_multi_thread! {

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

@@ -3,7 +3,8 @@
 use crate::loom::cell::UnsafeCell;
 use crate::loom::sync::Arc;
 use crate::runtime::scheduler::multi_thread::Stats;
-use crate::runtime::task::{self, Inject};
+use crate::runtime::scheduler::Inject;
+use crate::runtime::task;
 
 use std::mem::{self, MaybeUninit};
 use std::ptr;

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

@@ -62,8 +62,8 @@ use crate::runtime::context;
 use crate::runtime::scheduler::multi_thread::{
     idle, queue, Counters, Handle, Idle, Parker, Stats, Unparker,
 };
-use crate::runtime::scheduler::Defer;
-use crate::runtime::task::{Inject, OwnedTasks};
+use crate::runtime::scheduler::{Defer, Inject};
+use crate::runtime::task::OwnedTasks;
 use crate::runtime::{
     blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics,
 };

tokio/src/runtime/task/mod.rs

@@ -182,9 +182,6 @@ mod id;
 #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
 pub use id::{id, try_id, Id};
 
-mod inject;
-pub(super) use self::inject::Inject;
-
 #[cfg(feature = "rt")]
 mod abort;
 mod join;
@@ -198,7 +195,7 @@ mod list;
 pub(crate) use self::list::{LocalOwnedTasks, OwnedTasks};
 
 mod raw;
-use self::raw::RawTask;
+pub(crate) use self::raw::RawTask;
 
 mod state;
 use self::state::State;
@@ -335,13 +332,17 @@ cfg_rt! {
 }
 
 impl<S: 'static> Task<S> {
-    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
+    unsafe fn new(raw: RawTask) -> Task<S> {
         Task {
-            raw: RawTask::from_raw(ptr),
+            raw,
             _p: PhantomData,
         }
     }
 
+    unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
+        Task::new(RawTask::from_raw(ptr))
+    }
+
     #[cfg(all(
         tokio_unstable,
         tokio_taskdump,
@@ -369,22 +370,16 @@ impl<S: 'static> Notified<S> {
 }
 
 impl<S: 'static> Notified<S> {
-    unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> {
-        Notified(Task::from_raw(ptr))
-    }
-}
-
-impl<S: 'static> Task<S> {
-    fn into_raw(self) -> NonNull<Header> {
-        let ret = self.raw.header_ptr();
-        mem::forget(self);
-        ret
+    pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> {
+        Notified(Task::new(ptr))
     }
 }
 
 impl<S: 'static> Notified<S> {
-    fn into_raw(self) -> NonNull<Header> {
-        self.0.into_raw()
+    pub(crate) fn into_raw(self) -> RawTask {
+        let raw = self.0.raw;
+        mem::forget(self);
+        raw
     }
 }
 

tokio/src/runtime/task/raw.rs

@@ -6,7 +6,7 @@ use std::ptr::NonNull;
 use std::task::{Poll, Waker};
 
 /// Raw task handle
-pub(in crate::runtime) struct RawTask {
+pub(crate) struct RawTask {
     ptr: NonNull<Header>,
 }
 
@@ -240,6 +240,27 @@ impl RawTask {
     pub(super) fn ref_inc(self) {
         self.header().state.ref_inc();
     }
+
+    /// Get the queue-next pointer
+    ///
+    /// This is for usage by the injection queue
+    ///
+    /// Safety: make sure only one queue uses this and access is synchronized.
+    pub(crate) unsafe fn get_queue_next(self) -> Option<RawTask> {
+        self.header()
+            .queue_next
+            .with(|ptr| *ptr)
+            .map(|p| RawTask::from_raw(p))
+    }
+
+    /// Sets the queue-next pointer
+    ///
+    /// This is for usage by the injection queue
+    ///
+    /// Safety: make sure only one queue uses this and access is synchronized.
+    pub(crate) unsafe fn set_queue_next(self, val: Option<RawTask>) {
+        self.header().set_next(val.map(|task| task.ptr));
+    }
 }
 
 impl Clone for RawTask {

tokio/src/runtime/task/trace/mod.rs

@@ -1,7 +1,6 @@
 use crate::loom::sync::Arc;
 use crate::runtime::context;
-use crate::runtime::scheduler::{self, current_thread};
-use crate::runtime::task::Inject;
+use crate::runtime::scheduler::{self, current_thread, Inject};
 
 use backtrace::BacktraceFrame;
 use std::cell::Cell;

tokio/src/runtime/tests/inject.rs

@@ -1,4 +1,4 @@
-use crate::runtime::task::Inject;
+use crate::runtime::scheduler::Inject;
 
 #[test]
 fn push_and_pop() {

tokio/src/runtime/tests/loom_queue.rs

@@ -1,5 +1,5 @@
 use crate::runtime::scheduler::multi_thread::{queue, Stats};
-use crate::runtime::task::Inject;
+use crate::runtime::scheduler::Inject;
 use crate::runtime::tests::NoopSchedule;
 
 use loom::thread;

tokio/src/runtime/tests/queue.rs

@@ -1,5 +1,6 @@
 use crate::runtime::scheduler::multi_thread::{queue, Stats};
-use crate::runtime::task::{self, Inject, Schedule, Task};
+use crate::runtime::scheduler::Inject;
+use crate::runtime::task::{self, Schedule, Task};
 
 use std::thread;
 use std::time::Duration;