Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

sync: try to lock the parent first in CancellationToken #5561

Merged
merged 3 commits into from
Mar 21, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 31 additions & 35 deletions tokio-util/src/sync/cancellation_token/tree_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -151,47 +151,43 @@ fn with_locked_node_and_parent<F, Ret>(node: &Arc<TreeNode>, func: F) -> Ret
where
F: FnOnce(MutexGuard<'_, Inner>, Option<MutexGuard<'_, Inner>>) -> Ret,
{
let mut potential_parent = {
let locked_node = node.inner.lock().unwrap();
match locked_node.parent.clone() {
Some(parent) => parent,
// If we locked the node and its parent is `None`, we are in a valid state
// and can return.
None => return func(locked_node, None),
}
};
use std::sync::TryLockError;

let mut locked_node = node.inner.lock().unwrap();

// Every time this fails, the number of ancestors of the node decreases,
// so the loop must succeed after a finite number of iterations.
loop {
// Deadlock safety:
//
// Due to invariant #2, we know that we have to lock the parent first, and then the child.
// This is true even if the potential_parent is no longer the current parent or even its
// sibling, as the invariant still holds.
let locked_parent = potential_parent.inner.lock().unwrap();
let locked_node = node.inner.lock().unwrap();

let actual_parent = match locked_node.parent.clone() {
Some(parent) => parent,
// If we locked the node and its parent is `None`, we are in a valid state
// and can return.
None => {
// Was the wrong parent, so unlock it before calling `func`
drop(locked_parent);
return func(locked_node, None);
// 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() {
Ok(locked_parent) => locked_parent,
Err(TryLockError::WouldBlock) => {
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().unwrap();
locked_node = node.inner.lock().unwrap();
locked_parent
}
Err(TryLockError::Poisoned(err)) => Err(err).unwrap(),
};

// Loop until we managed to lock both the node and its parent
if Arc::ptr_eq(&actual_parent, &potential_parent) {
return func(locked_node, Some(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() {
if Arc::ptr_eq(actual_parent, &potential_parent) {
return func(locked_node, Some(locked_parent));
}
}

// Drop locked_parent before reassigning to potential_parent,
// as potential_parent is borrowed in it
drop(locked_node);
drop(locked_parent);

potential_parent = actual_parent;
}
}

Expand Down