From 37fb51d85f09ef3b6859c488f656af03b617595a Mon Sep 17 00:00:00 2001 From: Roman Varnava Date: Wed, 28 Oct 2026 16:34:51 +0300 Subject: [PATCH] perf: O(1) expected random pick for RRCache eviction and random_key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RRCache.evict() and random_key() chose an entry with table.iter().nth(random), which walks the table: O(n) per eviction. With a full cache of 100k entries an insert cost ~23 µs. Pick a random bucket index and keep the first occupied one instead. Conditioned on hitting an occupied bucket, every entry is equally likely, so the distribution stays uniform. hashbrown only grows by doubling, so a full cache stays roughly a quarter occupied or more even after eviction churn, which means a few probes on average. A sparser table (pre-sized via `capacity`, thinned by removals, or holding few heavy entries under `getsizeof`) falls back to the old uniform O(n) pick after 32 misses. Tests cover uniformity in a full and in a sparse table, draining a sparse table with popitem(), and keeping the cache full under eviction. --- src/policies/rrpolicy.rs | 48 +++++++++++++++++++++++++++++++--------- src/pyclasses/rrcache.rs | 14 +++++------- tests/test_impls.py | 39 ++++++++++++++++++++++++++++++++ 3 files changed, 81 insertions(+), 20 deletions(-) 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,