diff --git a/tokio/src/runtime/basic_scheduler.rs b/tokio/src/runtime/basic_scheduler.rs index cc91b1662..fe2e4a8c6 100644 --- a/tokio/src/runtime/basic_scheduler.rs +++ b/tokio/src/runtime/basic_scheduler.rs @@ -314,12 +314,10 @@ impl Drop for BasicScheduler

{ }; enter(&mut inner, |scheduler, context| { - // By closing the OwnedTasks, no new tasks can be spawned on it. - context.shared.owned.close(); - // Drain the OwnedTasks collection. - while let Some(task) = context.shared.owned.pop_back() { - task.shutdown(); - } + // Drain the OwnedTasks collection. This call also closes the + // collection, ensuring that no tasks are ever pushed after this + // call returns. + context.shared.owned.close_and_shutdown_all(); // Drain local queue // We already shut down every task, so we just need to drop the task. diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index 316680273..51b64966a 100644 --- a/tokio/src/runtime/task/core.rs +++ b/tokio/src/runtime/task/core.rs @@ -57,10 +57,10 @@ pub(crate) struct Header { /// Task state pub(super) state: State, - pub(crate) owned: UnsafeCell>, + pub(super) owned: UnsafeCell>, /// Pointer to next task, used with the injection queue - pub(crate) queue_next: UnsafeCell>>, + pub(super) queue_next: UnsafeCell>>, /// Table of function pointers for executing actions on the task. pub(super) vtable: &'static Vtable, @@ -239,18 +239,18 @@ impl Header { } impl Trailer { - pub(crate) unsafe fn set_waker(&self, waker: Option) { + pub(super) unsafe fn set_waker(&self, waker: Option) { self.waker.with_mut(|ptr| { *ptr = waker; }); } - pub(crate) unsafe fn will_wake(&self, waker: &Waker) -> bool { + pub(super) unsafe fn will_wake(&self, waker: &Waker) -> bool { self.waker .with(|ptr| (*ptr).as_ref().unwrap().will_wake(waker)) } - pub(crate) fn wake_join(&self) { + pub(super) fn wake_join(&self) { self.waker.with(|ptr| match unsafe { &*ptr } { Some(waker) => waker.wake_by_ref(), None => panic!("waker missing"), diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 5b0574a02..41b41938e 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -5,6 +5,7 @@ use crate::runtime::task::waker::waker_ref; use crate::runtime::task::{JoinError, Notified, Schedule, Task}; use std::mem; +use std::mem::ManuallyDrop; use std::panic; use std::ptr::NonNull; use std::task::{Context, Poll, Waker}; @@ -36,13 +37,6 @@ where fn core(&self) -> &Core { unsafe { &self.cell.as_ref().core } } - - fn scheduler_view(&self) -> SchedulerView<'_, S> { - SchedulerView { - header: self.header(), - scheduler: &self.core().scheduler, - } - } } impl Harness @@ -50,43 +44,103 @@ where T: Future, S: Schedule, { - /// Polls the inner future. + /// Polls the inner future. A ref-count is consumed. /// /// All necessary state checks and transitions are performed. - /// /// Panics raised while polling the future are handled. pub(super) fn poll(self) { + // We pass our ref-count to `poll_inner`. match self.poll_inner() { PollFuture::Notified => { - // Signal yield - self.core().scheduler.yield_now(Notified(self.to_task())); - // The ref-count was incremented as part of - // `transition_to_idle`. + // The `poll_inner` call has given us two ref-counts back. + // We give one of them to a new task and call `yield_now`. + self.core() + .scheduler + .yield_now(Notified(self.get_new_task())); + + // The remaining ref-count is now dropped. We kept the extra + // ref-count until now to ensure that even if the `yield_now` + // call drops the provided task, the task isn't deallocated + // before after `yield_now` returns. self.drop_reference(); } - PollFuture::DropReference => { - self.drop_reference(); + PollFuture::Complete => { + self.complete(); } - PollFuture::Complete(out, is_join_interested) => { - self.complete(out, is_join_interested); + PollFuture::Dealloc => { + self.dealloc(); } - PollFuture::None => (), + PollFuture::Done => (), } } - fn poll_inner(&self) -> PollFuture { - let snapshot = match self.scheduler_view().transition_to_running() { - TransitionToRunning::Ok(snapshot) => snapshot, - TransitionToRunning::DropReference => return PollFuture::DropReference, - }; + /// Poll the task and cancel it if necessary. This takes ownership of a + /// ref-count. + /// + /// If the return value is Notified, the caller is given ownership of two + /// ref-counts. + /// + /// If the return value is Complete, the caller is given ownership of a + /// single ref-count, which should be passed on to `complete`. + /// + /// If the return value is Dealloc, then this call consumed the last + /// ref-count and the caller should call `dealloc`. + /// + /// Otherwise the ref-count is consumed and the caller should not access + /// `self` again. + fn poll_inner(&self) -> PollFuture { + use super::state::{TransitionToIdle, TransitionToRunning}; - // The transition to `Running` done above ensures that a lock on the - // future has been obtained. This also ensures the `*mut T` pointer - // contains the future (as opposed to the output) and is initialized. + match self.header().state.transition_to_running() { + TransitionToRunning::Success => { + let waker_ref = waker_ref::(self.header()); + let cx = Context::from_waker(&*waker_ref); + let res = poll_future(&self.core().stage, cx); - let waker_ref = waker_ref::(self.header()); - let cx = Context::from_waker(&*waker_ref); - poll_future(self.header(), &self.core().stage, snapshot, cx) + if res == Poll::Ready(()) { + // The future completed. Move on to complete the task. + return PollFuture::Complete; + } + + match self.header().state.transition_to_idle() { + TransitionToIdle::Ok => PollFuture::Done, + TransitionToIdle::OkNotified => PollFuture::Notified, + TransitionToIdle::OkDealloc => PollFuture::Dealloc, + TransitionToIdle::Cancelled => { + // The transition to idle failed because the task was + // cancelled during the poll. + + cancel_task(&self.core().stage); + PollFuture::Complete + } + } + } + TransitionToRunning::Cancelled => { + cancel_task(&self.core().stage); + PollFuture::Complete + } + TransitionToRunning::Failed => PollFuture::Done, + TransitionToRunning::Dealloc => PollFuture::Dealloc, + } + } + + /// Forcibly shutdown the task + /// + /// Attempt to transition to `Running` in order to forcibly shutdown the + /// task. If the task is currently running or in a state of completion, then + /// 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() { + // The task is concurrently running. No further work needed. + self.drop_reference(); + return; + } + + // By transitioning the lifecycle to `Running`, we have permission to + // drop the future. + cancel_task(&self.core().stage); + self.complete(); } pub(super) fn dealloc(self) { @@ -124,6 +178,7 @@ where let panic = panic::catch_unwind(panic::AssertUnwindSafe(|| { self.core().stage.drop_future_or_output(); })); + if let Err(panic) = panic { maybe_panic = Some(panic); } @@ -137,16 +192,78 @@ where } } - // ===== waker behavior ===== - - pub(super) fn wake_by_val(self) { - self.wake_by_ref(); - self.drop_reference(); + /// Remotely abort the task. + /// + /// The caller should hold a ref-count, but we do not consume it. + /// + /// This is similar to `shutdown` except that it asks the runtime to perform + /// 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() { + // The transition has created a new ref-count, which we turn into + // a Notified and pass to the task. + // + // Since the caller holds a ref-count, the task cannot be destroyed + // before the call to `schedule` returns even if the call drops the + // `Notified` internally. + self.core() + .scheduler + .schedule(Notified(self.get_new_task())); + } } + // ===== waker behavior ===== + + /// This call consumes a ref-count and notifies the task. This will create a + /// new Notified and submit it if necessary. + /// + /// The caller does not need to hold a ref-count besides the one that was + /// passed to this call. + pub(super) fn wake_by_val(self) { + use super::state::TransitionToNotifiedByVal; + + match self.header().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 + // ref-count Notified and pass it to the call to `schedule`. + // + // The old ref-count is retained for now to ensure that the task + // is not dropped during the call to `schedule` if the call + // drops the task it was given. + self.core() + .scheduler + .schedule(Notified(self.get_new_task())); + + // Now that we have completed the call to schedule, we can + // release our ref-count. + self.drop_reference(); + } + TransitionToNotifiedByVal::Dealloc => { + self.dealloc(); + } + TransitionToNotifiedByVal::DoNothing => {} + } + } + + /// This call notifies the task. It will not consume any ref-counts, but the + /// caller should hold a ref-count. This will create a new Notified and + /// submit it if necessary. pub(super) fn wake_by_ref(&self) { - if self.header().state.transition_to_notified() { - self.core().scheduler.schedule(Notified(self.to_task())); + use super::state::TransitionToNotifiedByRef; + + match self.header().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 + // ensures that the task is not destroyed even if the new task + // is dropped before `schedule` returns. + self.core() + .scheduler + .schedule(Notified(self.get_new_task())); + } + TransitionToNotifiedByRef::DoNothing => {} } } @@ -161,153 +278,65 @@ where self.header().id.as_ref() } - /// Forcibly shutdown the task - /// - /// Attempt to transition to `Running` in order to forcibly shutdown the - /// task. If the task is currently running or in a state of completion, then - /// 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() { - // The task is concurrently running. No further work needed. - return; - } - - // By transitioning the lifecycle to `Running`, we have permission to - // drop the future. - let err = cancel_task(&self.core().stage); - self.complete(Err(err), true) - } - - /// Remotely abort the task - /// - /// This is similar to `shutdown` except that it asks the runtime to perform - /// 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() { - self.core().scheduler.schedule(Notified(self.to_task())); - } - } - // ====== internal ====== - fn complete(self, output: super::Result, is_join_interested: bool) { - // We catch panics here because dropping the output may panic. - // - // Dropping the output can also happen in the first branch inside - // transition_to_complete. - let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { - if is_join_interested { - // Store the output. The future has already been dropped - // - // Safety: Mutual exclusion is obtained by having transitioned the task - // state -> Running - let stage = &self.core().stage; - stage.store_output(output); + /// Complete the task. This method assumes that the state is RUNNING. + fn complete(self) { + // The future has completed and its output has been written to the task + // stage. We transition from running to complete. - // Transition to `Complete`, notifying the `JoinHandle` if necessary. - transition_to_complete(self.header(), stage, self.trailer()); - } else { - drop(output); + let snapshot = self.header().state.transition_to_complete(); + + // We catch panics here in case dropping the future or waking the + // JoinHandle panics. + let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { + if !snapshot.is_join_interested() { + // The `JoinHandle` is not interested in the output of + // this task. It is our responsibility to drop the + // output. + self.core().stage.drop_future_or_output(); + } else if snapshot.has_join_waker() { + // Notify the join handle. The previous transition obtains the + // lock on the waker cell. + self.trailer().wake_join(); } })); // The task has completed execution and will no longer be scheduled. - // - // Attempts to batch a ref-dec with the state transition below. + let num_release = self.release(); - if self - .scheduler_view() - .transition_to_terminal(is_join_interested) - { - self.dealloc() + if self.header().state.transition_to_terminal(num_release) { + self.dealloc(); } } - fn to_task(&self) -> Task { - self.scheduler_view().to_task() - } -} + /// Release the task from the scheduler. Returns the number of ref-counts + /// that should be decremented. + fn release(&self) -> usize { + // We don't actually increment the ref-count here, but the new task is + // never destroyed, so that's ok. + let me = ManuallyDrop::new(self.get_new_task()); -enum TransitionToRunning { - Ok(Snapshot), - DropReference, -} - -struct SchedulerView<'a, S> { - header: &'a Header, - scheduler: &'a S, -} - -impl<'a, S> SchedulerView<'a, S> -where - S: Schedule, -{ - fn to_task(&self) -> Task { - // SAFETY The header is from the same struct containing the scheduler `S` so the cast is safe - unsafe { Task::from_raw(self.header.into()) } - } - - /// Returns true if the task should be deallocated. - fn transition_to_terminal(&self, is_join_interested: bool) -> bool { - let me = self.to_task(); - - let ref_dec = if let Some(task) = self.scheduler.release(&me) { + if let Some(task) = self.core().scheduler.release(&me) { mem::forget(task); - true + 2 } else { - false - }; - - mem::forget(me); - - // This might deallocate - let snapshot = self - .header - .state - .transition_to_terminal(!is_join_interested, ref_dec); - - snapshot.ref_count() == 0 + 1 + } } - fn transition_to_running(&self) -> TransitionToRunning { - // Transition the task to the running state. - // - // A failure to transition here indicates the task has been cancelled - // while in the run queue pending execution. - let snapshot = match self.header.state.transition_to_running() { - Ok(snapshot) => snapshot, - Err(_) => { - // The task was shutdown while in the run queue. At this point, - // we just hold a ref counted reference. Since we do not have access to it here - // return `DropReference` so the caller drops it. - return TransitionToRunning::DropReference; - } - }; - - TransitionToRunning::Ok(snapshot) - } -} - -/// Transitions the task's lifecycle to `Complete`. Notifies the -/// `JoinHandle` if it still has interest in the completion. -fn transition_to_complete(header: &Header, stage: &CoreStage, trailer: &Trailer) -where - T: Future, -{ - // Transition the task's lifecycle to `Complete` and get a snapshot of - // the task's sate. - let snapshot = header.state.transition_to_complete(); - - if !snapshot.is_join_interested() { - // The `JoinHandle` is not interested in the output of this task. It - // is our responsibility to drop the output. - stage.drop_future_or_output(); - } else if snapshot.has_join_waker() { - // Notify the join handle. The previous transition obtains the - // lock on the waker cell. - trailer.wake_join(); + /// Create a new task that holds its own ref-count. + /// + /// # Safety + /// + /// Any use of `self` after this call must ensure that a ref-count to the + /// task holds the task alive until after the use of `self`. Passing the + /// returned Task to any method on `self` is unsound if dropping the Task + /// could drop `self` before the call on `self` returned. + fn get_new_task(&self) -> Task { + // safety: The header is at the beginning of the cell, so this cast is + // safe. + unsafe { Task::from_raw(self.cell.cast()) } } } @@ -389,73 +418,62 @@ fn set_join_waker( res } -enum PollFuture { - Complete(Result, bool), - DropReference, +enum PollFuture { + Complete, Notified, - None, + Done, + Dealloc, } -fn cancel_task(stage: &CoreStage) -> JoinError { +/// Cancel the task and store the appropriate error in the stage field. +fn cancel_task(stage: &CoreStage) { // Drop the future from a panic guard. let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { stage.drop_future_or_output(); })); - if let Err(err) = res { - // Dropping the future panicked, complete the join - // handle with the panic to avoid dropping the panic - // on the ground. - JoinError::panic(err) - } else { - JoinError::cancelled() - } -} - -fn poll_future( - header: &Header, - core: &CoreStage, - snapshot: Snapshot, - cx: Context<'_>, -) -> PollFuture { - if snapshot.is_cancelled() { - PollFuture::Complete(Err(cancel_task(core)), snapshot.is_join_interested()) - } else { - let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { - struct Guard<'a, T: Future> { - core: &'a CoreStage, - } - - impl Drop for Guard<'_, T> { - fn drop(&mut self) { - self.core.drop_future_or_output(); - } - } - - let guard = Guard { core }; - - let res = guard.core.poll(cx); - - // prevent the guard from dropping the future - mem::forget(guard); - - res - })); - match res { - Ok(Poll::Pending) => match header.state.transition_to_idle() { - Ok(snapshot) => { - if snapshot.is_notified() { - PollFuture::Notified - } else { - PollFuture::None - } - } - Err(_) => PollFuture::Complete(Err(cancel_task(core)), true), - }, - Ok(Poll::Ready(ok)) => PollFuture::Complete(Ok(ok), snapshot.is_join_interested()), - Err(err) => { - PollFuture::Complete(Err(JoinError::panic(err)), snapshot.is_join_interested()) - } + match res { + Ok(()) => { + stage.store_output(Err(JoinError::cancelled())); + } + Err(panic) => { + stage.store_output(Err(JoinError::panic(panic))); } } } + +/// Poll the future. If the future completes, the output is written to the +/// stage field. +fn poll_future(core: &CoreStage, cx: Context<'_>) -> Poll<()> { + // Poll the future. + let output = panic::catch_unwind(panic::AssertUnwindSafe(|| { + struct Guard<'a, T: Future> { + core: &'a CoreStage, + } + impl<'a, T: Future> Drop for Guard<'a, T> { + fn drop(&mut self) { + // If the future panics on poll, we drop it inside the panic + // guard. + self.core.drop_future_or_output(); + } + } + let guard = Guard { core }; + let res = guard.core.poll(cx); + mem::forget(guard); + res + })); + + // Prepare output for being placed in the core stage. + let output = match output { + Ok(Poll::Pending) => return Poll::Pending, + Ok(Poll::Ready(output)) => Ok(output), + Err(panic) => Err(JoinError::panic(panic)), + }; + + // Catch and ignore panics if the future panics on drop. + let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { + core.store_output(output); + })); + + Poll::Ready(()) +} diff --git a/tokio/src/runtime/task/inject.rs b/tokio/src/runtime/task/inject.rs index 640da648f..d1f0aee9d 100644 --- a/tokio/src/runtime/task/inject.rs +++ b/tokio/src/runtime/task/inject.rs @@ -92,6 +92,8 @@ impl Inject { debug_assert!(get_next(task).is_none()); if let Some(tail) = p.tail { + // safety: Holding the Notified for a task guarantees exclusive + // access to the `queue_next` field. set_next(tail, Some(task)); } else { p.head = Some(task); @@ -103,9 +105,6 @@ impl Inject { } /// Pushes several values into the queue. - /// - /// SAFETY: The caller should ensure that we have exclusive access to the - /// `queue_next` field in the provided tasks. #[inline] pub(crate) fn push_batch(&self, mut iter: I) where @@ -123,8 +122,11 @@ impl Inject { // We are going to be called with an `std::iter::Chain`, and that // iterator overrides `for_each` to something that is easier for the // compiler to optimize than a loop. - iter.map(|next| next.into_raw()).for_each(|next| { - // safety: The caller guarantees exclusive access to this field. + iter.for_each(|next| { + let next = next.into_raw(); + + // safety: Holding the Notified for a task guarantees exclusive + // access to the `queue_next` field. set_next(prev, Some(next)); prev = next; counter += 1; diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index 3be8ed3b6..edd3c4fac 100644 --- a/tokio/src/runtime/task/list.rs +++ b/tokio/src/runtime/task/list.rs @@ -7,6 +7,7 @@ //! the scheduler with the collection. use crate::future::Future; +use crate::loom::cell::UnsafeCell; use crate::loom::sync::Mutex; use crate::runtime::task::{JoinHandle, LocalNotified, Notified, Schedule, Task}; use crate::util::linked_list::{Link, LinkedList}; @@ -56,18 +57,16 @@ pub(crate) struct OwnedTasks { inner: Mutex>, id: u64, } +pub(crate) struct LocalOwnedTasks { + inner: UnsafeCell>, + id: u64, + _not_send_or_sync: PhantomData<*const ()>, +} struct OwnedTasksInner { list: LinkedList, as Link>::Target>, closed: bool, } -pub(crate) struct LocalOwnedTasks { - list: LinkedList, as Link>::Target>, - closed: bool, - id: u64, - _not_send_or_sync: PhantomData<*const ()>, -} - impl OwnedTasks { pub(crate) fn new() -> Self { Self { @@ -115,7 +114,7 @@ impl OwnedTasks { /// a LocalNotified, giving the thread permission to poll this task. #[inline] pub(crate) fn assert_owner(&self, task: Notified) -> LocalNotified { - assert_eq!(task.0.header().get_owner_id(), self.id); + assert_eq!(task.header().get_owner_id(), self.id); // safety: All tasks bound to this OwnedTasks are Send, so it is safe // to poll it on this thread no matter what thread we are on. @@ -125,8 +124,32 @@ impl OwnedTasks { } } - pub(crate) fn pop_back(&self) -> Option> { - self.inner.lock().list.pop_back() + /// Shut down all tasks in the collection. This call also closes the + /// collection, preventing new items from being added. + pub(crate) fn close_and_shutdown_all(&self) + where + S: Schedule, + { + // The first iteration of the loop was unrolled so it can set the + // closed bool. + let first_task = { + let mut lock = self.inner.lock(); + lock.closed = true; + lock.list.pop_back() + }; + match first_task { + Some(task) => task.shutdown(), + None => return, + } + + loop { + let task = match self.inner.lock().list.pop_back() { + Some(task) => task, + None => return, + }; + + task.shutdown(); + } } pub(crate) fn remove(&self, task: &Task) -> Option> { @@ -146,30 +169,22 @@ impl OwnedTasks { pub(crate) fn is_empty(&self) -> bool { self.inner.lock().list.is_empty() } - - #[cfg(feature = "rt-multi-thread")] - pub(crate) fn is_closed(&self) -> bool { - self.inner.lock().closed - } - - /// Close the OwnedTasks. This prevents adding new tasks to the collection. - pub(crate) fn close(&self) { - self.inner.lock().closed = true; - } } impl LocalOwnedTasks { pub(crate) fn new() -> Self { Self { - list: LinkedList::new(), - closed: false, + inner: UnsafeCell::new(OwnedTasksInner { + list: LinkedList::new(), + closed: false, + }), id: get_next_id(), _not_send_or_sync: PhantomData, } } pub(crate) fn bind( - &mut self, + &self, task: T, scheduler: S, ) -> (JoinHandle, Option>) @@ -186,21 +201,32 @@ impl LocalOwnedTasks { task.header().set_owner_id(self.id); } - if self.closed { + if self.is_closed() { drop(notified); task.shutdown(); (join, None) } else { - self.list.push_front(task); + self.with_inner(|inner| { + inner.list.push_front(task); + }); (join, Some(notified)) } } - pub(crate) fn pop_back(&mut self) -> Option> { - self.list.pop_back() + /// Shut down all tasks in the collection. This call also closes the + /// collection, preventing new items from being added. + pub(crate) fn close_and_shutdown_all(&self) + where + S: Schedule, + { + self.with_inner(|inner| inner.closed = true); + + while let Some(task) = self.with_inner(|inner| inner.list.pop_back()) { + task.shutdown(); + } } - pub(crate) fn remove(&mut self, task: &Task) -> Option> { + pub(crate) fn remove(&self, task: &Task) -> Option> { let task_id = task.header().get_owner_id(); if task_id == 0 { // The task is unowned. @@ -209,16 +235,17 @@ impl LocalOwnedTasks { assert_eq!(task_id, self.id); - // safety: We just checked that the provided task is not in some other - // linked list. - unsafe { self.list.remove(task.header().into()) } + self.with_inner(|inner| + // safety: We just checked that the provided task is not in some + // other linked list. + unsafe { inner.list.remove(task.header().into()) }) } /// Assert that the given task is owned by this LocalOwnedTasks and convert /// it to a LocalNotified, giving the thread permission to poll this task. #[inline] pub(crate) fn assert_owner(&self, task: Notified) -> LocalNotified { - assert_eq!(task.0.header().get_owner_id(), self.id); + assert_eq!(task.header().get_owner_id(), self.id); // safety: The task was bound to this LocalOwnedTasks, and the // LocalOwnedTasks is not Send or Sync, so we are on the right thread @@ -229,14 +256,23 @@ impl LocalOwnedTasks { } } - pub(crate) fn is_empty(&self) -> bool { - self.list.is_empty() + #[inline] + fn with_inner(&self, f: F) -> T + where + F: FnOnce(&mut OwnedTasksInner) -> T, + { + // safety: This type is not Sync, so concurrent calls of this method + // can't happen. Furthermore, all uses of this method in this file make + // sure that they don't call `with_inner` recursively. + self.inner.with_mut(|ptr| unsafe { f(&mut *ptr) }) } - /// Close the LocalOwnedTasks. This prevents adding new tasks to the - /// collection. - pub(crate) fn close(&mut self) { - self.closed = true; + pub(crate) fn is_closed(&self) -> bool { + self.with_inner(|inner| inner.closed) + } + + pub(crate) fn is_empty(&self) -> bool { + self.with_inner(|inner| inner.list.is_empty()) } } diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index adc91d81d..cc3910d4b 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -1,3 +1,140 @@ +//! The task module. +//! +//! The task module contains the code that manages spawned tasks and provides a +//! safe API for the rest of the runtime to use. Each task in a runtime is +//! stored in an OwnedTasks or LocalOwnedTasks object. +//! +//! # Task reference types +//! +//! A task is usually referenced by multiple handles, and there are several +//! types of handles. +//! +//! * OwnedTask - tasks stored in an OwnedTasks or LocalOwnedTasks are of this +//! reference type. +//! +//! * JoinHandle - each task has a JoinHandle that allows access to the output +//! of the task. +//! +//! * Waker - every waker for a task has this reference type. There can be any +//! number of waker references. +//! +//! * Notified - tracks whether the task is notified. +//! +//! * Unowned - this task reference type is used for tasks not stored in any +//! runtime. Mainly used for blocking tasks, but also in tests. +//! +//! The task uses a reference count to keep track of how many active references +//! exist. The Unowned reference type takes up two ref-counts. All other +//! reference types take pu a single ref-count. +//! +//! Besides the waker type, each task has at most one of each reference type. +//! +//! # State +//! +//! The task stores its state in an atomic usize with various bitfields for the +//! necessary information. The state has the following bitfields: +//! +//! * RUNNING - Tracks whether the task is currently being polled or cancelled. +//! This bit functions as a lock around the task. +//! +//! * COMPLETE - Is one once the future has fully completed and has been +//! dropped. Never unset once set. Never set together with RUNNING. +//! +//! * NOTIFIED - Tracks whether a Notified object currently exists. +//! +//! * CANCELLED - Is set to one for tasks that should be cancelled as soon as +//! possible. May take any value for completed tasks. +//! +//! * JOIN_INTEREST - Is set to one if there exists a JoinHandle. +//! +//! * JOIN_WAKER - Is set to one if the JoinHandle has set a waker. +//! +//! The rest of the bits are used for the ref-count. +//! +//! # Fields in the task +//! +//! The task has various fields. This section describes how and when it is safe +//! to access a field. +//! +//! * The state field is accessed with atomic instructions. +//! +//! * The OwnedTask reference has exclusive access to the `owned` field. +//! +//! * The Notified reference has exclusive access to the `queue_next` field. +//! +//! * The `owner_id` field can be set as part of construction of the task, but +//! is otherwise immutable and anyone can access the field immutably without +//! synchronization. +//! +//! * If COMPLETE is one, then the JoinHandle has exclusive access to the +//! stage field. If COMPLETE is zero, then the RUNNING bitfield functions as +//! a lock for the stage field, and it can be accessed only by the thread +//! that set RUNNING to one. +//! +//! * If JOIN_WAKER is zero, then the JoinHandle has exclusive access to the +//! join handle waker. If JOIN_WAKER and COMPLETE are both one, then the +//! thread that set COMPLETE to one has exclusive access to the join handle +//! waker. +//! +//! All other fields are immutable and can be accessed immutably without +//! synchronization by anyone. +//! +//! # Safety +//! +//! This section goes through various situations and explains why the API is +//! safe in that situation. +//! +//! ## Polling or dropping the future +//! +//! Any mutable access to the future happens after obtaining a lock by modifying +//! the RUNNING field, so exclusive access is ensured. +//! +//! When the task completes, exclusive access to the output is transferred to +//! the JoinHandle. If the JoinHandle is already dropped when the transition to +//! complete happens, the thread performing that transition retains exclusive +//! access to the output and should immediately drop it. +//! +//! ## Non-Send futures +//! +//! If a future is not Send, then it is bound to a LocalOwnedTasks. The future +//! will only ever be polled or dropped given a LocalNotified or inside a call +//! to LocalOwnedTasks::shutdown_all. In either case, it is guaranteed that the +//! future is on the right thread. +//! +//! If the task is never removed from the LocalOwnedTasks, then it is leaked, so +//! there is no risk that the task is dropped on some other thread when the last +//! ref-count drops. +//! +//! ## Non-Send output +//! +//! When a task completes, the output is placed in the stage of the task. Then, +//! a transition that sets COMPLETE to true is performed, and the value of +//! JOIN_INTEREST when this transition happens is read. +//! +//! If JOIN_INTEREST is zero when the transition to COMPLETE happens, then the +//! output is immediately dropped. +//! +//! If JOIN_INTEREST is one when the transition to COMPLETE happens, then the +//! JoinHandle is responsible for cleaning up the output. If the output is not +//! Send, then this happens: +//! +//! 1. The output is created on the thread that the future was polled on. Since +//! only non-Send futures can have non-Send output, the future was polled on +//! the thread that the future was spawned from. +//! 2. Since JoinHandle is not Send if Output is not Send, the +//! JoinHandle is also on the thread that the future was spawned from. +//! 3. Thus, the JoinHandle will not move the output across threads when it +//! takes or drops the output. +//! +//! ## Recursive poll/shutdown +//! +//! Calling poll from inside a shutdown call or vice-versa is not prevented by +//! the API exposed by the task module, so this has to be safe. In either case, +//! the lock in the RUNNING bitfield makes the inner call return immediately. If +//! the inner call is a `shutdown` call, then the CANCELLED bit is set, and the +//! poll call will notice it when the poll finishes, and the task is cancelled +//! at that point. + mod core; use self::core::Cell; use self::core::Header; @@ -161,6 +298,12 @@ impl Task { } } +impl Notified { + fn header(&self) -> &Header { + self.0.header() + } +} + cfg_rt_multi_thread! { impl Notified { unsafe fn from_raw(ptr: NonNull

) -> Notified { @@ -185,16 +328,19 @@ cfg_rt_multi_thread! { impl Task { /// Pre-emptively cancel the task as part of the shutdown process. - pub(crate) fn shutdown(&self) { - self.raw.shutdown(); + pub(crate) fn shutdown(self) { + let raw = self.raw; + mem::forget(self); + raw.shutdown(); } } impl LocalNotified { /// Run the task pub(crate) fn run(self) { - self.task.raw.poll(); + let raw = self.task.raw; mem::forget(self); + raw.poll(); } } @@ -220,11 +366,13 @@ impl UnownedTask { } pub(crate) fn run(self) { - // Decrement the ref-count - self.raw.header().state.ref_dec(); - // Poll the task - self.raw.poll(); + let raw = self.raw; mem::forget(self); + + // Poll the task + raw.poll(); + // Decrement our extra ref-count + raw.header().state.ref_dec(); } pub(crate) fn shutdown(self) { diff --git a/tokio/src/runtime/task/state.rs b/tokio/src/runtime/task/state.rs index 3641af083..059a7f9ee 100644 --- a/tokio/src/runtime/task/state.rs +++ b/tokio/src/runtime/task/state.rs @@ -64,6 +64,35 @@ const REF_ONE: usize = 1 << REF_COUNT_SHIFT; /// As the task starts with a `Notified`, `NOTIFIED` is set. const INITIAL_STATE: usize = (REF_ONE * 3) | JOIN_INTEREST | NOTIFIED; +#[must_use] +pub(super) enum TransitionToRunning { + Success, + Cancelled, + Failed, + Dealloc, +} + +#[must_use] +pub(super) enum TransitionToIdle { + Ok, + OkNotified, + OkDealloc, + Cancelled, +} + +#[must_use] +pub(super) enum TransitionToNotifiedByVal { + DoNothing, + Submit, + Dealloc, +} + +#[must_use] +pub(super) enum TransitionToNotifiedByRef { + DoNothing, + Submit, +} + /// All transitions are performed via RMW operations. This establishes an /// unambiguous modification order. impl State { @@ -81,51 +110,72 @@ impl State { Snapshot(self.val.load(Acquire)) } - /// Attempt to transition the lifecycle to `Running`. - /// - /// The `NOTIFIED` bit is always unset. - pub(super) fn transition_to_running(&self) -> UpdateResult { - self.fetch_update(|curr| { - assert!(curr.is_notified()); - - let mut next = curr; + /// Attempt to transition the lifecycle to `Running`. This sets the + /// notified bit to false so notifications during the poll can be detected. + pub(super) fn transition_to_running(&self) -> TransitionToRunning { + self.fetch_update_action(|mut next| { + let action; + assert!(next.is_notified()); if !next.is_idle() { - return None; - } + // This happens if the task is either currently running or if it + // has already completed, e.g. if it was cancelled during + // shutdown. Consume the ref-count and return. + next.ref_dec(); + if next.ref_count() == 0 { + action = TransitionToRunning::Dealloc; + } else { + action = TransitionToRunning::Failed; + } + } else { + // We are able to lock the RUNNING bit. + next.set_running(); + next.unset_notified(); - next.set_running(); - next.unset_notified(); - Some(next) + if next.is_cancelled() { + action = TransitionToRunning::Cancelled; + } else { + action = TransitionToRunning::Success; + } + } + (action, Some(next)) }) } /// Transitions the task from `Running` -> `Idle`. /// - /// Returns `Ok` if the transition to `Idle` is successful, `Err` otherwise. - /// In both cases, a snapshot of the state from **after** the transition is - /// returned. - /// + /// Returns `true` if the transition to `Idle` is successful, `false` otherwise. /// The transition to `Idle` fails if the task has been flagged to be /// cancelled. - pub(super) fn transition_to_idle(&self) -> UpdateResult { - self.fetch_update(|curr| { + pub(super) fn transition_to_idle(&self) -> TransitionToIdle { + self.fetch_update_action(|curr| { assert!(curr.is_running()); if curr.is_cancelled() { - return None; + return (TransitionToIdle::Cancelled, None); } let mut next = curr; + let action; next.unset_running(); - if next.is_notified() { - // The caller needs to schedule the task. To do this, it needs a - // waker. The waker requires a ref count. + if !next.is_notified() { + // Polling the future consumes the ref-count of the Notified. + next.ref_dec(); + if next.ref_count() == 0 { + action = TransitionToIdle::OkDealloc; + } else { + action = TransitionToIdle::Ok; + } + } else { + // The caller will schedule a new notification, so we create a + // new ref-count for the notification. Our own ref-count is kept + // for now, and the caller will drop it shortly. next.ref_inc(); + action = TransitionToIdle::OkNotified; } - Some(next) + (action, Some(next)) }) } @@ -141,47 +191,119 @@ impl State { } /// Transition from `Complete` -> `Terminal`, decrementing the reference - /// count by 1. + /// count the specified number of times. /// - /// When `ref_dec` is set, an additional ref count decrement is performed. - /// This is used to batch atomic ops when possible. - pub(super) fn transition_to_terminal(&self, complete: bool, ref_dec: bool) -> Snapshot { - self.fetch_update(|mut snapshot| { - if complete { - snapshot.set_complete(); - } else { - assert!(snapshot.is_complete()); - } - - // Decrement the primary handle - snapshot.ref_dec(); - - if ref_dec { - // Decrement a second time - snapshot.ref_dec(); - } - - Some(snapshot) - }) - .unwrap() + /// Returns true if the task should be deallocated. + pub(super) fn transition_to_terminal(&self, count: usize) -> bool { + let prev = Snapshot(self.val.fetch_sub(count * REF_ONE, AcqRel)); + assert!( + prev.ref_count() >= count, + "current: {}, sub: {}", + prev.ref_count(), + count + ); + prev.ref_count() == count } /// Transitions the state to `NOTIFIED`. /// - /// Returns `true` if the task needs to be submitted to the pool for - /// execution - pub(super) fn transition_to_notified(&self) -> bool { - let prev = Snapshot(self.val.fetch_or(NOTIFIED, AcqRel)); - prev.will_need_queueing() + /// If no task needs to be submitted, a ref-count is consumed. + /// + /// If a task needs to be submitted, the ref-count is incremented for the + /// new Notified. + pub(super) fn transition_to_notified_by_val(&self) -> TransitionToNotifiedByVal { + self.fetch_update_action(|mut snapshot| { + let action; + + if snapshot.is_running() { + // If the task is running, we mark it as notified, but we should + // not submit anything as the thread currently running the + // future is responsible for that. + snapshot.set_notified(); + snapshot.ref_dec(); + + // The thread that set the running bit also holds a ref-count. + assert!(snapshot.ref_count() > 0); + + action = TransitionToNotifiedByVal::DoNothing; + } else if snapshot.is_complete() || snapshot.is_notified() { + // We do not need to submit any notifications, but we have to + // decrement the ref-count. + snapshot.ref_dec(); + + if snapshot.ref_count() == 0 { + action = TransitionToNotifiedByVal::Dealloc; + } else { + action = TransitionToNotifiedByVal::DoNothing; + } + } else { + // We create a new notified that we can submit. The caller + // retains ownership of the ref-count they passed in. + snapshot.set_notified(); + snapshot.ref_inc(); + action = TransitionToNotifiedByVal::Submit; + } + + (action, Some(snapshot)) + }) } - /// Set the cancelled bit and transition the state to `NOTIFIED`. + /// Transitions the state to `NOTIFIED`. + pub(super) fn transition_to_notified_by_ref(&self) -> TransitionToNotifiedByRef { + self.fetch_update_action(|mut snapshot| { + if snapshot.is_complete() || snapshot.is_notified() { + // There is nothing to do in this case. + (TransitionToNotifiedByRef::DoNothing, None) + } else if snapshot.is_running() { + // If the task is running, we mark it as notified, but we should + // not submit as the thread currently running the future is + // responsible for that. + snapshot.set_notified(); + (TransitionToNotifiedByRef::DoNothing, Some(snapshot)) + } else { + // The task is idle and not notified. We should submit a + // notification. + snapshot.set_notified(); + snapshot.ref_inc(); + (TransitionToNotifiedByRef::Submit, Some(snapshot)) + } + }) + } + + /// Set the cancelled bit and transition the state to `NOTIFIED` if idle. /// /// Returns `true` if the task needs to be submitted to the pool for /// execution pub(super) fn transition_to_notified_and_cancel(&self) -> bool { - let prev = Snapshot(self.val.fetch_or(NOTIFIED | CANCELLED, AcqRel)); - prev.will_need_queueing() + self.fetch_update_action(|mut snapshot| { + if snapshot.is_cancelled() || snapshot.is_complete() { + // Aborts to completed or cancelled tasks are no-ops. + (false, None) + } else if snapshot.is_running() { + // If the task is running, we mark it as cancelled. The thread + // running the task will notice the cancelled bit when it + // stops polling and it will kill the task. + // + // The set_notified() call is not strictly necessary but it will + // in some cases let a wake_by_ref call return without having + // to perform a compare_exchange. + snapshot.set_notified(); + snapshot.set_cancelled(); + (false, Some(snapshot)) + } else { + // The task is idle. We set the cancelled and notified bits and + // submit a notification if the notified bit was not already + // set. + snapshot.set_cancelled(); + if !snapshot.is_notified() { + snapshot.set_notified(); + snapshot.ref_inc(); + (true, Some(snapshot)) + } else { + (false, Some(snapshot)) + } + } + }) } /// Set the `CANCELLED` bit and attempt to transition to `Running`. @@ -195,17 +317,11 @@ impl State { if snapshot.is_idle() { snapshot.set_running(); - - if snapshot.is_notified() { - // If the task is idle and notified, this indicates the task is - // in the run queue and is considered owned by the scheduler. - // The shutdown operation claims ownership of the task, which - // means we need to assign an additional ref-count to the task - // in the queue. - snapshot.ref_inc(); - } } + // If the task was not idle, the thread currently running the task + // will notice the cancelled bit and cancel it once the poll + // completes. snapshot.set_cancelled(); Some(snapshot) }); @@ -321,15 +437,39 @@ impl State { /// Returns `true` if the task should be released. pub(super) fn ref_dec(&self) -> bool { let prev = Snapshot(self.val.fetch_sub(REF_ONE, AcqRel)); + assert!(prev.ref_count() >= 1); prev.ref_count() == 1 } /// Returns `true` if the task should be released. pub(super) fn ref_dec_twice(&self) -> bool { let prev = Snapshot(self.val.fetch_sub(2 * REF_ONE, AcqRel)); + assert!(prev.ref_count() >= 2); prev.ref_count() == 2 } + fn fetch_update_action(&self, mut f: F) -> T + where + F: FnMut(Snapshot) -> (T, Option), + { + let mut curr = self.load(); + + loop { + let (output, next) = f(curr); + let next = match next { + Some(next) => next, + None => return output, + }; + + let res = self.val.compare_exchange(curr.0, next.0, AcqRel, Acquire); + + match res { + Ok(_) => return output, + Err(actual) => curr = Snapshot(actual), + } + } + } + fn fetch_update(&self, mut f: F) -> Result where F: FnMut(Snapshot) -> Option, @@ -369,6 +509,10 @@ impl Snapshot { self.0 &= !NOTIFIED } + fn set_notified(&mut self) { + self.0 |= NOTIFIED + } + pub(super) fn is_running(self) -> bool { self.0 & RUNNING == RUNNING } @@ -389,10 +533,6 @@ impl Snapshot { self.0 |= CANCELLED; } - fn set_complete(&mut self) { - self.0 |= COMPLETE; - } - /// Returns `true` if the task's future has completed execution. pub(super) fn is_complete(self) -> bool { self.0 & COMPLETE == COMPLETE @@ -431,10 +571,6 @@ impl Snapshot { assert!(self.ref_count() > 0); self.0 -= REF_ONE } - - fn will_need_queueing(self) -> bool { - !self.is_notified() && self.is_idle() - } } impl fmt::Debug for State { diff --git a/tokio/src/runtime/tests/loom_pool.rs b/tokio/src/runtime/tests/loom_pool.rs index d9ef52069..b3ecd4312 100644 --- a/tokio/src/runtime/tests/loom_pool.rs +++ b/tokio/src/runtime/tests/loom_pool.rs @@ -209,23 +209,6 @@ mod group_b { blocking_and_regular_inner(true); } - #[test] - fn pool_shutdown() { - loom::model(|| { - let pool = mk_pool(2); - - pool.spawn(track(async move { - gated2(true).await; - })); - - pool.spawn(track(async move { - gated2(false).await; - })); - - drop(pool); - }); - } - #[test] fn join_output() { loom::model(|| { @@ -274,10 +257,6 @@ mod group_b { }); }); } -} - -mod group_c { - use super::*; #[test] fn shutdown_with_notification() { @@ -306,6 +285,27 @@ mod group_c { } } +mod group_c { + use super::*; + + #[test] + fn pool_shutdown() { + loom::model(|| { + let pool = mk_pool(2); + + pool.spawn(track(async move { + gated2(true).await; + })); + + pool.spawn(track(async move { + gated2(false).await; + })); + + drop(pool); + }); + } +} + mod group_d { use super::*; diff --git a/tokio/src/runtime/tests/task.rs b/tokio/src/runtime/tests/task.rs index 39cb1d220..e93a1efe6 100644 --- a/tokio/src/runtime/tests/task.rs +++ b/tokio/src/runtime/tests/task.rs @@ -259,10 +259,7 @@ impl Runtime { fn shutdown(&self) { let mut core = self.0.core.try_lock().unwrap(); - self.0.owned.close(); - while let Some(task) = self.0.owned.pop_back() { - task.shutdown(); - } + self.0.owned.close_and_shutdown_all(); while let Some(task) = core.queue.pop_back() { drop(task); diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs index 608a7353a..f5004c0e4 100644 --- a/tokio/src/runtime/thread_pool/worker.rs +++ b/tokio/src/runtime/thread_pool/worker.rs @@ -599,13 +599,8 @@ impl Core { /// Signals all tasks to shut down, and waits for them to complete. Must run /// before we enter the single-threaded phase of shutdown processing. fn pre_shutdown(&mut self, worker: &Worker) { - // The OwnedTasks was closed in Shared::close. - debug_assert!(worker.shared.owned.is_closed()); - // Signal to all tasks to shut down. - while let Some(header) = worker.shared.owned.pop_back() { - header.shutdown(); - } + worker.shared.owned.close_and_shutdown_all(); } /// Shutdown the core @@ -707,7 +702,6 @@ impl Shared { pub(super) fn close(&self) { if self.inject.close() { - self.owned.close(); self.notify_all(); } } diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 6c0939b6b..a28d793c2 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -2,8 +2,9 @@ use crate::loom::sync::{Arc, Mutex}; use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, Task}; use crate::sync::AtomicWaker; +use crate::util::VecDequeCell; -use std::cell::{Cell, RefCell}; +use std::cell::Cell; use std::collections::VecDeque; use std::fmt; use std::future::Future; @@ -223,19 +224,14 @@ cfg_rt! { /// State available from the thread-local struct Context { - /// Owned task set and local run queue - tasks: RefCell, - - /// State shared between threads. - shared: Arc, -} - -struct Tasks { /// Collection of all active tasks spawned onto this executor. owned: LocalOwnedTasks>, /// Local run queue sender and receiver. - queue: VecDeque>>, + queue: VecDequeCell>>, + + /// State shared between threads. + shared: Arc, } /// LocalSet state shared between threads. @@ -308,7 +304,7 @@ cfg_rt! { let cx = maybe_cx .expect("`spawn_local` called from outside of a `task::LocalSet`"); - let (handle, notified) = cx.tasks.borrow_mut().owned.bind(future, cx.shared.clone()); + let (handle, notified) = cx.owned.bind(future, cx.shared.clone()); if let Some(notified) = notified { cx.shared.schedule(notified); @@ -334,10 +330,8 @@ impl LocalSet { LocalSet { tick: Cell::new(0), context: Context { - tasks: RefCell::new(Tasks { - owned: LocalOwnedTasks::new(), - queue: VecDeque::with_capacity(INITIAL_CAPACITY), - }), + owned: LocalOwnedTasks::new(), + queue: VecDequeCell::with_capacity(INITIAL_CAPACITY), shared: Arc::new(Shared { queue: Mutex::new(Some(VecDeque::with_capacity(INITIAL_CAPACITY))), waker: AtomicWaker::new(), @@ -391,12 +385,7 @@ impl LocalSet { { let future = crate::util::trace::task(future, "local", None); - let (handle, notified) = self - .context - .tasks - .borrow_mut() - .owned - .bind(future, self.context.shared.clone()); + let (handle, notified) = self.context.owned.bind(future, self.context.shared.clone()); if let Some(notified) = notified { self.context.shared.schedule(notified); @@ -551,24 +540,19 @@ impl LocalSet { .lock() .as_mut() .and_then(|queue| queue.pop_front()) - .or_else(|| self.context.tasks.borrow_mut().queue.pop_front()) + .or_else(|| self.context.queue.pop_front()) } else { - self.context - .tasks - .borrow_mut() - .queue - .pop_front() - .or_else(|| { - self.context - .shared - .queue - .lock() - .as_mut() - .and_then(|queue| queue.pop_front()) - }) + self.context.queue.pop_front().or_else(|| { + self.context + .shared + .queue + .lock() + .as_mut() + .and_then(|queue| queue.pop_front()) + }) }; - task.map(|task| self.context.tasks.borrow_mut().owned.assert_owner(task)) + task.map(|task| self.context.owned.assert_owner(task)) } fn with(&self, f: impl FnOnce() -> T) -> T { @@ -594,7 +578,7 @@ impl Future for LocalSet { // there are still tasks remaining in the run queue. cx.waker().wake_by_ref(); Poll::Pending - } else if self.context.tasks.borrow().owned.is_empty() { + } else if self.context.owned.is_empty() { // If the scheduler has no remaining futures, we're done! Poll::Ready(()) } else { @@ -615,27 +599,13 @@ impl Default for LocalSet { impl Drop for LocalSet { fn drop(&mut self) { self.with(|| { - // Close the LocalOwnedTasks. This ensures that any calls to - // spawn_local in the destructor of a future on this LocalSet will - // immediately cancel the task, and prevents the task from being - // added to `owned`. - self.context.tasks.borrow_mut().owned.close(); - - // Loop required here to ensure borrow is dropped between iterations - #[allow(clippy::while_let_loop)] - loop { - let task = match self.context.tasks.borrow_mut().owned.pop_back() { - Some(task) => task, - None => break, - }; - - // Safety: same as `run_unchecked`. - task.shutdown(); - } + // Shut down all tasks in the LocalOwnedTasks and close it to + // prevent new tasks from ever being added. + self.context.owned.close_and_shutdown_all(); // We already called shutdown on all tasks above, so there is no // need to call shutdown. - for task in self.context.tasks.borrow_mut().queue.drain(..) { + for task in self.context.queue.take() { drop(task); } @@ -646,7 +616,7 @@ impl Drop for LocalSet { drop(task); } - assert!(self.context.tasks.borrow().owned.is_empty()); + assert!(self.context.owned.is_empty()); }); } } @@ -689,7 +659,7 @@ impl Shared { fn schedule(&self, task: task::Notified>) { CURRENT.with(|maybe_cx| match maybe_cx { Some(cx) if cx.shared.ptr_eq(self) => { - cx.tasks.borrow_mut().queue.push_back(task); + cx.queue.push_back(task); } _ => { // First check whether the queue is still there (if not, the @@ -716,7 +686,7 @@ impl task::Schedule for Arc { CURRENT.with(|maybe_cx| { let cx = maybe_cx.expect("scheduler context missing"); assert!(cx.shared.ptr_eq(self)); - cx.tasks.borrow_mut().owned.remove(task) + cx.owned.remove(task) }) } diff --git a/tokio/src/util/mod.rs b/tokio/src/util/mod.rs index bcc73ed68..9065f50a8 100644 --- a/tokio/src/util/mod.rs +++ b/tokio/src/util/mod.rs @@ -24,6 +24,9 @@ cfg_rt! { mod sync_wrapper; pub(crate) use sync_wrapper::SyncWrapper; + + mod vec_deque_cell; + pub(crate) use vec_deque_cell::VecDequeCell; } cfg_rt_multi_thread! { diff --git a/tokio/src/util/vec_deque_cell.rs b/tokio/src/util/vec_deque_cell.rs new file mode 100644 index 000000000..12883abdc --- /dev/null +++ b/tokio/src/util/vec_deque_cell.rs @@ -0,0 +1,53 @@ +use crate::loom::cell::UnsafeCell; + +use std::collections::VecDeque; +use std::marker::PhantomData; + +/// This type is like VecDeque, except that it is not Sync and can be modified +/// through immutable references. +pub(crate) struct VecDequeCell { + inner: UnsafeCell>, + _not_sync: PhantomData<*const ()>, +} + +// This is Send for the same reasons that RefCell> is Send. +unsafe impl Send for VecDequeCell {} + +impl VecDequeCell { + pub(crate) fn with_capacity(cap: usize) -> Self { + Self { + inner: UnsafeCell::new(VecDeque::with_capacity(cap)), + _not_sync: PhantomData, + } + } + + /// Safety: This method may not be called recursively. + #[inline] + unsafe fn with_inner(&self, f: F) -> R + where + F: FnOnce(&mut VecDeque) -> R, + { + // safety: This type is not Sync, so concurrent calls of this method + // cannot happen. Furthermore, the caller guarantees that the method is + // not called recursively. Finally, this is the only place that can + // create mutable references to the inner VecDeque. This ensures that + // any mutable references created here are exclusive. + self.inner.with_mut(|ptr| f(&mut *ptr)) + } + + pub(crate) fn pop_front(&self) -> Option { + unsafe { self.with_inner(VecDeque::pop_front) } + } + + pub(crate) fn push_back(&self, item: T) { + unsafe { + self.with_inner(|inner| inner.push_back(item)); + } + } + + /// Replace the inner VecDeque with an empty VecDeque and return the current + /// contents. + pub(crate) fn take(&self) -> VecDeque { + unsafe { self.with_inner(|inner| std::mem::take(inner)) } + } +}