From 7d98dec32795dbf1028016dc4fe4e3b88004b6bd Mon Sep 17 00:00:00 2001
From: Simon <141160134+so-schen@users.noreply.github.com>
Date: Sun, 10 Nov 2024 17:20:53 +0100
Subject: [PATCH] Update registry cache

improve cache
---
 types/src/state_store/state_key/registry.rs | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/types/src/state_store/state_key/registry.rs b/types/src/state_store/state_key/registry.rs
index 222664b2b1b01..c56577eda6653 100644
--- a/types/src/state_store/state_key/registry.rs
+++ b/types/src/state_store/state_key/registry.rs
@@ -151,9 +151,13 @@ where
     fn maybe_remove(&self, key1: &Key1, key2: &Key2) {
         let mut locked = self.inner.write();
         if let Some(map2) = locked.get_mut(key1) {
-            map2.remove(key2);
-            if map2.is_empty() {
-                locked.remove(key1);
+            if let Some(entry) = map2.get(key2) {
+                if entry.strong_count() == 0 {
+                    map2.remove(key2);
+                    if map2.is_empty() {
+                        locked.remove(key1);
+                    }
+                }
             }
         }
     }