diff --git a/tokio-util/src/sync/cancellation_token/tree_node.rs b/tokio-util/src/sync/cancellation_token/tree_node.rs index 8f97dee8de5..f1abb4a8d33 100644 --- a/tokio-util/src/sync/cancellation_token/tree_node.rs +++ b/tokio-util/src/sync/cancellation_token/tree_node.rs @@ -151,47 +151,43 @@ fn with_locked_node_and_parent(node: &Arc, func: F) -> Ret where F: FnOnce(MutexGuard<'_, Inner>, Option>) -> 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; } }