From 91caac6aacd24de5938a174378a6b9c92e2f5acd Mon Sep 17 00:00:00 2001 From: aldenhu Date: Wed, 18 Sep 2024 17:20:08 +0000 Subject: [PATCH] Fix potential deadlock in state_key::Entry::drop If before Entry::drop() calls Registry::maybe_remove, the weak ref under the same key1 and key2 is replaced, Entry::drop can be called recursively resulting in a deadlock. The crux is we try to see if the entry has been replaced by trying to upgrading it to a strong ref, and it possible that the upgrade is successful and the result becomes a only ref to the new entry. This tries to fix it by determining the same thing by seeing if Weak::strong_count() is 0 instead of upgrading the weak ref. --- types/src/state_store/state_key/registry.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/types/src/state_store/state_key/registry.rs b/types/src/state_store/state_key/registry.rs index 1fad8c4d90a15..c56577eda6653 100644 --- a/types/src/state_store/state_key/registry.rs +++ b/types/src/state_store/state_key/registry.rs @@ -152,13 +152,13 @@ where let mut locked = self.inner.write(); if let Some(map2) = locked.get_mut(key1) { if let Some(entry) = map2.get(key2) { - if entry.upgrade().is_none() { + if entry.strong_count() == 0 { map2.remove(key2); + if map2.is_empty() { + locked.remove(key1); + } } } - if map2.is_empty() { - locked.remove(key1); - } } }