fix: build the handles for update() outside the cache lock, in chunks

src/policies/wrapped.rs

@@ -93,6 +93,25 @@ impl<P: PolicyExt> Wrapped<P> {
     }
 }
 
+/// 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::<PyList>()
+        || obj.is_exact_instance_of::<PyTuple>()
+        || obj.is_exact_instance_of::<PySet>()
+        || obj.is_exact_instance_of::<PyFrozenSet>()
+}
+
 #[inline(always)]
 fn insert_inner<P: PolicyExt>(
     lock: &mut parking_lot::MutexGuard<'_, P>,
@@ -214,6 +233,14 @@ impl<P: PolicyExt> Wrapped<P> {
     ///   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<F>(&self, iterable: alias::BoundObject, mut transform: F) -> pyo3::PyResult<()>
     where
@@ -222,45 +249,84 @@ impl<P: PolicyExt> Wrapped<P> {
         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 {
             pyo3::ffi::PyObject_TypeCheck(iterable.as_ptr(), crate::typeref::STD_DICT_TYPE) == 1
         };
         if is_dictionary {
             let dict = unsafe { iterable.cast_unchecked::<pyo3::types::PyDict>() };
+            let mut lock = self.policy_py_attached(py);
 
             for pair in dict.items() {
                 let (key, value) = unsafe {
                     pair.extract::<(alias::PyObject, alias::PyObject)>()
                         .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<P::Handle>) -> 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.

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