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 //!
41 use 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.
46 pub(crate) struct TreeNode {
47     inner: Mutex<Inner>,
48     waker: tokio::sync::Notify,
49 }
50 impl TreeNode {
new() -> Self51     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 
notified(&self) -> tokio::sync::futures::Notified<'_>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.
73 struct 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
is_cancelled(node: &Arc<TreeNode>) -> bool82 pub(crate) fn is_cancelled(node: &Arc<TreeNode>) -> bool {
83     node.inner.lock().unwrap().is_cancelled
84 }
85 
86 /// Creates a child node
child_node(parent: &Arc<TreeNode>) -> Arc<TreeNode>87 pub(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.
disconnect_children(node: &mut Inner)125 fn disconnect_children(node: &mut Inner) {
126     for child in std::mem::take(&mut node.children) {
127         let mut locked_child = 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)).
with_locked_node_and_parent<F, Ret>(node: &Arc<TreeNode>, func: F) -> Ret where F: FnOnce(MutexGuard<'_, Inner>, Option<MutexGuard<'_, Inner>>) -> Ret,150 fn with_locked_node_and_parent<F, Ret>(node: &Arc<TreeNode>, func: F) -> Ret
151 where
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`].
move_children_to_parent(node: &mut Inner, parent: &mut Inner)202 fn move_children_to_parent(node: &mut Inner, parent: &mut Inner) {
203     // Pre-allocate in the parent, for performance
204     parent.children.reserve(node.children.len());
205 
206     for child in std::mem::take(&mut node.children) {
207         {
208             let mut child_locked = 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`].
remove_child(parent: &mut Inner, mut node: MutexGuard<'_, Inner>)220 fn 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.
increase_handle_refcount(node: &Arc<TreeNode>)249 pub(crate) fn increase_handle_refcount(node: &Arc<TreeNode>) {
250     let mut locked_node = 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.
decrease_handle_refcount(node: &Arc<TreeNode>)263 pub(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.
cancel(node: &Arc<TreeNode>)297 pub(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