From f6aa231f77fb4d2d23af6bda867b64ee2bf4cfa9 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Sun, 6 Nov 2022 00:50:57 +0100 Subject: [PATCH 1/4] task: elaborate safety comments in task deallocation --- tokio/src/runtime/task/harness.rs | 39 ++++++++++++++++++++++--------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 206cdf2695a..708443552b8 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::{State, Snapshot}; use crate::runtime::task::waker::waker_ref; use crate::runtime::task::{JoinError, Notified, Schedule, Task}; @@ -34,6 +34,10 @@ where unsafe { &self.cell.as_ref().header } } + fn state(&self) -> &State { + &self.header().state + } + fn trailer(&self) -> &Trailer { unsafe { &self.cell.as_ref().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(); } } From 8e88f0b767fde61bca05ab8a0a65c2d80138f86c Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Sun, 6 Nov 2022 13:28:11 +0100 Subject: [PATCH 2/4] rustfmt --- tokio/src/runtime/task/harness.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 708443552b8..0182a467d9a 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::{State, Snapshot}; +use crate::runtime::task::state::{Snapshot, State}; use crate::runtime::task::waker::waker_ref; use crate::runtime::task::{JoinError, Notified, Schedule, Task}; From 0ec8d519c7227f26bac9676fa36325e35b527600 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Sun, 6 Nov 2022 14:05:13 +0100 Subject: [PATCH 3/4] Simplify header fn --- tokio/src/runtime/task/harness.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 0182a467d9a..680e413f01b 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -31,7 +31,7 @@ where } fn header(&self) -> &Header { - unsafe { &self.cell.as_ref().header } + unsafe { self.header_ptr().as_ref() } } fn state(&self) -> &State { From 28b7312c7f8847dee8e870664f31e508ed4cd494 Mon Sep 17 00:00:00 2001 From: Alice Ryhl Date: Sun, 6 Nov 2022 14:44:41 +0100 Subject: [PATCH 4/4] Fix --- tokio/src/runtime/task/harness.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 680e413f01b..70415799225 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -31,7 +31,7 @@ where } fn header(&self) -> &Header { - unsafe { self.header_ptr().as_ref() } + unsafe { &*self.header_ptr().as_ptr() } } fn state(&self) -> &State {