perf: O(1) expected random pick for RRCache eviction and random_key

src/policies/rrpolicy.rs

@@ -115,8 +115,38 @@ impl RRPolicy {
     pub fn table(&self) -> &hashbrown::raw::RawTable<Handle> {
         &self.table
     }
+
+    /// Returns a uniformly random occupied bucket, or `None` if the table is empty.
+    ///
+    /// Samples random bucket indices and keeps the first occupied one; conditioned on
+    /// hitting an occupied bucket, every entry is equally likely.
+    ///
+    /// Expected O(1) for a full cache: hashbrown only grows by doubling, so a table
+    /// that grew through insertions stays roughly a quarter occupied or more, even
+    /// after eviction churn (a few probes on average). A table pre-sized via
+    /// `capacity`, thinned by removals, or holding few heavy entries under
+    /// `getsizeof` can be much sparser; it falls back to a uniform O(n) pick after
+    /// [`RANDOM_PROBES`] misses.
+    pub fn random_bucket(&self) -> Option<hashbrown::raw::Bucket<Handle>> {
+        if self.table.is_empty() {
+            return None;
+        }
+
+        let buckets = self.table.num_buckets();
+        for _ in 0..RANDOM_PROBES {
+            if let Some(bucket) = self.table.checked_bucket(fastrand::usize(0..buckets)) {
+                return Some(bucket);
+            }
+        }
+
+        let nth = fastrand::usize(0..self.table.len());
+        unsafe { self.table.iter().nth(nth) }
+    }
 }
 
+/// Random bucket probes in [`RRPolicy::random_bucket`] before the O(n) fallback.
+const RANDOM_PROBES: usize = 32;
+
 impl PolicyExt for RRPolicy {
     type Shared = Shared;
     type Handle = Handle;
@@ -177,19 +207,15 @@ impl PolicyExt for RRPolicy {
 
     #[inline]
     fn evict(&mut self, shared: &Self::Shared) -> pyo3::PyResult<Self::Handle> {
-        if self.table.is_empty() {
-            Err(new_py_error!(PyKeyError, "cache is empty"))
-        } else {
-            let nth = fastrand::usize(0..self.table.len());
+        let bucket = self
+            .random_bucket()
+            .ok_or_else(|| new_py_error!(PyKeyError, "cache is empty"))?;
 
-            let bucket = unsafe { self.table.iter().nth(nth).unwrap_unchecked() };
-
-            shared.generation_version().increment();
+        shared.generation_version().increment();
 
-            let (handle, _) = unsafe { self.table.remove(bucket) };
-            self.currsize = self.currsize.saturating_sub(handle.size());
-            Ok(handle)
-        }
+        let (handle, _) = unsafe { self.table.remove(bucket) };
+        self.currsize = self.currsize.saturating_sub(handle.size());
+        Ok(handle)
     }
 
     #[inline]

src/pyclasses/rrcache.rs

@@ -665,16 +665,12 @@ impl PyRRCache {
         let inner = self.0.get()?;
         let policy = inner.policy_py_attached(py);
 
-        if policy.table().is_empty() {
-            Err(new_py_error!(PyKeyError, "cache is empty"))
-        } else {
-            let nth = fastrand::usize(0..policy.table().len());
+        let bucket = policy
+            .random_bucket()
+            .ok_or_else(|| new_py_error!(PyKeyError, "cache is empty"))?;
 
-            let bucket = unsafe { policy.table().iter().nth(nth).unwrap_unchecked() };
-
-            let handle = unsafe { bucket.as_ref() };
-            Ok(handle.key().clone_ref(py).into())
-        }
+        let handle = unsafe { bucket.as_ref() };
+        Ok(handle.key().clone_ref(py).into())
     }
 
     fn __traverse__(&self, visit: pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> {

tests/test_impls.py

@@ -1,3 +1,4 @@
+import collections
 import subprocess
 import sys
 import sysconfig
@@ -497,6 +498,44 @@ def test_random_key_method(self):
         cache["d"] = 4
         assert cache.random_key() in ("a", "b", "c", "d")
 
+    @staticmethod
+    def _assert_roughly_uniform(draws, keys):
+        counts = collections.Counter(draws)
+        assert set(counts) == set(keys)
+        expected = len(draws) / len(keys)
+        # ~10 standard deviations: catches a biased pick, never flakes.
+        assert all(abs(n - expected) < expected * 0.3 for n in counts.values())
+
+    def test_random_key_is_uniform(self):
+        keys = list(range(8))
+        cache = self.create_cache(8, {k: k for k in keys})
+        self._assert_roughly_uniform([cache.random_key() for _ in range(8000)], keys)
+
+    def test_random_key_is_uniform_in_sparse_table(self):
+        # Many removals leave few entries in a large table (the fallback path).
+        cache = self.create_cache(0, {k: k for k in range(10_000)})
+        keys = [3, 1_234, 9_999]
+        for k in range(10_000):
+            if k not in keys:
+                del cache[k]
+        assert cache.capacity() >= 10_000
+        self._assert_roughly_uniform([cache.random_key() for _ in range(3000)], keys)
+
+    def test_popitem_drains_sparse_table(self):
+        cache = self.create_cache(0, {k: k for k in range(5_000)})
+        for k in range(4_990):
+            del cache[k]
+        popped = {cache.popitem()[0] for _ in range(10)}
+        assert popped == set(range(4_990, 5_000))
+        assert len(cache) == 0
+
+    def test_eviction_keeps_cache_full(self):
+        cache = self.create_cache(1_000)
+        for k in range(10_000):
+            cache[k] = k
+            assert len(cache) == min(k + 1, 1_000)
+        assert all(cache[k] == k for k in cache)
+
 
 class TestLRUCache(
     mixins.InitializeMixin,