A pit full of rusty nails
Something went wrong. Try again.
15 kB · 368 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369//! This mod provides the logic for the inner tree structure of the `CancelToken`.//!//! `CancelToken`s are only light handles with references to [`TreeNode`].//! All the logic is actually implemented in the [`TreeNode`].//!//! A [`TreeNode`] is part of the cancellation tree and may have one parent and an arbitrary number of//! children.//!//! A [`TreeNode`] can receive the request to perform a cancellation through a `CancelToken`.//! This cancellation request will cancel the node and all of its descendants.//!//! As soon as a node cannot get cancelled any more (because it was already cancelled or it has no//! more `CancelTokens` pointing to it any more), it gets removed from the tree, to keep the//! tree as small as possible.//!//! # Invariants//!//! Those invariants shall be true at any time.//!//! 1. A node that has no parents and no handles can no longer be cancelled.//! This is important during both cancellation and refcounting.//!//! 2. If node B *is* or *was* a child of node A, then node B was created *after* node A.//! This is important for deadlock safety, as it is used for lock order.//! Node B can only become the child of node A in two ways://! - being created with `child_node()`, in which case it is trivially true that//! node A already existed when node B was created//! - being moved A->C->B to A->B because node C was removed in `decrease_handle_refcount()`//! or `cancel()`. In this case the invariant still holds, as B was younger than C, and C//! was younger than A, therefore B is also younger than A.//!//! 3. If two nodes are both unlocked and node A is the parent of node B, then node B is a child of//! node A. It is important to always restore that invariant before dropping the lock of a node.//!//! # Deadlock safety//!//! We always lock in the order of creation time. We can prove this through invariant #2.//! Specifically, through invariant #2, we know that we always have to lock a parent//! before its child.//!use parking_lot::{Mutex, MutexGuard};use std::sync::Arc;
use event_listener::{Event, EventListener};
/// A node of the cancellation tree structure////// The actual data it holds is wrapped inside a mutex for synchronization.pub(crate) struct TreeNode { inner: Mutex<Inner>, waker: Event,}
impl TreeNode { pub(crate) const fn new() -> Self { Self { inner: Mutex::new(Inner { parent: None, parent_idx: 0, children: Vec::new(), is_cancelled: false, num_handles: 1, }), waker: Event::new(), } }
pub(crate) fn notified(&self) -> EventListener<()> { self.waker.listen() }
/// Returns whether or not the node is cancelled pub(crate) fn is_cancelled(self: &Arc<Self>) -> bool { self.inner.lock().is_cancelled }
/// Creates a child node pub(crate) fn child_node(self: &Arc<Self>) -> Arc<Self> { let mut locked_parent = self.inner.lock();
// Do not register as child if we are already cancelled. // Cancelled trees can never be uncancelled and therefore // need no connection to parents or children any more. if locked_parent.is_cancelled { return Arc::new(TreeNode { inner: Mutex::new(Inner { parent: None, parent_idx: 0, children: Vec::new(), is_cancelled: true, num_handles: 1, }), waker: Event::new(), }); }
let child = Arc::new(TreeNode { inner: Mutex::new(Inner { parent: Some(self.clone()), parent_idx: locked_parent.children.len(), children: Vec::new(), is_cancelled: false, num_handles: 1, }), waker: Event::new(), });
locked_parent.children.push(child.clone());
child }
/// Increases the reference count of handles. pub(crate) fn increase_handle_refcount(self: &Arc<Self>) { let mut locked_node = self.inner.lock();
// Once no handles are left over, the node gets detached from the tree. // There should never be a new handle once all handles are dropped. assert!(locked_node.num_handles > 0);
locked_node.num_handles += 1; }
/// Decreases the reference count of handles. /// /// Once no handle is left, we can remove the node from the /// tree and connect its parent directly to its children. pub(crate) fn decrease_handle_refcount(self: &Arc<Self>) { let num_handles = { let mut locked_node = self.inner.lock(); locked_node.num_handles -= 1; locked_node.num_handles };
if num_handles == 0 { self.with_locked_node_and_parent(|mut node, parent| { // Remove the node from the tree match parent { Some(mut parent) => { // As we want to remove ourselves from the tree, // we have to move the children to the parent, so that // they still receive the cancellation event without us. // Moving them does not violate invariant #1. node.move_children_to_parent(&mut parent);
// Remove the node from the parent parent.remove_child(node); } None => { // Due to invariant #1, we can assume that our // children can no longer be cancelled through us. // (as we now have neither a parent nor handles) // Therefore we can disconnect them. node.disconnect_children(); } } }); } }
/// Cancels a node and its children. pub(crate) fn cancel(self: &Arc<Self>) { let mut locked_node = self.inner.lock();
if locked_node.is_cancelled { return; }
// One by one, adopt grandchildren and then cancel and detach the child while let Some(child) = locked_node.children.pop() { // This can't deadlock because the mutex we are already // holding is the parent of child. let mut locked_child = child.inner.lock();
// Detach the child from node // No need to modify node.children, as the child already got removed with `.pop` locked_child.parent = None; locked_child.parent_idx = 0;
// If child is already cancelled, detaching is enough if locked_child.is_cancelled { continue; }
// Cancel or adopt grandchildren while let Some(grandchild) = locked_child.children.pop() { // This can't deadlock because the two mutexes we are already // holding is the parent and grandparent of grandchild. let mut locked_grandchild = grandchild.inner.lock();
// Detach the grandchild locked_grandchild.parent = None; locked_grandchild.parent_idx = 0;
// If grandchild is already cancelled, detaching is enough if locked_grandchild.is_cancelled { continue; }
// For performance reasons, only adopt grandchildren that have children. // Otherwise, just cancel them right away, no need for another iteration. if locked_grandchild.children.is_empty() { // Cancel the grandchild locked_grandchild.is_cancelled = true; locked_grandchild.children = Vec::new(); drop(locked_grandchild); grandchild.waker.notify(usize::MAX); } else { // Otherwise, adopt grandchild locked_grandchild.parent = Some(self.clone()); locked_grandchild.parent_idx = locked_node.children.len(); drop(locked_grandchild); locked_node.children.push(grandchild); } }
// Cancel the child locked_child.is_cancelled = true; locked_child.children = Vec::new(); drop(locked_child); child.waker.notify(usize::MAX);
// Now the child is cancelled and detached and all its children are adopted. // Just continue until all (including adopted) children are cancelled and detached. }
// Cancel the node itself. locked_node.is_cancelled = true; locked_node.children = Vec::new(); drop(locked_node); self.waker.notify(usize::MAX); }
/// Figures out the parent of the node and locks the node and its parent atomically. /// /// The basic principle of preventing deadlocks in the tree is /// that we always lock the parent first, and then the child. /// For more info look at *deadlock safety* and *invariant #2*. /// /// Sadly, it's impossible to figure out the parent of a node without /// locking it. To then achieve locking order consistency, the node /// has to be unlocked before the parent gets locked. /// This leaves a small window where we already assume that we know the parent, /// but neither the parent nor the node is locked. Therefore, the parent could change. /// /// To prevent that this problem leaks into the rest of the code, it is abstracted /// in this function. /// /// The locked child and optionally its locked parent, if a parent exists, get passed /// to the `func` argument via (node, None) or (node, Some(parent)). fn with_locked_node_and_parent<F, Ret>(self: &Arc<Self>, func: F) -> Ret where F: FnOnce(MutexGuard<'_, Inner>, Option<MutexGuard<'_, Inner>>) -> Ret, { let mut locked_node = self.inner.lock();
// Every time this fails, the number of ancestors of the node decreases, // so the loop must succeed after a finite number of iterations. loop { // Look up the parent of the currently locked node. let potential_parent = match locked_node.parent.as_ref() { Some(potential_parent) => potential_parent.clone(), None => return func(locked_node, None), };
// Lock the parent. This may require unlocking the child first. let locked_parent = match potential_parent.inner.try_lock() { Some(locked_parent) => locked_parent, None => { drop(locked_node); // Deadlock safety: // // Due to invariant #2, the potential parent must come before // the child in the creation order. Therefore, we can safely // lock the child while holding the parent lock. let locked_parent = potential_parent.inner.lock(); locked_node = self.inner.lock(); locked_parent } };
// If we unlocked the child, then the parent may have changed. Check // that we still have the right parent. if let Some(actual_parent) = locked_node.parent.as_ref() && Arc::ptr_eq(actual_parent, &potential_parent) { return func(locked_node, Some(locked_parent)); } } }}
/// The data contained inside a `TreeNode`.////// This struct exists so that the data of the node can be wrapped/// in a Mutex.struct Inner { parent: Option<Arc<TreeNode>>, parent_idx: usize, children: Vec<Arc<TreeNode>>, is_cancelled: bool, num_handles: usize,}
impl Inner { /// Disconnects the given parent from all of its children. /// /// Takes a reference to [`Inner`] to make sure the parent is already locked. fn disconnect_children(&mut self) { for child in std::mem::take(&mut self.children) { let mut locked_child = child.inner.lock(); locked_child.parent_idx = 0; locked_child.parent = None; } }
/// Moves all children from `node` to `parent`. /// /// `parent` MUST have been a parent of the node when they both got locked, /// otherwise there is a potential for a deadlock as invariant #2 would be violated. /// /// To acquire the locks for node and parent, use [`with_locked_node_and_parent`]. fn move_children_to_parent(&mut self, parent: &mut Inner) { // Pre-allocate in the parent, for performance parent.children.reserve(self.children.len());
for child in std::mem::take(&mut self.children) { { let mut child_locked = child.inner.lock(); child_locked.parent.clone_from(&self.parent); child_locked.parent_idx = parent.children.len(); } parent.children.push(child); } }
/// Removes a child from the parent. /// /// `parent` MUST be the parent of `node`. /// To acquire the locks for node and parent, use [`with_locked_node_and_parent`]. fn remove_child(&mut self, mut node: MutexGuard<'_, Inner>) { // Query the position from where to remove a node let pos = node.parent_idx; node.parent = None; node.parent_idx = 0;
// Unlock node, so that only one child at a time is locked. // Otherwise we would violate the lock order (see 'deadlock safety') as we // don't know the creation order of the child nodes drop(node);
// If `node` is the last element in the list, we don't need any swapping if self.children.len() == pos + 1 { self.children.pop(); } else { // If `node` is not the last element in the list, we need to // replace it with the last element let replacement_child = self.children.pop().unwrap(); replacement_child.inner.lock().parent_idx = pos; self.children[pos] = replacement_child; }
let len = self.children.len(); if 4 * len <= self.children.capacity() { self.children.shrink_to(2 * len); } }}