diff --git a/src/policies/wrapped.rs b/src/policies/wrapped.rs index badc6e9..7eabc6a 100644 --- a/src/policies/wrapped.rs +++ b/src/policies/wrapped.rs @@ -93,6 +93,25 @@ impl Wrapped

{ } } +/// How many handles [`Wrapped::extend`] builds before inserting them, which +/// bounds its memory when the input is much larger than the cache. +const EXTEND_CHUNK: usize = 1024; + +/// Built-in containers whose iteration runs no Python code. Exact types only: +/// a subclass can override `__iter__`. +fn iterates_without_python(obj: &alias::BoundObject) -> bool { + use pyo3::types::PyAnyMethods; + use pyo3::types::PyFrozenSet; + use pyo3::types::PyList; + use pyo3::types::PySet; + use pyo3::types::PyTuple; + + obj.is_exact_instance_of::() + || obj.is_exact_instance_of::() + || obj.is_exact_instance_of::() + || obj.is_exact_instance_of::() +} + #[inline(always)] fn insert_inner( lock: &mut parking_lot::MutexGuard<'_, P>, @@ -214,6 +233,14 @@ impl Wrapped

{ /// other dict-like types; `.items()` is called and the result is iterated. /// - **Any other iterable** — iterated directly, with each element expected to /// unpack as a `(key, value)` pair. + /// + /// A `dict`, `list`, `tuple`, `set` or `frozenset` is inserted in one pass under + /// one lock: iterating it runs no Python code. Any other iterable may run Python + /// code that uses this cache or locks another one (a generator, another + /// cache's `.items()`), so it is never iterated under the lock: handles are + /// built in chunks of [`EXTEND_CHUNK`] and each chunk is inserted under one + /// lock. Entries read before an error are still inserted, like `dict.update`. + // https://github.com/awolverp/cachebox/issues/118 #[inline] pub fn extend(&self, iterable: alias::BoundObject, mut transform: F) -> pyo3::PyResult<()> where @@ -222,7 +249,7 @@ impl Wrapped

{ use pyo3::types::PyAnyMethods; use pyo3::types::PyDictMethods; - let mut lock = self.policy_py_attached(iterable.py()); + let py = iterable.py(); // Using [pyo3::ffi::PyObject_TypeCheck] and [Bound::cast_unchecked] is so faster than [Bound::cast] let is_dictionary = unsafe { @@ -230,6 +257,7 @@ impl Wrapped

{ }; if is_dictionary { let dict = unsafe { iterable.cast_unchecked::() }; + let mut lock = self.policy_py_attached(py); for pair in dict.items() { let (key, value) = unsafe { @@ -237,30 +265,68 @@ impl Wrapped

{ .unwrap_unchecked() }; - insert_inner(&mut lock, &self.shared, pair.py(), transform(key, value)?)?; + insert_inner(&mut lock, &self.shared, py, transform(key, value)?)?; } return Ok(()); } - // By this we will support everything has `.items()` attribute, - // including our cache classes - let items_iterable = { - if let Some(items_attribute) = iterable.getattr_opt(c"items")? { - items_attribute.call0()? - } else { - iterable + if iterates_without_python(&iterable) { + let mut lock = self.policy_py_attached(py); + + for pair in iterable.try_iter()? { + let pair = pair?; + let (key, value) = pair.extract::<(alias::PyObject, alias::PyObject)>()?; + + insert_inner(&mut lock, &self.shared, py, transform(key, value)?)?; } - }; - for pair in items_iterable.try_iter()? { - let pair = pair?; - let (key, value) = pair.extract::<(alias::PyObject, alias::PyObject)>()?; + return Ok(()); + } - insert_inner(&mut lock, &self.shared, pair.py(), transform(key, value)?)?; + let mut handles = Vec::new(); + + let collected = (|| { + // By this we will support everything has `.items()` attribute, + // including our cache classes + let items_iterable = { + if let Some(items_attribute) = iterable.getattr_opt(c"items")? { + items_attribute.call0()? + } else { + iterable + } + }; + + // Only a hint: a generator has no len() (its TypeError is dropped here). + handles.reserve(items_iterable.len().unwrap_or(0).min(EXTEND_CHUNK)); + + for pair in items_iterable.try_iter()? { + let pair = pair?; + let (key, value) = pair.extract::<(alias::PyObject, alias::PyObject)>()?; + + handles.push(transform(key, value)?); + if handles.len() == EXTEND_CHUNK { + self.insert_all(py, &mut handles)?; + } + } + + Ok(()) + })(); + + if !handles.is_empty() { + self.insert_all(py, &mut handles)?; } + collected + } + + /// Inserts `handles` under one lock, leaving `handles` empty. + fn insert_all(&self, py: pyo3::Python<'_>, handles: &mut Vec) -> pyo3::PyResult<()> { + // Declared before the lock, so after an error the handles not inserted + // are dropped (and their finalizers run) once the lock is released. + let mut pending = handles.drain(..); + let mut lock = self.policy_py_attached(py); - Ok(()) + pending.try_for_each(|handle| insert_inner(&mut lock, &self.shared, py, handle).map(drop)) } /// Calls the `evict()` `n` times and returns count of removed items. diff --git a/tests/mixins.py b/tests/mixins.py index 33731cc..94bc495 100644 --- a/tests/mixins.py +++ b/tests/mixins.py @@ -417,6 +417,64 @@ def test_update_overwrites_existing(self): cache.update({"a": 99}) assert cache.get("a") == 99 + def test_update_from_tuple_set_and_list_subclass(self): + cache = self.create_cache() + + cache.update((("t", 1),)) + cache.update({("s", 2)}) + cache.update(frozenset({("f", 3)})) + + class Pairs(list): + # a subclass may run Python code (even this cache) while iterated + def __iter__(self): + assert len(cache) >= 0 + return super().__iter__() + + cache.update(Pairs([("l", 4)])) + assert [cache.get(k) for k in "tsfl"] == [1, 2, 3, 4] + + def test_update_keeps_entries_read_before_an_error(self): + cache = self.create_cache() + + def pairs(): + yield "a", 1 + yield "b", 2 + raise ValueError("stop") + + with pytest.raises(ValueError, match="stop"): + cache.update(pairs()) + assert cache.get("a") == 1 + assert cache.get("b") == 2 + + def test_update_memory_follows_the_cache_not_the_input(self): + tracemalloc = pytest.importorskip("tracemalloc") # not on PyPy + cache = self.create_cache(10) + if isinstance(cache, cachebox.Cache): + pytest.skip("Cache does not evict") + + tracemalloc.start() + try: + cache.update((i, bytes(1000)) for i in range(20_000)) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + # all 20k values alive at once would be about 20 MB + assert peak < 5_000_000, peak + + def test_update_ignores_a_wrong_len(self): + class Lying: + def __len__(self): + return 2**62 + + def __iter__(self): + return iter([("a", 1)]) + + cache = self.create_cache() + + cache.update(Lying()) + assert cache.get("a") == 1 + def test_update_invalid_argument(self): cache = self.create_cache() @@ -1142,6 +1200,64 @@ def insert(base): print("finished") """ +UPDATE_RUNNING_PYTHON_CODE = """ +import faulthandler +import sys +import threading + +import cachebox + +name = sys.argv[1] +cls = getattr(cachebox, name) +# a hang prints every thread's stack before the parent's timeout +faulthandler.dump_traceback_later(50, exit=True) +# switch threads as often as possible, also inside Key.__eq__ +sys.setswitchinterval(1e-6) + + +def make(n): + cache = cls(0, global_ttl=60) if name == "TTLCache" else cls(0) + for i in range(n): + cache[Key(i)] = i + return cache + + +class Key: + def __init__(self, v): + self.v = v + + def __hash__(self): + return self.v + + def __eq__(self, other): + return self.v == other.v + + +# the iterable reads the cache that is being updated +a = make(10) +a.update((Key(i), a.get(Key(i), 0) + 1) for i in range(20)) +assert a.get(Key(19)) == 1, a.get(Key(19)) + +# two caches updated from each other +a, b = make(200), make(200) + + +def update(x, y): + for _ in range(100): + x.update(y) + + +threads = [ + threading.Thread(target=update, args=(a, b)), + threading.Thread(target=update, args=(b, a)), +] +for t in threads: + t.start() +for t in threads: + t.join() +print("ok") +""" + LOCK_WAITERS_WITH_FINALIZERS = """ import faulthandler import random @@ -1331,6 +1447,29 @@ def test_concurrent_inserts_with_finalizers(self): done.stderr or f"exit code {done.returncode}" ) + def test_update_running_python_code(self): + # https://github.com/awolverp/cachebox/issues/118 + # + # update() iterated its argument while holding the cache lock, so an + # iterable that used the same cache hung, and so did a.update(b) with + # b.update(a) on another thread. + name = type(self.create_cache()).__name__ + + try: + done = subprocess.run( + [sys.executable, "-c", UPDATE_RUNNING_PYTHON_CODE, name], + capture_output=True, + text=True, + timeout=60, + check=False, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"{name}: update() never returned") + + assert done.stdout.strip() == "ok", ( + done.stderr or f"exit code {done.returncode}" + ) + def test_lock_waiters_with_finalizers(self): # Follow-up to #95: methods such as len() and clear() waited for the # cache lock without detaching. With the GIL, a finalizer running under