tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-24 00:00:11 +02:00
Code Issues Packages Projects Releases Wiki Activity

runtime: give Notified a safe API (#4005)

Browse Source
This commit is contained in:
Alice Ryhl
2021-08-12 10:06:05 +02:00
committed by GitHub
parent 032c55e77f
commit b501f25202
13 changed files with 803 additions and 448 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/basic_scheduler.rs
+4 -6
View File
@@ -314,12 +314,10 @@ impl<P: Park> Drop for BasicScheduler<P> {
};
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.
tokio/src/runtime/task/core.rs
+5 -5
View File
@@ -57,10 +57,10 @@ pub(crate) struct Header {
/// Task state
pub(super) state: State,
pub(crate) owned: UnsafeCell<linked_list::Pointers<Header>>,
pub(super) owned: UnsafeCell<linked_list::Pointers<Header>>,
/// Pointer to next task, used with the injection queue
pub(crate) queue_next: UnsafeCell<Option<NonNull<Header>>>,
pub(super) queue_next: UnsafeCell<Option<NonNull<Header>>>,
/// 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<Waker>) {
pub(super) unsafe fn set_waker(&self, waker: Option<Waker>) {
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"),
tokio/src/runtime/task/harness.rs
+244 -226
View File
@@ -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<T, S> {
unsafe { &self.cell.as_ref().core }
}
fn scheduler_view(&self) -> SchedulerView<'_, S> {
SchedulerView {
header: self.header(),
scheduler: &self.core().scheduler,
}
}
}
impl<T, S> Harness<T, S>
@@ -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<T::Output> {
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::<T, S>(self.header());
let cx = Context::from_waker(&*waker_ref);
let res = poll_future(&self.core().stage, cx);
let waker_ref = waker_ref::<T, S>(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<T::Output>, 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<S> {
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<S> {
// 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<T>(header: &Header, stage: &CoreStage<T>, 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<S> {
// 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<T> {
Complete(Result<T, JoinError>, bool),
DropReference,
enum PollFuture {
Complete,
Notified,
None,
Done,
Dealloc,
}
fn cancel_task<T: Future>(stage: &CoreStage<T>) -> JoinError {
/// Cancel the task and store the appropriate error in the stage field.
fn cancel_task<T: Future>(stage: &CoreStage<T>) {
// 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<T: Future>(
header: &Header,
core: &CoreStage<T>,
snapshot: Snapshot,
cx: Context<'_>,
) -> PollFuture<T::Output> {
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<T>,
}
impl<T: Future> 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<T: Future>(core: &CoreStage<T>, cx: Context<'_>) -> Poll<()> {
// Poll the future.
let output = panic::catch_unwind(panic::AssertUnwindSafe(|| {
struct Guard<'a, T: Future> {
core: &'a CoreStage<T>,
}
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(())
}
tokio/src/runtime/task/inject.rs
+7 -5
View File
@@ -92,6 +92,8 @@ impl<T: 'static> Inject<T> {
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<T: 'static> Inject<T> {
}
/// 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<I>(&self, mut iter: I)
where
@@ -123,8 +122,11 @@ impl<T: 'static> Inject<T> {
// 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;
tokio/src/runtime/task/list.rs
+74 -38
View File
@@ -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<S: 'static> {
inner: Mutex<OwnedTasksInner<S>>,
id: u64,
}
pub(crate) struct LocalOwnedTasks<S: 'static> {
inner: UnsafeCell<OwnedTasksInner<S>>,
id: u64,
_not_send_or_sync: PhantomData<*const ()>,
}
struct OwnedTasksInner<S: 'static> {
list: LinkedList<Task<S>, <Task<S> as Link>::Target>,
closed: bool,
}
pub(crate) struct LocalOwnedTasks<S: 'static> {
list: LinkedList<Task<S>, <Task<S> as Link>::Target>,
closed: bool,
id: u64,
_not_send_or_sync: PhantomData<*const ()>,
}
impl<S: 'static> OwnedTasks<S> {
pub(crate) fn new() -> Self {
Self {
@@ -115,7 +114,7 @@ impl<S: 'static> OwnedTasks<S> {
/// a LocalNotified, giving the thread permission to poll this task.
#[inline]
pub(crate) fn assert_owner(&self, task: Notified<S>) -> LocalNotified<S> {
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<S: 'static> OwnedTasks<S> {
}
}
pub(crate) fn pop_back(&self) -> Option<Task<S>> {
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<S>) -> Option<Task<S>> {
@@ -146,30 +169,22 @@ impl<S: 'static> OwnedTasks<S> {
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<S: 'static> LocalOwnedTasks<S> {
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<T>(
&mut self,
&self,
task: T,
scheduler: S,
) -> (JoinHandle<T::Output>, Option<Notified<S>>)
@@ -186,21 +201,32 @@ impl<S: 'static> LocalOwnedTasks<S> {
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<Task<S>> {
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<S>) -> Option<Task<S>> {
pub(crate) fn remove(&self, task: &Task<S>) -> Option<Task<S>> {
let task_id = task.header().get_owner_id();
if task_id == 0 {
// The task is unowned.
@@ -209,16 +235,17 @@ impl<S: 'static> LocalOwnedTasks<S> {
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<S>) -> LocalNotified<S> {
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<S: 'static> LocalOwnedTasks<S> {
}
}
pub(crate) fn is_empty(&self) -> bool {
self.list.is_empty()
#[inline]
fn with_inner<F, T>(&self, f: F) -> T
where
F: FnOnce(&mut OwnedTasksInner<S>) -> 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())
}
}
tokio/src/runtime/task/mod.rs
+155 -7
View File
@@ -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<Output> 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<S: 'static> Task<S> {
}
}
impl<S: 'static> Notified<S> {
fn header(&self) -> &Header {
self.0.header()
}
}
cfg_rt_multi_thread! {
impl<S: 'static> Notified<S> {
unsafe fn from_raw(ptr: NonNull<Header>) -> Notified<S> {
@@ -185,16 +328,19 @@ cfg_rt_multi_thread! {
impl<S: Schedule> Task<S> {
/// 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<S: Schedule> LocalNotified<S> {
/// 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<S: Schedule> UnownedTask<S> {
}
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) {
tokio/src/runtime/task/state.rs
+207 -71
View File
@@ -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<F, T>(&self, mut f: F) -> T
where
F: FnMut(Snapshot) -> (T, Option<Snapshot>),
{
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<F>(&self, mut f: F) -> Result<Snapshot, Snapshot>
where
F: FnMut(Snapshot) -> Option<Snapshot>,
@@ -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 {
tokio/src/runtime/tests/loom_pool.rs
+21 -21
View File
@@ -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::*;
tokio/src/runtime/tests/task.rs
+1 -4
View File
@@ -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);
tokio/src/runtime/thread_pool/worker.rs
+1 -7
View File
@@ -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();
}
}
tokio/src/task/local.rs
+28 -58
View File
@@ -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<Tasks>,
/// State shared between threads.
shared: Arc<Shared>,
}
struct Tasks {
/// Collection of all active tasks spawned onto this executor.
owned: LocalOwnedTasks<Arc<Shared>>,
/// Local run queue sender and receiver.
queue: VecDeque<task::Notified<Arc<Shared>>>,
queue: VecDequeCell<task::Notified<Arc<Shared>>>,
/// State shared between threads.
shared: Arc<Shared>,
}
/// 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<T>(&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<Arc<Self>>) {
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<Shared> {
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)
})
}
tokio/src/util/mod.rs
+3
View File
@@ -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! {
tokio/src/util/vec_deque_cell.rs
+53
View File
@@ -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<T> {
inner: UnsafeCell<VecDeque<T>>,
_not_sync: PhantomData<*const ()>,
}
// This is Send for the same reasons that RefCell<VecDeque<T>> is Send.
unsafe impl<T: Send> Send for VecDequeCell<T> {}
impl<T> VecDequeCell<T> {
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<F, R>(&self, f: F) -> R
where
F: FnOnce(&mut VecDeque<T>) -> 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<T> {
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<T> {
unsafe { self.with_inner(|inner| std::mem::take(inner)) }
}
}