diff --git a/src/policies/rrpolicy.rs b/src/policies/rrpolicy.rs index b0ad113..46e4930 100644 --- a/src/policies/rrpolicy.rs +++ b/src/policies/rrpolicy.rs @@ -115,8 +115,38 @@ impl RRPolicy { pub fn table(&self) -> &hashbrown::raw::RawTable { &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> { + 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 { - 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] diff --git a/src/pyclasses/rrcache.rs b/src/pyclasses/rrcache.rs index b433dd5..3eb9984 100644 --- a/src/pyclasses/rrcache.rs +++ b/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> { diff --git a/tests/test_impls.py b/tests/test_impls.py index 34e6943..f733623 100644 --- a/tests/test_impls.py +++ b/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,