task: elaborate safety comments in task deallocation (#5172)

tokio/src/runtime/task/harness.rs

@@ -1,6 +1,6 @@
 use crate::future::Future;
 use crate::runtime::task::core::{Cell, Core, CoreStage, Header, Trailer};
-use crate::runtime::task::state::Snapshot;
+use crate::runtime::task::state::{Snapshot, State};
 use crate::runtime::task::waker::waker_ref;
 use crate::runtime::task::{JoinError, Notified, Schedule, Task};
 
@@ -31,7 +31,11 @@ where
     }
 
     fn header(&self) -> &Header {
-        unsafe { &self.cell.as_ref().header }
+        unsafe { &*self.header_ptr().as_ptr() }
+    }
+
+    fn state(&self) -> &State {
+        &self.header().state
     }
 
     fn trailer(&self) -> &Trailer {
@@ -95,7 +99,7 @@ where
     fn poll_inner(&self) -> PollFuture {
         use super::state::{TransitionToIdle, TransitionToRunning};
 
-        match self.header().state.transition_to_running() {
+        match self.state().transition_to_running() {
             TransitionToRunning::Success => {
                 let header_ptr = self.header_ptr();
                 let waker_ref = waker_ref::<T, S>(&header_ptr);
@@ -113,7 +117,7 @@ where
                     return PollFuture::Complete;
                 }
 
-                match self.header().state.transition_to_idle() {
+                match self.state().transition_to_idle() {
                     TransitionToIdle::Ok => PollFuture::Done,
                     TransitionToIdle::OkNotified => PollFuture::Notified,
                     TransitionToIdle::OkDealloc => PollFuture::Dealloc,
@@ -143,7 +147,7 @@ where
     /// there is nothing further to do. When the task completes running, it will
     /// notice the `CANCELLED` bit and finalize the task.
     pub(super) fn shutdown(self) {
-        if !self.header().state.transition_to_shutdown() {
+        if !self.state().transition_to_shutdown() {
             // The task is concurrently running. No further work needed.
             self.drop_reference();
             return;
@@ -163,6 +167,19 @@ where
         // Check causality
         self.core().stage.with_mut(drop);
 
+        // Safety: The caller of this method just transitioned our ref-count to
+        // zero, so it is our responsibility to release the allocation.
+        //
+        // We don't hold any references into the allocation at this point, but
+        // it is possible for another thread to still hold a `&State` into the
+        // allocation if that other thread has decremented its last ref-count,
+        // but has not yet returned from the relevant method on `State`.
+        //
+        // However, the `State` type consists of just an `AtomicUsize`, and an
+        // `AtomicUsize` wraps the entirety of its contents in an `UnsafeCell`.
+        // As explained in the documentation for `UnsafeCell`, such references
+        // are allowed to be dangling after their last use, even if the
+        // reference has not yet gone out of scope.
         unsafe {
             drop(Box::from_raw(self.cell.as_ptr()));
         }
@@ -187,7 +204,7 @@ where
     pub(super) fn drop_join_handle_slow(self) {
         // Try to unset `JOIN_INTEREST`. This must be done as a first step in
         // case the task concurrently completed.
-        if self.header().state.unset_join_interested().is_err() {
+        if self.state().unset_join_interested().is_err() {
             // It is our responsibility to drop the output. This is critical as
             // the task output may not be `Send` and as such must remain with
             // the scheduler or `JoinHandle`. i.e. if the output remains in the
@@ -214,7 +231,7 @@ where
     /// the shutdown. This is necessary to avoid the shutdown happening in the
     /// wrong thread for non-Send tasks.
     pub(super) fn remote_abort(self) {
-        if self.header().state.transition_to_notified_and_cancel() {
+        if self.state().transition_to_notified_and_cancel() {
             // The transition has created a new ref-count, which we turn into
             // a Notified and pass to the task.
             //
@@ -237,7 +254,7 @@ where
     pub(super) fn wake_by_val(self) {
         use super::state::TransitionToNotifiedByVal;
 
-        match self.header().state.transition_to_notified_by_val() {
+        match self.state().transition_to_notified_by_val() {
             TransitionToNotifiedByVal::Submit => {
                 // The caller has given us a ref-count, and the transition has
                 // created a new ref-count, so we now hold two. We turn the new
@@ -267,7 +284,7 @@ where
     pub(super) fn wake_by_ref(&self) {
         use super::state::TransitionToNotifiedByRef;
 
-        match self.header().state.transition_to_notified_by_ref() {
+        match self.state().transition_to_notified_by_ref() {
             TransitionToNotifiedByRef::Submit => {
                 // The transition above incremented the ref-count for a new task
                 // and the caller also holds a ref-count. The caller's ref-count
@@ -282,7 +299,7 @@ where
     }
 
     pub(super) fn drop_reference(self) {
-        if self.header().state.ref_dec() {
+        if self.state().ref_dec() {
             self.dealloc();
         }
     }
@@ -299,7 +316,7 @@ where
         // The future has completed and its output has been written to the task
         // stage. We transition from running to complete.
 
-        let snapshot = self.header().state.transition_to_complete();
+        let snapshot = self.state().transition_to_complete();
 
         // We catch panics here in case dropping the future or waking the
         // JoinHandle panics.
@@ -319,7 +336,7 @@ where
         // The task has completed execution and will no longer be scheduled.
         let num_release = self.release();
 
-        if self.header().state.transition_to_terminal(num_release) {
+        if self.state().transition_to_terminal(num_release) {
             self.dealloc();
         }
     }