fix alice's findings

tokio/src/runtime/context.rs

@@ -20,6 +20,8 @@ cfg_rt! {
     use scoped::Scoped;
 
     use crate::runtime::{scheduler, task::Id};
+    #[cfg(tokio_unstable)]
+    use crate::runtime::task::Header;
 
     #[cfg(tokio_unstable)]
     use std::ptr::NonNull;
@@ -52,7 +54,7 @@ struct Context {
     current_task_id: Cell<Option<Id>>,
 
     #[cfg(all(feature = "rt", tokio_unstable))]
-    current_task: Cell<Option<NonNull<()>>>,
+    current_task: Cell<Option<NonNull<Header>>>,
 
     /// Tracks if the current thread is currently driving a runtime.
     /// Note, that if this is set to "entered", the current scheduler
@@ -168,14 +170,28 @@ cfg_rt! {
     }
 
     #[cfg(tokio_unstable)]
-    pub(crate) fn set_current_task(task: Option<NonNull<()>>) -> Option<NonNull<()>> {
+    pub(crate) fn set_current_task(task: Option<NonNull<Header>>) -> Option<NonNull<Header>> {
         CONTEXT
             .try_with(|ctx| ctx.current_task.replace(task))
             .unwrap_or(None)
     }
 
     #[cfg(tokio_unstable)]
-    pub(crate) fn current_task() -> Option<NonNull<()>> {
+    pub(crate) fn set_current_task_id_and_task(
+        id: Option<Id>,
+        task: Option<NonNull<Header>>,
+    ) -> (Option<Id>, Option<NonNull<Header>>) {
+        CONTEXT
+            .try_with(|ctx| {
+                let parent_task_id = ctx.current_task_id.replace(id);
+                let parent_task = ctx.current_task.replace(task);
+                (parent_task_id, parent_task)
+            })
+            .unwrap_or((None, None))
+    }
+
+    #[cfg(tokio_unstable)]
+    pub(crate) fn current_task() -> Option<NonNull<Header>> {
         CONTEXT
             .try_with(|ctx| ctx.current_task.get())
             .unwrap_or(None)

tokio/src/runtime/task/core.rs

@@ -350,6 +350,32 @@ impl Drop for TaskIdGuard {
     }
 }
 
+#[cfg(tokio_unstable)]
+struct TaskContextGuard {
+    parent_task_id: Option<Id>,
+    parent_task: Option<NonNull<Header>>,
+}
+
+#[cfg(tokio_unstable)]
+impl TaskContextGuard {
+    fn enter(id: Id, header: NonNull<Header>) -> Self {
+        let (parent_task_id, parent_task) =
+            context::set_current_task_id_and_task(Some(id), Some(header));
+
+        TaskContextGuard {
+            parent_task_id,
+            parent_task,
+        }
+    }
+}
+
+#[cfg(tokio_unstable)]
+impl Drop for TaskContextGuard {
+    fn drop(&mut self) {
+        context::set_current_task_id_and_task(self.parent_task_id, self.parent_task);
+    }
+}
+
 impl<T: Future, S: Schedule> Core<T, S> {
     /// Polls the future.
     ///
@@ -365,12 +391,11 @@ impl<T: Future, S: Schedule> Core<T, S> {
     /// `self` must also be pinned. This is handled by storing the task on the
     /// heap.
     ///
-    /// When `tokio_unstable` is enabled, `header` must point to the header for
+    /// `header` must point to the header for this exact task allocation, and
-    /// this exact task allocation, and the allocation must remain live until
+    /// the allocation must remain live until this function returns.
-    /// this function returns.
     pub(super) unsafe fn poll(
         &self,
-        #[cfg(tokio_unstable)] header: NonNull<Header>,
+        header: NonNull<Header>,
         mut cx: Context<'_>,
     ) -> Poll<T::Output> {
         let res = {
@@ -384,18 +409,16 @@ impl<T: Future, S: Schedule> Core<T, S> {
                 // Safety: The caller ensures the future is pinned.
                 let future = unsafe { Pin::new_unchecked(future) };
 
-                let _guard = TaskIdGuard::enter(self.task_id);
                 #[cfg(tokio_unstable)]
-                let _current_task = CurrentTaskGuard::enter(header);
+                let _guard = TaskContextGuard::enter(self.task_id, header);
+                #[cfg(not(tokio_unstable))]
+                let _guard = TaskIdGuard::enter(self.task_id);
                 future.poll(&mut cx)
             })
         };
 
         if res.is_ready() {
-            self.drop_future_or_output(
+            self.drop_future_or_output(header);
-                #[cfg(tokio_unstable)]
-                header,
-            );
         }
 
         res
@@ -406,7 +429,10 @@ impl<T: Future, S: Schedule> Core<T, S> {
     /// # Safety
     ///
     /// The caller must ensure it is safe to mutate the `stage` field.
-    pub(super) fn drop_future_or_output(&self, #[cfg(tokio_unstable)] header: NonNull<Header>) {
+    pub(super) fn drop_future_or_output(&self, header: NonNull<Header>) {
+        #[cfg(not(tokio_unstable))]
+        let _ = header;
+
         #[cfg(tokio_unstable)]
         let _current_task = {
             let dropping_future = self.stage.stage.with(|ptr| {
@@ -464,14 +490,14 @@ impl<T: Future, S: Schedule> Core<T, S> {
 
 #[cfg(tokio_unstable)]
 pub(crate) struct CurrentTaskGuard {
-    parent_task: Option<NonNull<()>>,
+    parent_task: Option<NonNull<Header>>,
 }
 
 #[cfg(tokio_unstable)]
 impl CurrentTaskGuard {
     fn enter(header: NonNull<Header>) -> Self {
         CurrentTaskGuard {
-            parent_task: context::set_current_task(Some(header.cast())),
+            parent_task: context::set_current_task(Some(header)),
         }
     }
 }

tokio/src/runtime/task/harness.rs

@@ -242,11 +242,7 @@ where
                     }
 
                     if self.state().load().is_cancelled() {
-                        cancel_task(
+                        cancel_task(self.core(), header_ptr);
-                            self.core(),
-                            #[cfg(tokio_unstable)]
-                            header_ptr,
-                        );
                         return PollFuture::Complete;
                     }
                 }
@@ -256,14 +252,7 @@ where
                 // Safety: `transition_to_running` succeeded, so this thread has
                 // exclusive access to the future/output storage. The header pointer
                 // comes from this harness and remains live while the task is running.
-                let res = unsafe {
+                let res = unsafe { poll_future(self.core(), header_ptr, cx) };
-                    poll_future(
-                        self.core(),
-                        #[cfg(tokio_unstable)]
-                        header_ptr,
-                        cx,
-                    )
-                };
 
                 #[cfg(tokio_unstable)]
                 {
@@ -294,20 +283,12 @@ where
                 if let TransitionToIdle::Cancelled = transition_res {
                     // The transition to idle failed because the task was
                     // cancelled during the poll.
-                    cancel_task(
+                    cancel_task(self.core(), header_ptr);
-                        self.core(),
-                        #[cfg(tokio_unstable)]
-                        header_ptr,
-                    );
                 }
                 transition_result_to_poll_future(transition_res)
             }
             TransitionToRunning::Cancelled => {
-                cancel_task(
+                cancel_task(self.core(), self.header_ptr());
-                    self.core(),
-                    #[cfg(tokio_unstable)]
-                    self.header_ptr(),
-                );
                 PollFuture::Complete
             }
             TransitionToRunning::Failed => PollFuture::Done,
@@ -330,11 +311,7 @@ where
 
         // By transitioning the lifecycle to `Running`, we have permission to
         // drop the future.
-        cancel_task(
+        cancel_task(self.core(), self.header_ptr());
-            self.core(),
-            #[cfg(tokio_unstable)]
-            self.header_ptr(),
-        );
         self.complete();
     }
 
@@ -390,10 +367,7 @@ where
             // they are dropping the `JoinHandle`, we assume they are not
             // interested in the panic and swallow it.
             let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
-                self.core().drop_future_or_output(
+                self.core().drop_future_or_output(self.header_ptr());
-                    #[cfg(tokio_unstable)]
-                    self.header_ptr(),
-                );
             }));
         }
 
@@ -434,10 +408,7 @@ where
                 // this task. It is our responsibility to drop the
                 // output. The join waker was already dropped by the
                 // `JoinHandle` before.
-                self.core().drop_future_or_output(
+                self.core().drop_future_or_output(self.header_ptr());
-                    #[cfg(tokio_unstable)]
-                    self.header_ptr(),
-                );
             } else if snapshot.is_join_waker_set() {
                 // Notify the waker. Reading the waker field is safe per rule 4
                 // in task/mod.rs, since the JOIN_WAKER bit is set and the call
@@ -592,16 +563,10 @@ enum PollFuture {
 }
 
 /// Cancels the task and store the appropriate error in the stage field.
-fn cancel_task<T: Future, S: Schedule>(
+fn cancel_task<T: Future, S: Schedule>(core: &Core<T, S>, header: NonNull<Header>) {
-    core: &Core<T, S>,
-    #[cfg(tokio_unstable)] header: NonNull<Header>,
-) {
     // Drop the future from a panic guard.
     let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
-        core.drop_future_or_output(
+        core.drop_future_or_output(header);
-            #[cfg(tokio_unstable)]
-            header,
-        );
     }));
 
     core.store_output(Err(panic_result_to_join_error(core.task_id, res)));
@@ -653,45 +618,33 @@ unsafe fn poll_hook_panic<T: Future, S: Schedule>(
 ///
 /// The caller must satisfy the mutual-exclusion requirements of `Core::poll`.
 ///
-/// When `tokio_unstable` is enabled, `header` must point to the header for this
+/// `header` must point to the header for this exact task allocation, and the
-/// exact task allocation, and the allocation must remain live until this
+/// allocation must remain live until this function returns.
-/// function returns.
 unsafe fn poll_future<T: Future, S: Schedule>(
     core: &Core<T, S>,
-    #[cfg(tokio_unstable)] header: NonNull<Header>,
+    header: NonNull<Header>,
     cx: Context<'_>,
 ) -> Poll<()> {
     // Poll the future.
     let output = panic::catch_unwind(panic::AssertUnwindSafe(|| {
         struct Guard<'a, T: Future, S: Schedule> {
             core: &'a Core<T, S>,
-            #[cfg(tokio_unstable)]
             header: NonNull<Header>,
         }
         impl<'a, T: Future, S: Schedule> Drop for Guard<'a, T, S> {
             fn drop(&mut self) {
                 // If the future panics on poll, we drop it inside the panic
                 // guard.
-                self.core.drop_future_or_output(
+                self.core.drop_future_or_output(self.header);
-                    #[cfg(tokio_unstable)]
-                    self.header,
-                );
             }
         }
         let guard = Guard {
             core,
-            #[cfg(tokio_unstable)]
             header,
         };
         // Safety: the caller guarantees the mutual-exclusion requirements of
         // `Core::poll` and that `header` identifies this live task allocation.
-        let res = unsafe {
+        let res = unsafe { guard.core.poll(header, cx) };
-            guard.core.poll(
-                #[cfg(tokio_unstable)]
-                header,
-                cx,
-            )
-        };
         mem::forget(guard);
         res
     }));

tokio/src/runtime/task/mod.rs

@@ -186,7 +186,7 @@
 
 mod core;
 use self::core::Cell;
-use self::core::Header;
+pub(crate) use self::core::Header;
 
 mod error;
 pub use self::error::JoinError;
@@ -687,7 +687,7 @@ pub(crate) fn with_current_task_meta<R>(
 
     // Safety: the context stores this pointer only while the referenced task is
     // being polled, so the allocation is alive for this synchronous call.
-    let raw = unsafe { RawTask::from_raw(ptr.cast()) };
+    let raw = unsafe { RawTask::from_raw(ptr) };
     // Safety: parent metadata is exposed read-only during this synchronous call
     // while no mutable parent hook metadata is live. The closure-bound lifetime
     // prevents references exposed through the metadata from escaping.