lrucache: allow replacing an entry, make it thread-safe

src/borg/helpers/lrucache.py

@@ -1,16 +1,31 @@
+import threading
 from collections import OrderedDict
 from collections.abc import Callable, ItemsView, Iterator, KeysView, MutableMapping, ValuesView
-from typing import TypeVar
+from typing import Any, TypeVar
 
 K = TypeVar("K")
 V = TypeVar("V")
 
+_MISSING = object()
+
 
 class LRUCache(MutableMapping[K, V]):
     """
     Mapping which maintains a maximum size by removing the least recently used value.
-    Items are passed to dispose before being removed and setting an item which is
-    already in the cache has to be done using the replace method.
+
+    A value that leaves the cache - evicted, deleted, or replaced by another value under
+    the same key - is passed to *dispose* first, so that a cache of e.g. open files can
+    close them. Use replace() to update an entry while keeping its old value alive.
+
+    Thread safety: getting, setting, deleting, popping, replacing and clearing are each
+    atomic, so several threads can share one cache without corrupting it. A *sequence* of
+    them is not: two threads can miss on the same key and both compute and store a value,
+    which for a cache is wasteful, not wrong - and the mixin methods built on top of the
+    above (setdefault, popitem, update) inherit that. Iterating (also via keys(), values(),
+    items()) yields a live view of the underlying dict, so it needs external
+    synchronization if another thread may write meanwhile.
+
+    *dispose* is called while the lock is held, so it must not use the cache.
     """
 
     _cache: OrderedDict[K, V]
@@ -23,38 +38,67 @@ def __init__(self, capacity: int, dispose: Callable[[V], None] = lambda _: None)
         self._cache = OrderedDict()
         self._capacity = capacity
         self._dispose = dispose
+        self._lock = threading.Lock()
 
     def __setitem__(self, key: K, value: V) -> None:
-        assert (
-            key not in self._cache
-        ), "Unexpected attempt to replace a cached item without first deleting the old item."
-        while len(self._cache) >= self._capacity:
-            self._dispose(self._cache.popitem(last=False)[1])
-        self._cache[key] = value  # add new entry at the end
+        with self._lock:
+            try:
+                previous = self._cache.pop(key)
+            except KeyError:
+                pass
+            else:
+                # the old value is no longer in the cache, so it is disposed like a deleted one.
+                if previous is not value:
+                    self._dispose(previous)
+            while len(self._cache) >= self._capacity:
+                self._dispose(self._cache.popitem(last=False)[1])
+            self._cache[key] = value  # add the new (or refreshed) entry at the end
 
     def __getitem__(self, key: K) -> V:
-        self._cache.move_to_end(key)  # raise KeyError if not found
-        return self._cache[key]
+        with self._lock:
+            self._cache.move_to_end(key)  # raise KeyError if not found
+            return self._cache[key]
 
     def __delitem__(self, key: K) -> None:
-        self._dispose(self._cache.pop(key))
+        with self._lock:
+            self._dispose(self._cache.pop(key))
 
     def __contains__(self, key: object) -> bool:
         return key in self._cache
 
     def __len__(self) -> int:
         return len(self._cache)
 
+    def pop(self, key: K, default: Any = _MISSING) -> Any:
+        """Remove *key* and return its value, without disposing it."""
+        # MutableMapping would implement this as __getitem__ + __delitem__, which is not
+        # atomic (and would dispose the value).
+        with self._lock:
+            try:
+                return self._cache.pop(key)
+            except KeyError:
+                if default is _MISSING:
+                    raise
+                return default
+
     def replace(self, key: K, value: V) -> None:
-        """Replace an item that is already present, not disposing it in the process."""
-        # this method complements __setitem__ which should be used for the normal use case.
-        assert key in self._cache, "Unexpected attempt to update a non-existing item."
-        self._cache[key] = value
+        """Replace the value of an entry that is present, without disposing the old value.
+
+        This is for the rare case where the old value must stay alive, e.g. because the new
+        value still refers to it (the legacy repository's fd cache re-times-stamps its
+        entries this way, keeping the open file object). It also keeps the entry where it
+        is in the LRU order, i.e. it does not count as a use. Everything else should just
+        assign to the cache.
+        """
+        with self._lock:
+            assert key in self._cache, "Unexpected attempt to update a non-existing item."
+            self._cache[key] = value
 
     def clear(self) -> None:
-        for value in self._cache.values():
-            self._dispose(value)
-        self._cache.clear()
+        with self._lock:
+            for value in self._cache.values():
+                self._dispose(value)
+            self._cache.clear()
 
     def __iter__(self) -> Iterator[K]:
         return iter(self._cache)

src/borg/testsuite/helpers/lrucache_test.py

@@ -1,3 +1,5 @@
+import sys
+import threading
 from tempfile import TemporaryFile
 
 import pytest
@@ -35,6 +37,33 @@ def test_lrucache(self):
         c.clear()
         assert c.items() == set()
 
+    def test_assign_to_existing_key(self):
+        # assigning to a key that is present replaces the value and counts as a use,
+        # so the entry becomes the most recently used one.
+        c = LRUCache(2)
+        c["a"] = 1
+        c["b"] = 2
+        c["a"] = 3  # replaces, and refreshes "a"
+        assert len(c) == 2
+        assert c["a"] == 3
+        c["c"] = 4  # evicts the least recently used entry, which is "b" now
+        assert set(c.keys()) == {"a", "c"}
+
+    def test_dispose_on_replacement(self):
+        disposed = []
+        c = LRUCache(2, dispose=disposed.append)
+        c["a"] = "first"
+        c["a"] = "second"  # the replaced value leaves the cache, so it is disposed
+        assert disposed == ["first"]
+        assert c["a"] == "second"
+        value = "same object"
+        c["b"] = value
+        c["b"] = value  # assigning the identical object does not dispose it
+        assert disposed == ["first"]
+        c.replace("b", "replacement")  # replace() never disposes, see there
+        assert disposed == ["first"]
+        assert c["b"] == "replacement"
+
     def test_dispose(self):
         c = LRUCache(2, dispose=lambda f: f.close())
         f1 = TemporaryFile()
@@ -54,3 +83,50 @@ def test_dispose(self):
         c.clear()
         assert c.items() == set()
         assert f3.closed
+
+    def test_pop(self):
+        disposed = []
+        c = LRUCache(2, dispose=disposed.append)
+        c["a"] = 1
+        assert c.pop("a") == 1  # pop hands the value to the caller, so it is not disposed
+        assert disposed == []
+        assert "a" not in c
+        assert c.pop("a", "default") == "default"
+        with pytest.raises(KeyError):
+            c.pop("a")
+
+    def test_threaded_access(self):
+        # several threads sharing one cache must not corrupt it or raise: they hit the same
+        # keys, so they race on inserting, refreshing, evicting and removing the same entries.
+        # e.g. popping via MutableMapping (get, then delete) raises KeyError here when
+        # another thread removes the key in between.
+        c: LRUCache = LRUCache(8)
+        keys = list(range(16))  # more keys than capacity, so entries keep getting evicted
+        errors: list[Exception] = []
+        start = threading.Barrier(8)
+
+        def worker(n):
+            start.wait()
+            try:
+                for _round in range(2000):
+                    for key in keys:
+                        c[key] = n
+                        c.get(key)
+                        c.pop(key, None)
+            except Exception as e:
+                errors.append(e)
+
+        # the default switch interval (5 ms) is longer than these loops, so threads would
+        # hardly ever switch inside one cache operation and the races would not show up.
+        previous_interval = sys.getswitchinterval()
+        sys.setswitchinterval(1e-6)
+        try:
+            threads = [threading.Thread(target=worker, args=(n,)) for n in range(8)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+        finally:
+            sys.setswitchinterval(previous_interval)
+        assert errors == []
+        assert len(c) <= 8