Fix a bug that eviction listener's key-level locks are not removed

when immediate notification mode is used

src/sync/cache.rs

@@ -2054,6 +2054,10 @@ where
     pub(crate) fn set_expiration_clock(&self, clock: Option<crate::common::time::Clock>) {
         self.base.set_expiration_clock(clock);
     }
+
+    pub(crate) fn key_locks_map_is_empty(&self) -> bool {
+        self.base.key_locks_map_is_empty()
+    }
 }
 
 // To see the debug prints, run test as `cargo test -- --nocapture`
@@ -4339,6 +4343,8 @@ mod tests {
             assert_eq!(a[0], (Arc::new("alice"), "a3", RemovalCause::Expired));
             a.clear();
         }
+
+        assert!(cache.key_locks_map_is_empty());
     }
 
     // This test ensures the key-level lock for the immediate notification
@@ -4469,6 +4475,8 @@ mod tests {
         for (i, (actual, expected)) in actual.iter().zip(&expected).enumerate() {
             assert_eq!(actual, expected, "expected[{}]", i);
         }
+
+        assert!(cache.key_locks_map_is_empty());
     }
 
     // NOTE: To enable the panic logging, run the following command:

src/sync_base/base_cache.rs

@@ -785,6 +785,10 @@ where
     pub(crate) fn set_expiration_clock(&self, clock: Option<Clock>) {
         self.inner.set_expiration_clock(clock);
     }
+
+    pub(crate) fn key_locks_map_is_empty(&self) -> bool {
+        self.inner.key_locks_map_is_empty()
+    }
 }
 
 struct EvictionState<'a, K, V> {
@@ -2415,6 +2419,10 @@ where
             *exp_clock = None;
         }
     }
+
+    fn key_locks_map_is_empty(&self) -> bool {
+        self.key_locks.as_ref().map(|m| m.is_empty()).unwrap_or_default()
+    }
 }
 
 //

src/sync_base/key_lock.rs

@@ -30,11 +30,11 @@ where
     S: BuildHasher,
 {
     fn drop(&mut self) {
-        if TrioArc::count(&self.lock) <= 1 {
+        if TrioArc::count(&self.lock) <= 2 {
             self.map.remove_if(
                 self.hash,
                 |k| k == &self.key,
-                |_k, v| TrioArc::count(v) <= 1,
+                |_k, v| TrioArc::count(v) <= 2,
             );
         }
     }
@@ -86,3 +86,10 @@ where
         }
     }
 }
+
+#[cfg(test)]
+impl<K, S> KeyLockMap<K, S> {
+    pub(crate) fn is_empty(&self) -> bool {
+        self.locks.len() == 0
+    }
+}