1//! This mod provides the logic for the inner tree structure of the `CancellationToken`.
2//!
3//! `CancellationTokens` are only light handles with references to [`TreeNode`].
4//! All the logic is actually implemented in the [`TreeNode`].
5//!
6//! A [`TreeNode`] is part of the cancellation tree and may have one parent and an arbitrary number of
7//! children.
8//!
9//! A [`TreeNode`] can receive the request to perform a cancellation through a `CancellationToken`.
10//! This cancellation request will cancel the node and all of its descendants.
11//!
12//! As soon as a node cannot get cancelled any more (because it was already cancelled or it has no
13//! more `CancellationTokens` pointing to it any more), it gets removed from the tree, to keep the
14//! tree as small as possible.
15//!
16//! # Invariants
17//!
18//! Those invariants shall be true at any time.
19//!
20//! 1. A node that has no parents and no handles can no longer be cancelled.
21//! This is important during both cancellation and refcounting.
22//!
23//! 2. If node B *is* or *was* a child of node A, then node B was created *after* node A.
24//! This is important for deadlock safety, as it is used for lock order.
25//! Node B can only become the child of node A in two ways:
26//! - being created with `child_node()`, in which case it is trivially true that
27//! node A already existed when node B was created
28//! - being moved A->C->B to A->B because node C was removed in `decrease_handle_refcount()`
29//! or `cancel()`. In this case the invariant still holds, as B was younger than C, and C
30//! was younger than A, therefore B is also younger than A.
31//!
32//! 3. If two nodes are both unlocked and node A is the parent of node B, then node B is a child of
33//! node A. It is important to always restore that invariant before dropping the lock of a node.
34//!
35//! # Deadlock safety
36//!
37//! We always lock in the order of creation time. We can prove this through invariant #2.
38//! Specifically, through invariant #2, we know that we always have to lock a parent
39//! before its child.
40//!
41use crate::loom::sync::{Arc, Mutex, MutexGuard};
42
43/// A node of the cancellation tree structure
44///
45/// The actual data it holds is wrapped inside a mutex for synchronization.
46pub(crate) struct TreeNode {
47 inner: Mutex<Inner>,
48 waker: tokio::sync::Notify,
49}
50impl TreeNode {
51 pub(crate) fn new() -> Self {
52 Self {
53 inner: Mutex::new(Inner {
54 parent: None,
55 parent_idx: 0,
56 children: vec![],
57 is_cancelled: false,
58 num_handles: 1,
59 }),
60 waker: tokio::sync::Notify::new(),
61 }
62 }
63
64 pub(crate) fn notified(&self) -> tokio::sync::futures::Notified<'_> {
65 self.waker.notified()
66 }
67}
68
69/// The data contained inside a `TreeNode`.
70///
71/// This struct exists so that the data of the node can be wrapped
72/// in a Mutex.
73struct Inner {
74 parent: Option<Arc<TreeNode>>,
75 parent_idx: usize,
76 children: Vec<Arc<TreeNode>>,
77 is_cancelled: bool,
78 num_handles: usize,
79}
80
81/// Returns whether or not the node is cancelled
82pub(crate) fn is_cancelled(node: &Arc<TreeNode>) -> bool {
83 node.inner.lock().unwrap().is_cancelled
84}
85
86/// Creates a child node
87pub(crate) fn child_node(parent: &Arc<TreeNode>) -> Arc<TreeNode> {
88 let mut locked_parent = parent.inner.lock().unwrap();
89
90 // Do not register as child if we are already cancelled.
91 // Cancelled trees can never be uncancelled and therefore
92 // need no connection to parents or children any more.
93 if locked_parent.is_cancelled {
94 return Arc::new(TreeNode {
95 inner: Mutex::new(Inner {
96 parent: None,
97 parent_idx: 0,
98 children: vec![],
99 is_cancelled: true,
100 num_handles: 1,
101 }),
102 waker: tokio::sync::Notify::new(),
103 });
104 }
105
106 let child = Arc::new(TreeNode {
107 inner: Mutex::new(Inner {
108 parent: Some(parent.clone()),
109 parent_idx: locked_parent.children.len(),
110 children: vec![],
111 is_cancelled: false,
112 num_handles: 1,
113 }),
114 waker: tokio::sync::Notify::new(),
115 });
116
117 locked_parent.children.push(child.clone());
118
119 child
120}
121
122/// Disconnects the given parent from all of its children.
123///
124/// Takes a reference to [Inner] to make sure the parent is already locked.
125fn disconnect_children(node: &mut Inner) {
126 for child: Arc in std::mem::take(&mut node.children) {
127 let mut locked_child: MutexGuard<'_, Inner> = child.inner.lock().unwrap();
128 locked_child.parent_idx = 0;
129 locked_child.parent = None;
130 }
131}
132
133/// Figures out the parent of the node and locks the node and its parent atomically.
134///
135/// The basic principle of preventing deadlocks in the tree is
136/// that we always lock the parent first, and then the child.
137/// For more info look at *deadlock safety* and *invariant #2*.
138///
139/// Sadly, it's impossible to figure out the parent of a node without
140/// locking it. To then achieve locking order consistency, the node
141/// has to be unlocked before the parent gets locked.
142/// This leaves a small window where we already assume that we know the parent,
143/// but neither the parent nor the node is locked. Therefore, the parent could change.
144///
145/// To prevent that this problem leaks into the rest of the code, it is abstracted
146/// in this function.
147///
148/// The locked child and optionally its locked parent, if a parent exists, get passed
149/// to the `func` argument via (node, None) or (node, Some(parent)).
150fn with_locked_node_and_parent<F, Ret>(node: &Arc<TreeNode>, func: F) -> Ret
151where
152 F: FnOnce(MutexGuard<'_, Inner>, Option<MutexGuard<'_, Inner>>) -> Ret,
153{
154 use std::sync::TryLockError;
155
156 let mut locked_node = node.inner.lock().unwrap();
157
158 // Every time this fails, the number of ancestors of the node decreases,
159 // so the loop must succeed after a finite number of iterations.
160 loop {
161 // Look up the parent of the currently locked node.
162 let potential_parent = match locked_node.parent.as_ref() {
163 Some(potential_parent) => potential_parent.clone(),
164 None => return func(locked_node, None),
165 };
166
167 // Lock the parent. This may require unlocking the child first.
168 let locked_parent = match potential_parent.inner.try_lock() {
169 Ok(locked_parent) => locked_parent,
170 Err(TryLockError::WouldBlock) => {
171 drop(locked_node);
172 // Deadlock safety:
173 //
174 // Due to invariant #2, the potential parent must come before
175 // the child in the creation order. Therefore, we can safely
176 // lock the child while holding the parent lock.
177 let locked_parent = potential_parent.inner.lock().unwrap();
178 locked_node = node.inner.lock().unwrap();
179 locked_parent
180 }
181 // https://github.com/tokio-rs/tokio/pull/6273#discussion_r1443752911
182 #[allow(clippy::unnecessary_literal_unwrap)]
183 Err(TryLockError::Poisoned(err)) => Err(err).unwrap(),
184 };
185
186 // If we unlocked the child, then the parent may have changed. Check
187 // that we still have the right parent.
188 if let Some(actual_parent) = locked_node.parent.as_ref() {
189 if Arc::ptr_eq(actual_parent, &potential_parent) {
190 return func(locked_node, Some(locked_parent));
191 }
192 }
193 }
194}
195
196/// Moves all children from `node` to `parent`.
197///
198/// `parent` MUST have been a parent of the node when they both got locked,
199/// otherwise there is a potential for a deadlock as invariant #2 would be violated.
200///
201/// To acquire the locks for node and parent, use [`with_locked_node_and_parent`].
202fn move_children_to_parent(node: &mut Inner, parent: &mut Inner) {
203 // Pre-allocate in the parent, for performance
204 parent.children.reserve(additional:node.children.len());
205
206 for child: Arc in std::mem::take(&mut node.children) {
207 {
208 let mut child_locked: MutexGuard<'_, Inner> = child.inner.lock().unwrap();
209 child_locked.parent.clone_from(&node.parent);
210 child_locked.parent_idx = parent.children.len();
211 }
212 parent.children.push(child);
213 }
214}
215
216/// Removes a child from the parent.
217///
218/// `parent` MUST be the parent of `node`.
219/// To acquire the locks for node and parent, use [`with_locked_node_and_parent`].
220fn remove_child(parent: &mut Inner, mut node: MutexGuard<'_, Inner>) {
221 // Query the position from where to remove a node
222 let pos = node.parent_idx;
223 node.parent = None;
224 node.parent_idx = 0;
225
226 // Unlock node, so that only one child at a time is locked.
227 // Otherwise we would violate the lock order (see 'deadlock safety') as we
228 // don't know the creation order of the child nodes
229 drop(node);
230
231 // If `node` is the last element in the list, we don't need any swapping
232 if parent.children.len() == pos + 1 {
233 parent.children.pop().unwrap();
234 } else {
235 // If `node` is not the last element in the list, we need to
236 // replace it with the last element
237 let replacement_child = parent.children.pop().unwrap();
238 replacement_child.inner.lock().unwrap().parent_idx = pos;
239 parent.children[pos] = replacement_child;
240 }
241
242 let len = parent.children.len();
243 if 4 * len <= parent.children.capacity() {
244 parent.children.shrink_to(2 * len);
245 }
246}
247
248/// Increases the reference count of handles.
249pub(crate) fn increase_handle_refcount(node: &Arc<TreeNode>) {
250 let mut locked_node: MutexGuard<'_, Inner> = node.inner.lock().unwrap();
251
252 // Once no handles are left over, the node gets detached from the tree.
253 // There should never be a new handle once all handles are dropped.
254 assert!(locked_node.num_handles > 0);
255
256 locked_node.num_handles += 1;
257}
258
259/// Decreases the reference count of handles.
260///
261/// Once no handle is left, we can remove the node from the
262/// tree and connect its parent directly to its children.
263pub(crate) fn decrease_handle_refcount(node: &Arc<TreeNode>) {
264 let num_handles = {
265 let mut locked_node = node.inner.lock().unwrap();
266 locked_node.num_handles -= 1;
267 locked_node.num_handles
268 };
269
270 if num_handles == 0 {
271 with_locked_node_and_parent(node, |mut node, parent| {
272 // Remove the node from the tree
273 match parent {
274 Some(mut parent) => {
275 // As we want to remove ourselves from the tree,
276 // we have to move the children to the parent, so that
277 // they still receive the cancellation event without us.
278 // Moving them does not violate invariant #1.
279 move_children_to_parent(&mut node, &mut parent);
280
281 // Remove the node from the parent
282 remove_child(&mut parent, node);
283 }
284 None => {
285 // Due to invariant #1, we can assume that our
286 // children can no longer be cancelled through us.
287 // (as we now have neither a parent nor handles)
288 // Therefore we can disconnect them.
289 disconnect_children(&mut node);
290 }
291 }
292 });
293 }
294}
295
296/// Cancels a node and its children.
297pub(crate) fn cancel(node: &Arc<TreeNode>) {
298 let mut locked_node = node.inner.lock().unwrap();
299
300 if locked_node.is_cancelled {
301 return;
302 }
303
304 // One by one, adopt grandchildren and then cancel and detach the child
305 while let Some(child) = locked_node.children.pop() {
306 // This can't deadlock because the mutex we are already
307 // holding is the parent of child.
308 let mut locked_child = child.inner.lock().unwrap();
309
310 // Detach the child from node
311 // No need to modify node.children, as the child already got removed with `.pop`
312 locked_child.parent = None;
313 locked_child.parent_idx = 0;
314
315 // If child is already cancelled, detaching is enough
316 if locked_child.is_cancelled {
317 continue;
318 }
319
320 // Cancel or adopt grandchildren
321 while let Some(grandchild) = locked_child.children.pop() {
322 // This can't deadlock because the two mutexes we are already
323 // holding is the parent and grandparent of grandchild.
324 let mut locked_grandchild = grandchild.inner.lock().unwrap();
325
326 // Detach the grandchild
327 locked_grandchild.parent = None;
328 locked_grandchild.parent_idx = 0;
329
330 // If grandchild is already cancelled, detaching is enough
331 if locked_grandchild.is_cancelled {
332 continue;
333 }
334
335 // For performance reasons, only adopt grandchildren that have children.
336 // Otherwise, just cancel them right away, no need for another iteration.
337 if locked_grandchild.children.is_empty() {
338 // Cancel the grandchild
339 locked_grandchild.is_cancelled = true;
340 locked_grandchild.children = Vec::new();
341 drop(locked_grandchild);
342 grandchild.waker.notify_waiters();
343 } else {
344 // Otherwise, adopt grandchild
345 locked_grandchild.parent = Some(node.clone());
346 locked_grandchild.parent_idx = locked_node.children.len();
347 drop(locked_grandchild);
348 locked_node.children.push(grandchild);
349 }
350 }
351
352 // Cancel the child
353 locked_child.is_cancelled = true;
354 locked_child.children = Vec::new();
355 drop(locked_child);
356 child.waker.notify_waiters();
357
358 // Now the child is cancelled and detached and all its children are adopted.
359 // Just continue until all (including adopted) children are cancelled and detached.
360 }
361
362 // Cancel the node itself.
363 locked_node.is_cancelled = true;
364 locked_node.children = Vec::new();
365 drop(locked_node);
366 node.waker.notify_waiters();
367}
368