diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 206cdf2695a..70415799225 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/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::(&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(); } }