mirror of
https://github.com/tokio-rs/tokio.git
synced 2026-09-09 00:00:08 +02:00
sync: rewrite CancellationToken (#4652)
This commit is contained in:
Finomnis
@@ -1,18 +1,15 @@
|
||||
//! An asynchronously awaitable `CancellationToken`.
|
||||
//! The token allows to signal a cancellation request to one or more tasks.
|
||||
pub(crate) mod guard;
|
||||
mod tree_node;
|
||||
|
||||
use crate::loom::sync::atomic::AtomicUsize;
|
||||
use crate::loom::sync::Mutex;
|
||||
use crate::sync::intrusive_double_linked_list::{LinkedList, ListNode};
|
||||
|
||||
use crate::loom::sync::Arc;
|
||||
use core::future::Future;
|
||||
use core::pin::Pin;
|
||||
use core::ptr::NonNull;
|
||||
use core::sync::atomic::Ordering;
|
||||
use core::task::{Context, Poll, Waker};
|
||||
use core::task::{Context, Poll};
|
||||
|
||||
use guard::DropGuard;
|
||||
use pin_project_lite::pin_project;
|
||||
|
||||
/// A token which can be used to signal a cancellation request to one or more
|
||||
/// tasks.
|
||||
@@ -55,31 +52,20 @@ use guard::DropGuard;
|
||||
/// }
|
||||
/// ```
|
||||
pub struct CancellationToken {
|
||||
inner: NonNull<CancellationTokenState>,
|
||||
inner: Arc<tree_node::TreeNode>,
|
||||
}
|
||||
|
||||
// Safety: The CancellationToken is thread-safe and can be moved between threads,
|
||||
// since all methods are internally synchronized.
|
||||
unsafe impl Send for CancellationToken {}
|
||||
unsafe impl Sync for CancellationToken {}
|
||||
|
||||
/// A Future that is resolved once the corresponding [`CancellationToken`]
|
||||
/// was cancelled
|
||||
#[must_use = "futures do nothing unless polled"]
|
||||
pub struct WaitForCancellationFuture<'a> {
|
||||
/// The CancellationToken that is associated with this WaitForCancellationFuture
|
||||
cancellation_token: Option<&'a CancellationToken>,
|
||||
/// Node for waiting at the cancellation_token
|
||||
wait_node: ListNode<WaitQueueEntry>,
|
||||
/// Whether this future was registered at the token yet as a waiter
|
||||
is_registered: bool,
|
||||
pin_project! {
|
||||
/// A Future that is resolved once the corresponding [`CancellationToken`]
|
||||
/// is cancelled.
|
||||
#[must_use = "futures do nothing unless polled"]
|
||||
pub struct WaitForCancellationFuture<'a> {
|
||||
cancellation_token: &'a CancellationToken,
|
||||
#[pin]
|
||||
future: tokio::sync::futures::Notified<'a>,
|
||||
}
|
||||
}
|
||||
|
||||
// Safety: Futures can be sent between threads as long as the underlying
|
||||
// cancellation_token is thread-safe (Sync),
|
||||
// which allows to poll/register/unregister from a different thread.
|
||||
unsafe impl<'a> Send for WaitForCancellationFuture<'a> {}
|
||||
|
||||
// ===== impl CancellationToken =====
|
||||
|
||||
impl core::fmt::Debug for CancellationToken {
|
||||
@@ -92,43 +78,16 @@ impl core::fmt::Debug for CancellationToken {
|
||||
|
||||
impl Clone for CancellationToken {
|
||||
fn clone(&self) -> Self {
|
||||
// Safety: The state inside a `CancellationToken` is always valid, since
|
||||
// is reference counted
|
||||
let inner = self.state();
|
||||
|
||||
// Tokens are cloned by increasing their refcount
|
||||
let current_state = inner.snapshot();
|
||||
inner.increment_refcount(current_state);
|
||||
|
||||
CancellationToken { inner: self.inner }
|
||||
tree_node::increase_handle_refcount(&self.inner);
|
||||
CancellationToken {
|
||||
inner: self.inner.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CancellationToken {
|
||||
fn drop(&mut self) {
|
||||
let token_state_pointer = self.inner;
|
||||
|
||||
// Safety: The state inside a `CancellationToken` is always valid, since
|
||||
// is reference counted
|
||||
let inner = unsafe { &mut *self.inner.as_ptr() };
|
||||
|
||||
let mut current_state = inner.snapshot();
|
||||
|
||||
// We need to safe the parent, since the state might be released by the
|
||||
// next call
|
||||
let parent = inner.parent;
|
||||
|
||||
// Drop our own refcount
|
||||
current_state = inner.decrement_refcount(current_state);
|
||||
|
||||
// If this was the last reference, unregister from the parent
|
||||
if current_state.refcount == 0 {
|
||||
if let Some(mut parent) = parent {
|
||||
// Safety: Since we still retain a reference on the parent, it must be valid.
|
||||
let parent = unsafe { parent.as_mut() };
|
||||
parent.unregister_child(token_state_pointer, current_state);
|
||||
}
|
||||
}
|
||||
tree_node::decrease_handle_refcount(&self.inner);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,29 +100,11 @@ impl Default for CancellationToken {
|
||||
impl CancellationToken {
|
||||
/// Creates a new CancellationToken in the non-cancelled state.
|
||||
pub fn new() -> CancellationToken {
|
||||
let state = Box::new(CancellationTokenState::new(
|
||||
None,
|
||||
StateSnapshot {
|
||||
cancel_state: CancellationState::NotCancelled,
|
||||
has_parent_ref: false,
|
||||
refcount: 1,
|
||||
},
|
||||
));
|
||||
|
||||
// Safety: We just created the Box. The pointer is guaranteed to be
|
||||
// not null
|
||||
CancellationToken {
|
||||
inner: unsafe { NonNull::new_unchecked(Box::into_raw(state)) },
|
||||
inner: Arc::new(tree_node::TreeNode::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a reference to the utilized `CancellationTokenState`.
|
||||
fn state(&self) -> &CancellationTokenState {
|
||||
// Safety: The state inside a `CancellationToken` is always valid, since
|
||||
// is reference counted
|
||||
unsafe { &*self.inner.as_ptr() }
|
||||
}
|
||||
|
||||
/// Creates a `CancellationToken` which will get cancelled whenever the
|
||||
/// current token gets cancelled.
|
||||
///
|
||||
@@ -203,56 +144,8 @@ impl CancellationToken {
|
||||
/// }
|
||||
/// ```
|
||||
pub fn child_token(&self) -> CancellationToken {
|
||||
let inner = self.state();
|
||||
|
||||
// Increment the refcount of this token. It will be referenced by the
|
||||
// child, independent of whether the child is immediately cancelled or
|
||||
// not.
|
||||
let _current_state = inner.increment_refcount(inner.snapshot());
|
||||
|
||||
let mut unpacked_child_state = StateSnapshot {
|
||||
has_parent_ref: true,
|
||||
refcount: 1,
|
||||
cancel_state: CancellationState::NotCancelled,
|
||||
};
|
||||
let mut child_token_state = Box::new(CancellationTokenState::new(
|
||||
Some(self.inner),
|
||||
unpacked_child_state,
|
||||
));
|
||||
|
||||
{
|
||||
let mut guard = inner.synchronized.lock().unwrap();
|
||||
if guard.is_cancelled {
|
||||
// This task was already cancelled. In this case we should not
|
||||
// insert the child into the list, since it would never get removed
|
||||
// from the list.
|
||||
(*child_token_state.synchronized.lock().unwrap()).is_cancelled = true;
|
||||
unpacked_child_state.cancel_state = CancellationState::Cancelled;
|
||||
// Since it's not in the list, the parent doesn't need to retain
|
||||
// a reference to it.
|
||||
unpacked_child_state.has_parent_ref = false;
|
||||
child_token_state
|
||||
.state
|
||||
.store(unpacked_child_state.pack(), Ordering::SeqCst);
|
||||
} else {
|
||||
if let Some(mut first_child) = guard.first_child {
|
||||
child_token_state.from_parent.next_peer = Some(first_child);
|
||||
// Safety: We manipulate other child task inside the Mutex
|
||||
// and retain a parent reference on it. The child token can't
|
||||
// get invalidated while the Mutex is held.
|
||||
unsafe {
|
||||
first_child.as_mut().from_parent.prev_peer =
|
||||
Some((&mut *child_token_state).into())
|
||||
};
|
||||
}
|
||||
guard.first_child = Some((&mut *child_token_state).into());
|
||||
}
|
||||
};
|
||||
|
||||
let child_token_ptr = Box::into_raw(child_token_state);
|
||||
// Safety: We just created the pointer from a `Box`
|
||||
CancellationToken {
|
||||
inner: unsafe { NonNull::new_unchecked(child_token_ptr) },
|
||||
inner: tree_node::child_node(&self.inner),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -260,21 +153,33 @@ impl CancellationToken {
|
||||
/// derived from it.
|
||||
///
|
||||
/// This will wake up all tasks which are waiting for cancellation.
|
||||
///
|
||||
/// Be aware that cancellation is not an atomic operation. It is possible
|
||||
/// for another thread running in parallel with a call to `cancel` to first
|
||||
/// receive `true` from `is_cancelled` on one child node, and then receive
|
||||
/// `false` from `is_cancelled` on another child node. However, once the
|
||||
/// call to `cancel` returns, all child nodes have been fully cancelled.
|
||||
pub fn cancel(&self) {
|
||||
self.state().cancel();
|
||||
tree_node::cancel(&self.inner);
|
||||
}
|
||||
|
||||
/// Returns `true` if the `CancellationToken` had been cancelled
|
||||
/// Returns `true` if the `CancellationToken` is cancelled.
|
||||
pub fn is_cancelled(&self) -> bool {
|
||||
self.state().is_cancelled()
|
||||
tree_node::is_cancelled(&self.inner)
|
||||
}
|
||||
|
||||
/// Returns a `Future` that gets fulfilled when cancellation is requested.
|
||||
///
|
||||
/// The future will complete immediately if the token is already cancelled
|
||||
/// when this method is called.
|
||||
///
|
||||
/// # Cancel safety
|
||||
///
|
||||
/// This method is cancel safe.
|
||||
pub fn cancelled(&self) -> WaitForCancellationFuture<'_> {
|
||||
WaitForCancellationFuture {
|
||||
cancellation_token: Some(self),
|
||||
wait_node: ListNode::new(WaitQueueEntry::new()),
|
||||
is_registered: false,
|
||||
cancellation_token: self,
|
||||
future: self.inner.notified(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -285,26 +190,6 @@ impl CancellationToken {
|
||||
pub fn drop_guard(self) -> DropGuard {
|
||||
DropGuard { inner: Some(self) }
|
||||
}
|
||||
|
||||
unsafe fn register(
|
||||
&self,
|
||||
wait_node: &mut ListNode<WaitQueueEntry>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<()> {
|
||||
self.state().register(wait_node, cx)
|
||||
}
|
||||
|
||||
fn check_for_cancellation(
|
||||
&self,
|
||||
wait_node: &mut ListNode<WaitQueueEntry>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<()> {
|
||||
self.state().check_for_cancellation(wait_node, cx)
|
||||
}
|
||||
|
||||
fn unregister(&self, wait_node: &mut ListNode<WaitQueueEntry>) {
|
||||
self.state().unregister(wait_node)
|
||||
}
|
||||
}
|
||||
|
||||
// ===== impl WaitForCancellationFuture =====
|
||||
@@ -319,560 +204,21 @@ impl<'a> Future for WaitForCancellationFuture<'a> {
|
||||
type Output = ();
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
|
||||
// Safety: We do not move anything out of `WaitForCancellationFuture`
|
||||
let mut_self: &mut WaitForCancellationFuture<'_> = unsafe { Pin::get_unchecked_mut(self) };
|
||||
|
||||
let cancellation_token = mut_self
|
||||
.cancellation_token
|
||||
.expect("polled WaitForCancellationFuture after completion");
|
||||
|
||||
let poll_res = if !mut_self.is_registered {
|
||||
// Safety: The `ListNode` is pinned through the Future,
|
||||
// and we will unregister it in `WaitForCancellationFuture::drop`
|
||||
// before the Future is dropped and the memory reference is invalidated.
|
||||
unsafe { cancellation_token.register(&mut mut_self.wait_node, cx) }
|
||||
} else {
|
||||
cancellation_token.check_for_cancellation(&mut mut_self.wait_node, cx)
|
||||
};
|
||||
|
||||
if let Poll::Ready(()) = poll_res {
|
||||
// The cancellation_token was signalled
|
||||
mut_self.cancellation_token = None;
|
||||
// A signalled Token means the Waker won't be enqueued anymore
|
||||
mut_self.is_registered = false;
|
||||
mut_self.wait_node.task = None;
|
||||
} else {
|
||||
// This `Future` and its stored `Waker` stay registered at the
|
||||
// `CancellationToken`
|
||||
mut_self.is_registered = true;
|
||||
}
|
||||
|
||||
poll_res
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> Drop for WaitForCancellationFuture<'a> {
|
||||
fn drop(&mut self) {
|
||||
// If this WaitForCancellationFuture has been polled and it was added to the
|
||||
// wait queue at the cancellation_token, it must be removed before dropping.
|
||||
// Otherwise the cancellation_token would access invalid memory.
|
||||
if let Some(token) = self.cancellation_token {
|
||||
if self.is_registered {
|
||||
token.unregister(&mut self.wait_node);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tracks how the future had interacted with the [`CancellationToken`]
|
||||
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
|
||||
enum PollState {
|
||||
/// The task has never interacted with the [`CancellationToken`].
|
||||
New,
|
||||
/// The task was added to the wait queue at the [`CancellationToken`].
|
||||
Waiting,
|
||||
/// The task has been polled to completion.
|
||||
Done,
|
||||
}
|
||||
|
||||
/// Tracks the WaitForCancellationFuture waiting state.
|
||||
/// Access to this struct is synchronized through the mutex in the CancellationToken.
|
||||
struct WaitQueueEntry {
|
||||
/// The task handle of the waiting task
|
||||
task: Option<Waker>,
|
||||
// Current polling state. This state is only updated inside the Mutex of
|
||||
// the CancellationToken.
|
||||
state: PollState,
|
||||
}
|
||||
|
||||
impl WaitQueueEntry {
|
||||
/// Creates a new WaitQueueEntry
|
||||
fn new() -> WaitQueueEntry {
|
||||
WaitQueueEntry {
|
||||
task: None,
|
||||
state: PollState::New,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SynchronizedState {
|
||||
waiters: LinkedList<WaitQueueEntry>,
|
||||
first_child: Option<NonNull<CancellationTokenState>>,
|
||||
is_cancelled: bool,
|
||||
}
|
||||
|
||||
impl SynchronizedState {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
waiters: LinkedList::new(),
|
||||
first_child: None,
|
||||
is_cancelled: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Information embedded in child tokens which is synchronized through the Mutex
|
||||
/// in their parent.
|
||||
struct SynchronizedThroughParent {
|
||||
next_peer: Option<NonNull<CancellationTokenState>>,
|
||||
prev_peer: Option<NonNull<CancellationTokenState>>,
|
||||
}
|
||||
|
||||
/// Possible states of a `CancellationToken`
|
||||
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
|
||||
enum CancellationState {
|
||||
NotCancelled = 0,
|
||||
Cancelling = 1,
|
||||
Cancelled = 2,
|
||||
}
|
||||
|
||||
impl CancellationState {
|
||||
fn pack(self) -> usize {
|
||||
self as usize
|
||||
}
|
||||
|
||||
fn unpack(value: usize) -> Self {
|
||||
match value {
|
||||
0 => CancellationState::NotCancelled,
|
||||
1 => CancellationState::Cancelling,
|
||||
2 => CancellationState::Cancelled,
|
||||
_ => unreachable!("Invalid value"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
|
||||
struct StateSnapshot {
|
||||
/// The amount of references to this particular CancellationToken.
|
||||
/// `CancellationToken` structs hold these references to a `CancellationTokenState`.
|
||||
/// Also the state is referenced by the state of each child.
|
||||
refcount: usize,
|
||||
/// Whether the state is still referenced by it's parent and can therefore
|
||||
/// not be freed.
|
||||
has_parent_ref: bool,
|
||||
/// Whether the token is cancelled
|
||||
cancel_state: CancellationState,
|
||||
}
|
||||
|
||||
impl StateSnapshot {
|
||||
/// Packs the snapshot into a `usize`
|
||||
fn pack(self) -> usize {
|
||||
self.refcount << 3 | if self.has_parent_ref { 4 } else { 0 } | self.cancel_state.pack()
|
||||
}
|
||||
|
||||
/// Unpacks the snapshot from a `usize`
|
||||
fn unpack(value: usize) -> Self {
|
||||
let refcount = value >> 3;
|
||||
let has_parent_ref = value & 4 != 0;
|
||||
let cancel_state = CancellationState::unpack(value & 0x03);
|
||||
|
||||
StateSnapshot {
|
||||
refcount,
|
||||
has_parent_ref,
|
||||
cancel_state,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether this `CancellationTokenState` is still referenced by any
|
||||
/// `CancellationToken`.
|
||||
fn has_refs(&self) -> bool {
|
||||
self.refcount != 0 || self.has_parent_ref
|
||||
}
|
||||
}
|
||||
|
||||
/// The maximum permitted amount of references to a CancellationToken. This
|
||||
/// is derived from the intent to never use more than 32bit in the `Snapshot`.
|
||||
const MAX_REFS: u32 = (std::u32::MAX - 7) >> 3;
|
||||
|
||||
/// Internal state of the `CancellationToken` pair above
|
||||
struct CancellationTokenState {
|
||||
state: AtomicUsize,
|
||||
parent: Option<NonNull<CancellationTokenState>>,
|
||||
from_parent: SynchronizedThroughParent,
|
||||
synchronized: Mutex<SynchronizedState>,
|
||||
}
|
||||
|
||||
impl CancellationTokenState {
|
||||
fn new(
|
||||
parent: Option<NonNull<CancellationTokenState>>,
|
||||
state: StateSnapshot,
|
||||
) -> CancellationTokenState {
|
||||
CancellationTokenState {
|
||||
parent,
|
||||
from_parent: SynchronizedThroughParent {
|
||||
prev_peer: None,
|
||||
next_peer: None,
|
||||
},
|
||||
state: AtomicUsize::new(state.pack()),
|
||||
synchronized: Mutex::new(SynchronizedState::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a snapshot of the current atomic state of the token
|
||||
fn snapshot(&self) -> StateSnapshot {
|
||||
StateSnapshot::unpack(self.state.load(Ordering::SeqCst))
|
||||
}
|
||||
|
||||
fn atomic_update_state<F>(&self, mut current_state: StateSnapshot, func: F) -> StateSnapshot
|
||||
where
|
||||
F: Fn(StateSnapshot) -> StateSnapshot,
|
||||
{
|
||||
let mut current_packed_state = current_state.pack();
|
||||
let mut this = self.project();
|
||||
loop {
|
||||
let next_state = func(current_state);
|
||||
match self.state.compare_exchange(
|
||||
current_packed_state,
|
||||
next_state.pack(),
|
||||
Ordering::SeqCst,
|
||||
Ordering::SeqCst,
|
||||
) {
|
||||
Ok(_) => {
|
||||
return next_state;
|
||||
}
|
||||
Err(actual) => {
|
||||
current_packed_state = actual;
|
||||
current_state = StateSnapshot::unpack(actual);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn increment_refcount(&self, current_state: StateSnapshot) -> StateSnapshot {
|
||||
self.atomic_update_state(current_state, |mut state: StateSnapshot| {
|
||||
if state.refcount >= MAX_REFS as usize {
|
||||
eprintln!("[ERROR] Maximum reference count for CancellationToken was exceeded");
|
||||
std::process::abort();
|
||||
}
|
||||
state.refcount += 1;
|
||||
state
|
||||
})
|
||||
}
|
||||
|
||||
fn decrement_refcount(&self, current_state: StateSnapshot) -> StateSnapshot {
|
||||
let current_state = self.atomic_update_state(current_state, |mut state: StateSnapshot| {
|
||||
state.refcount -= 1;
|
||||
state
|
||||
});
|
||||
|
||||
// Drop the State if it is not referenced anymore
|
||||
if !current_state.has_refs() {
|
||||
// Safety: `CancellationTokenState` is always stored in refcounted
|
||||
// Boxes
|
||||
let _ = unsafe { Box::from_raw(self as *const Self as *mut Self) };
|
||||
}
|
||||
|
||||
current_state
|
||||
}
|
||||
|
||||
fn remove_parent_ref(&self, current_state: StateSnapshot) -> StateSnapshot {
|
||||
let current_state = self.atomic_update_state(current_state, |mut state: StateSnapshot| {
|
||||
state.has_parent_ref = false;
|
||||
state
|
||||
});
|
||||
|
||||
// Drop the State if it is not referenced anymore
|
||||
if !current_state.has_refs() {
|
||||
// Safety: `CancellationTokenState` is always stored in refcounted
|
||||
// Boxes
|
||||
let _ = unsafe { Box::from_raw(self as *const Self as *mut Self) };
|
||||
}
|
||||
|
||||
current_state
|
||||
}
|
||||
|
||||
/// Unregisters a child from the parent token.
|
||||
/// The child tokens state is not exactly known at this point in time.
|
||||
/// If the parent token is cancelled, the child token gets removed from the
|
||||
/// parents list, and might therefore already have been freed. If the parent
|
||||
/// token is not cancelled, the child token is still valid.
|
||||
fn unregister_child(
|
||||
&mut self,
|
||||
mut child_state: NonNull<CancellationTokenState>,
|
||||
current_child_state: StateSnapshot,
|
||||
) {
|
||||
let removed_child = {
|
||||
// Remove the child toke from the parents linked list
|
||||
let mut guard = self.synchronized.lock().unwrap();
|
||||
if !guard.is_cancelled {
|
||||
// Safety: Since the token was not cancelled, the child must
|
||||
// still be in the list and valid.
|
||||
let mut child_state = unsafe { child_state.as_mut() };
|
||||
debug_assert!(child_state.snapshot().has_parent_ref);
|
||||
|
||||
if guard.first_child == Some(child_state.into()) {
|
||||
guard.first_child = child_state.from_parent.next_peer;
|
||||
}
|
||||
// Safety: If peers wouldn't be valid anymore, they would try
|
||||
// to remove themselves from the list. This would require locking
|
||||
// the Mutex that we currently own.
|
||||
unsafe {
|
||||
if let Some(mut prev_peer) = child_state.from_parent.prev_peer {
|
||||
prev_peer.as_mut().from_parent.next_peer =
|
||||
child_state.from_parent.next_peer;
|
||||
}
|
||||
if let Some(mut next_peer) = child_state.from_parent.next_peer {
|
||||
next_peer.as_mut().from_parent.prev_peer =
|
||||
child_state.from_parent.prev_peer;
|
||||
}
|
||||
}
|
||||
child_state.from_parent.prev_peer = None;
|
||||
child_state.from_parent.next_peer = None;
|
||||
|
||||
// The child is no longer referenced by the parent, since we were able
|
||||
// to remove its reference from the parents list.
|
||||
true
|
||||
} else {
|
||||
// Do not touch the linked list anymore. If the parent is cancelled
|
||||
// it will move all childs outside of the Mutex and manipulate
|
||||
// the pointers there. Manipulating the pointers here too could
|
||||
// lead to races. Therefore leave them just as as and let the
|
||||
// parent deal with it. The parent will make sure to retain a
|
||||
// reference to this state as long as it manipulates the list
|
||||
// pointers. Therefore the pointers are not dangling.
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
if removed_child {
|
||||
// If the token removed itself from the parents list, it can reset
|
||||
// the parent ref status. If it is isn't able to do so, because the
|
||||
// parent removed it from the list, there is no need to do this.
|
||||
// The parent ref acts as as another reference count. Therefore
|
||||
// removing this reference can free the object.
|
||||
// Safety: The token was in the list. This means the parent wasn't
|
||||
// cancelled before, and the token must still be alive.
|
||||
unsafe { child_state.as_mut().remove_parent_ref(current_child_state) };
|
||||
}
|
||||
|
||||
// Decrement the refcount on the parent and free it if necessary
|
||||
self.decrement_refcount(self.snapshot());
|
||||
}
|
||||
|
||||
fn cancel(&self) {
|
||||
// Move the state of the CancellationToken from `NotCancelled` to `Cancelling`
|
||||
let mut current_state = self.snapshot();
|
||||
|
||||
let state_after_cancellation = loop {
|
||||
if current_state.cancel_state != CancellationState::NotCancelled {
|
||||
// Another task already initiated the cancellation
|
||||
return;
|
||||
if this.cancellation_token.is_cancelled() {
|
||||
return Poll::Ready(());
|
||||
}
|
||||
|
||||
let mut next_state = current_state;
|
||||
next_state.cancel_state = CancellationState::Cancelling;
|
||||
match self.state.compare_exchange(
|
||||
current_state.pack(),
|
||||
next_state.pack(),
|
||||
Ordering::SeqCst,
|
||||
Ordering::SeqCst,
|
||||
) {
|
||||
Ok(_) => break next_state,
|
||||
Err(actual) => current_state = StateSnapshot::unpack(actual),
|
||||
// No wakeups can be lost here because there is always a call to
|
||||
// `is_cancelled` between the creation of the future and the call to
|
||||
// `poll`, and the code that sets the cancelled flag does so before
|
||||
// waking the `Notified`.
|
||||
if this.future.as_mut().poll(cx).is_pending() {
|
||||
return Poll::Pending;
|
||||
}
|
||||
};
|
||||
|
||||
// This task cancelled the token
|
||||
|
||||
// Take the task list out of the Token
|
||||
// We do not want to cancel child token inside this lock. If one of the
|
||||
// child tasks would have additional child tokens, we would recursively
|
||||
// take locks.
|
||||
|
||||
// Doing this action has an impact if the child token is dropped concurrently:
|
||||
// It will try to deregister itself from the parent task, but can not find
|
||||
// itself in the task list anymore. Therefore it needs to assume the parent
|
||||
// has extracted the list and will process it. It may not modify the list.
|
||||
// This is OK from a memory safety perspective, since the parent still
|
||||
// retains a reference to the child task until it finished iterating over
|
||||
// it.
|
||||
|
||||
let mut first_child = {
|
||||
let mut guard = self.synchronized.lock().unwrap();
|
||||
// Save the cancellation also inside the Mutex
|
||||
// This allows child tokens which want to detach themselves to detect
|
||||
// that this is no longer required since the parent cleared the list.
|
||||
guard.is_cancelled = true;
|
||||
|
||||
// Wakeup all waiters
|
||||
// This happens inside the lock to make cancellation reliable
|
||||
// If we would access waiters outside of the lock, the pointers
|
||||
// may no longer be valid.
|
||||
// Typically this shouldn't be an issue, since waking a task should
|
||||
// only move it from the blocked into the ready state and not have
|
||||
// further side effects.
|
||||
|
||||
// Use a reverse iterator, so that the oldest waiter gets
|
||||
// scheduled first
|
||||
guard.waiters.reverse_drain(|waiter| {
|
||||
// We are not allowed to move the `Waker` out of the list node.
|
||||
// The `Future` relies on the fact that the old `Waker` stays there
|
||||
// as long as the `Future` has not completed in order to perform
|
||||
// the `will_wake()` check.
|
||||
// Therefore `wake_by_ref` is used instead of `wake()`
|
||||
if let Some(handle) = &mut waiter.task {
|
||||
handle.wake_by_ref();
|
||||
}
|
||||
// Mark the waiter to have been removed from the list.
|
||||
waiter.state = PollState::Done;
|
||||
});
|
||||
|
||||
guard.first_child.take()
|
||||
};
|
||||
|
||||
while let Some(mut child) = first_child {
|
||||
// Safety: We know this is a valid pointer since it is in our child pointer
|
||||
// list. It can't have been freed in between, since we retain a a reference
|
||||
// to each child.
|
||||
let mut_child = unsafe { child.as_mut() };
|
||||
|
||||
// Get the next child and clean up list pointers
|
||||
first_child = mut_child.from_parent.next_peer;
|
||||
mut_child.from_parent.prev_peer = None;
|
||||
mut_child.from_parent.next_peer = None;
|
||||
|
||||
// Cancel the child task
|
||||
mut_child.cancel();
|
||||
|
||||
// Drop the parent reference. This `CancellationToken` is not interested
|
||||
// in interacting with the child anymore.
|
||||
// This is ONLY allowed once we promised not to touch the state anymore
|
||||
// after this interaction.
|
||||
mut_child.remove_parent_ref(mut_child.snapshot());
|
||||
this.future.set(this.cancellation_token.inner.notified());
|
||||
}
|
||||
|
||||
// The cancellation has completed
|
||||
// At this point in time tasks which registered a wait node can be sure
|
||||
// that this wait node already had been dequeued from the list without
|
||||
// needing to inspect the list.
|
||||
self.atomic_update_state(state_after_cancellation, |mut state| {
|
||||
state.cancel_state = CancellationState::Cancelled;
|
||||
state
|
||||
});
|
||||
}
|
||||
|
||||
/// Returns `true` if the `CancellationToken` had been cancelled
|
||||
fn is_cancelled(&self) -> bool {
|
||||
let current_state = self.snapshot();
|
||||
current_state.cancel_state != CancellationState::NotCancelled
|
||||
}
|
||||
|
||||
/// Registers a waiting task at the `CancellationToken`.
|
||||
/// Safety: This method is only safe as long as the waiting waiting task
|
||||
/// will properly unregister the wait node before it gets moved.
|
||||
unsafe fn register(
|
||||
&self,
|
||||
wait_node: &mut ListNode<WaitQueueEntry>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<()> {
|
||||
debug_assert_eq!(PollState::New, wait_node.state);
|
||||
let current_state = self.snapshot();
|
||||
|
||||
// Perform an optimistic cancellation check before. This is not strictly
|
||||
// necessary since we also check for cancellation in the Mutex, but
|
||||
// reduces the necessary work to be performed for tasks which already
|
||||
// had been cancelled.
|
||||
if current_state.cancel_state != CancellationState::NotCancelled {
|
||||
return Poll::Ready(());
|
||||
}
|
||||
|
||||
// So far the token is not cancelled. However it could be cancelled before
|
||||
// we get the chance to store the `Waker`. Therefore we need to check
|
||||
// for cancellation again inside the mutex.
|
||||
let mut guard = self.synchronized.lock().unwrap();
|
||||
if guard.is_cancelled {
|
||||
// Cancellation was signalled
|
||||
wait_node.state = PollState::Done;
|
||||
Poll::Ready(())
|
||||
} else {
|
||||
// Added the task to the wait queue
|
||||
wait_node.task = Some(cx.waker().clone());
|
||||
wait_node.state = PollState::Waiting;
|
||||
guard.waiters.add_front(wait_node);
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
fn check_for_cancellation(
|
||||
&self,
|
||||
wait_node: &mut ListNode<WaitQueueEntry>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<()> {
|
||||
debug_assert!(
|
||||
wait_node.task.is_some(),
|
||||
"Method can only be called after task had been registered"
|
||||
);
|
||||
|
||||
let current_state = self.snapshot();
|
||||
|
||||
if current_state.cancel_state != CancellationState::NotCancelled {
|
||||
// If the cancellation had been fully completed we know that our `Waker`
|
||||
// is no longer registered at the `CancellationToken`.
|
||||
// Otherwise the cancel call may or may not yet have iterated
|
||||
// through the waiters list and removed the wait nodes.
|
||||
// If it hasn't yet, we need to remove it. Otherwise an attempt to
|
||||
// reuse the `wait_node´ might get freed due to the `WaitForCancellationFuture`
|
||||
// getting dropped before the cancellation had interacted with it.
|
||||
if current_state.cancel_state != CancellationState::Cancelled {
|
||||
self.unregister(wait_node);
|
||||
}
|
||||
Poll::Ready(())
|
||||
} else {
|
||||
// Check if we need to swap the `Waker`. This will make the check more
|
||||
// expensive, since the `Waker` is synchronized through the Mutex.
|
||||
// If we don't need to perform a `Waker` update, an atomic check for
|
||||
// cancellation is sufficient.
|
||||
let need_waker_update = wait_node
|
||||
.task
|
||||
.as_ref()
|
||||
.map(|waker| !waker.will_wake(cx.waker()))
|
||||
.unwrap_or(true);
|
||||
|
||||
if need_waker_update {
|
||||
let guard = self.synchronized.lock().unwrap();
|
||||
if guard.is_cancelled {
|
||||
// Cancellation was signalled. Since this cancellation signal
|
||||
// is set inside the Mutex, the old waiter must already have
|
||||
// been removed from the waiting list
|
||||
debug_assert_eq!(PollState::Done, wait_node.state);
|
||||
wait_node.task = None;
|
||||
Poll::Ready(())
|
||||
} else {
|
||||
// The WaitForCancellationFuture is already in the queue.
|
||||
// The CancellationToken can't have been cancelled,
|
||||
// since this would change the is_cancelled flag inside the mutex.
|
||||
// Therefore we just have to update the Waker. A follow-up
|
||||
// cancellation will always use the new waker.
|
||||
wait_node.task = Some(cx.waker().clone());
|
||||
Poll::Pending
|
||||
}
|
||||
} else {
|
||||
// Do nothing. If the token gets cancelled, this task will get
|
||||
// woken again and can fetch the cancellation.
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn unregister(&self, wait_node: &mut ListNode<WaitQueueEntry>) {
|
||||
debug_assert!(
|
||||
wait_node.task.is_some(),
|
||||
"waiter can not be active without task"
|
||||
);
|
||||
|
||||
let mut guard = self.synchronized.lock().unwrap();
|
||||
// WaitForCancellationFuture only needs to get removed if it has been added to
|
||||
// the wait queue of the CancellationToken.
|
||||
// This has happened in the PollState::Waiting case.
|
||||
if let PollState::Waiting = wait_node.state {
|
||||
// Safety: Due to the state, we know that the node must be part
|
||||
// of the waiter list
|
||||
if !unsafe { guard.waiters.remove(wait_node) } {
|
||||
// Panic if the address isn't found. This can only happen if the contract was
|
||||
// violated, e.g. the WaitQueueEntry got moved after the initial poll.
|
||||
panic!("Future could not be removed from wait queue");
|
||||
}
|
||||
wait_node.state = PollState::Done;
|
||||
}
|
||||
wait_node.task = None;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user