beam: hardened thread-safe LRU cache (sentinel + RLock)

wizard-work/beam/lru_cache.py

@@ -1,17 +1,22 @@
-"""O(1) LRU cache: doubly linked list + hash map.
+"""Thread-safe O(1) LRU cache (doubly linked list + hash map).
 
-Design notes:
-- Each node holds key, value, prev, next.
-- `_map` maps key -> node.
-- Head is MRU sentinel, tail is LRU sentinel. Most-recently-used sits right after head.
-- `get` and `put`-on-existing must BOTH promote the node to MRU.
-- `put`-on-new inserts at MRU; if over capacity, evicts the node before tail (LRU).
-- Eviction must delete from `_map` (common bug: forget this).
+Subtle bugs fixed vs. naive versions:
+- `get`/`put`-on-existing BOTH promote to MRU (forgetting promote on get is the
+  classic LRU bug).
+- Eviction deletes from `_map` (forgetting this leaks memory and breaks size).
+- Missing-key sentinel uses a unique `_MISSING` object instead of `-1` so a
+  stored value of -1 is not mistaken for "absent".
+- All public ops hold a single RLock so concurrent access cannot corrupt the
+  linked list or hand out a value for an evicted key.
+- Capacity validated as positive int (bool rejected).
 """
 from __future__ import annotations
 
+import threading
 from typing import Any, Optional
 
+_MISSING = object()
+
 
 class _Node:
     __slots__ = ("key", "value", "prev", "next")
@@ -25,22 +30,20 @@ class _Node:
 
 class LRUCache:
     def __init__(self, capacity: int) -> None:
-        if not isinstance(capacity, int) or isinstance(capacity, bool):
+        if isinstance(capacity, bool) or not isinstance(capacity, int):
             raise TypeError("capacity must be an int")
         if capacity <= 0:
             raise ValueError("capacity must be positive")
         self.capacity = capacity
         self._map: dict[Any, _Node] = {}
-        # sentinels
-        self._head = _Node(None, None)  # MRU side
-        self._tail = _Node(None, None)  # LRU side
+        self._head = _Node(None, None)  # MRU side sentinel
+        self._tail = _Node(None, None)  # LRU side sentinel
         self._head.next = self._tail
         self._tail.prev = self._head
+        self._lock = threading.RLock()
 
-    # ---- internal list ops ----
     def _detach(self, node: _Node) -> None:
-        p = node.prev
-        n = node.next
+        p, n = node.prev, node.next
         assert p is not None and n is not None
         p.next = n
         n.prev = p
@@ -48,7 +51,6 @@ class LRUCache:
         node.next = None
 
     def _attach_mru(self, node: _Node) -> None:
-        # insert right after head
         first = self._head.next
         self._head.next = node
         node.prev = self._head
@@ -60,39 +62,43 @@ class LRUCache:
         self._detach(node)
         self._attach_mru(node)
 
-    # ---- public API ----
-    def get(self, key: Any) -> Any:
-        node = self._map.get(key)
-        if node is None:
-            return -1
-        self._promote(node)
-        return node.value
+    def get(self, key: Any, default: Any = _MISSING) -> Any:
+        with self._lock:
+            node = self._map.get(key)
+            if node is None:
+                return None if default is _MISSING else default
+            self._promote(node)
+            return node.value
 
     def put(self, key: Any, value: Any) -> None:
-        node = self._map.get(key)
-        if node is not None:
-            # update value AND recency (subtle bug if you skip promote)
-            node.value = value
-            self._promote(node)
-            return
-        node = _Node(key, value)
-        self._map[key] = node
-        self._attach_mru(node)
-        if len(self._map) > self.capacity:
-            # evict LRU: node just before tail
-            lru = self._tail.prev
-            assert lru is not None and lru is not self._head
-            self._detach(lru)
-            del self._map[lru.key]
+        with self._lock:
+            node = self._map.get(key)
+            if node is not None:
+                node.value = value
+                self._promote(node)
+                return
+            node = _Node(key, value)
+            self._map[key] = node
+            self._attach_mru(node)
+            if len(self._map) > self.capacity:
+                lru = self._tail.prev
+                assert lru is not None and lru is not self._head
+                self._detach(lru)
+                del self._map[lru.key]
+
+    def contains(self, key: Any) -> bool:
+        with self._lock:
+            return key in self._map
 
     def size(self) -> int:
-        return len(self._map)
+        with self._lock:
+            return len(self._map)
 
-    # helpers for tests / introspection
     def order_mru_to_lru(self) -> list[Any]:
-        out = []
-        cur = self._head.next
-        while cur is not None and cur is not self._tail:
-            out.append(cur.key)
-            cur = cur.next
-        return out
+        with self._lock:
+            out = []
+            cur = self._head.next
+            while cur is not None and cur is not self._tail:
+                out.append(cur.key)
+                cur = cur.next
+            return out