Reduce context stack contention under free-threaded Python

Author: RazerM

Cargo.lock

@@ -13,4 +13,5 @@ rust-version = "1.83"
 license = "BSD-3-Clause"
 
 [workspace.dependencies]
+arc-swap = "1.7"
 pyo3 = "0.29"

Cargo.toml

@@ -11,20 +11,15 @@
 from __future__ import annotations
 
 import threading
-from collections.abc import Callable, Iterable, Iterator, Sequence
+from collections.abc import Callable, Iterator
 from contextvars import ContextVar
 from itertools import chain, count
-from typing import Any, Generic, SupportsIndex, TypeVar, overload
-from weakref import WeakKeyDictionary
-
-from typing_extensions import TypeAliasType
+from typing import Any, Generic, TypeVar
 
 from logbook.helpers import get_iterator_next_method
 
 _missing = object()
-_MAX_CONTEXT_OBJECT_CACHE = 256
 
-T_co = TypeVar("T_co", covariant=True)
 T = TypeVar("T")
 
 
@@ -105,57 +100,36 @@ def applicationbound(self):
         return ApplicationBound(self)
 
 
-class FrozenSequence(Sequence[T_co]):
-    __slots__ = ("__weakref__", "_hash", "_items")
-
-    def __init__(self, iterable: Iterable[T_co] = ()) -> None:
-        self._items = tuple(iterable)
-        self._hash: int | None = None
-
-    @overload
-    def __getitem__(self, index: SupportsIndex) -> T_co: ...
-
-    @overload
-    def __getitem__(self, index: slice) -> FrozenSequence[T_co]: ...
-
-    def __getitem__(self, index: SupportsIndex | slice) -> T_co | FrozenSequence[T_co]:
-        if isinstance(index, slice):
-            return FrozenSequence(self._items[index])
-        return self._items[index]
-
-    def __len__(self) -> int:
-        return len(self._items)
+class _StackState(Generic[T]):
+    """A state of the persistent context stack.
 
-    def __iter__(self) -> Iterator[T_co]:
-        return iter(self._items)
-
-    def __reversed__(self) -> Iterator[T_co]:
-        return reversed(self._items)
-
-    def __contains__(self, item: object) -> bool:
-        return item in self._items
-
-    def __eq__(self, other: object) -> bool:
-        if isinstance(other, FrozenSequence):
-            return self._items == other._items
-        return NotImplemented
+    Each state is a node in a parent-linked chain: pushing creates a child
+    node and popping restores the parent, so a stack state keeps its
+    identity (and its ``merged`` memo) across push/pop cycles.  The root
+    sentinel has ``item = parent = None``.
+    """
 
-    def __hash__(self) -> int:
-        if self._hash is None:
-            self._hash = hash(self._items)
-        return self._hash
+    __slots__ = ("item", "merged", "parent", "size")
+
+    def __init__(
+        self, item: tuple[int, T] | None, parent: _StackState[T] | None
+    ) -> None:
+        self.item = item
+        self.parent = parent
+        self.size = 0 if parent is None else parent.size + 1
+        self.merged: tuple[tuple[tuple[int, T], ...], tuple[T, ...]] | None = None
+
+    def __iter__(self) -> Iterator[tuple[int, T]]:
+        items = []
+        node = self
+        while node.parent is not None:
+            assert node.item is not None
+            items.append(node.item)
+            node = node.parent
+        return iter(reversed(items))
 
     def __repr__(self) -> str:
-        if self._items:
-            items = repr(self._items)
-        else:
-            items = ""
-        return f"{self.__class__.__name__}({items})"
-
-
-FrozenStack = TypeAliasType(
-    "FrozenStack", FrozenSequence[tuple[int, T]], type_params=(T,)
-)
+        return f"<_StackState len={self.size}>"
 
 
 class ContextStackManager(Generic[T]):
@@ -164,57 +138,82 @@ class ContextStackManager(Generic[T]):
     """
 
     def __init__(self) -> None:
-        self._global: list[T] = []
-        self._context_stack: ContextVar[FrozenStack[T]] = ContextVar(
-            "stack", default=FrozenSequence()
+        self._write_lock = threading.Lock()
+        # () is a singleton, so empty globals from different eras are
+        # identity-equal; harmless for the memo check because equal contents
+        # produce equal merges.
+        self._global: tuple[tuple[int, T], ...] = ()
+        # All contexts share the root node; safe because nodes are immutable
+        # apart from the atomically-published memo slot.
+        self._root: _StackState[T] = _StackState(None, None)
+        self._context_stack: ContextVar[_StackState[T]] = ContextVar(
+            "stack", default=self._root
         )
-        self._cache: WeakKeyDictionary[FrozenStack[T], list[T]] = WeakKeyDictionary()
         self._stackop: Callable[[], int] = get_iterator_next_method(count())
-        self._lock = threading.Lock()
 
     def iter_context_objects(self) -> Iterator[T]:
         """Returns an iterator over all objects for the combined
         application and context cache.
         """
+        node = self._context_stack.get()
+        current_global = self._global
+
+        memo = node.merged
+        # The memo holds a strong reference to the global it was computed
+        # against, so the identity check cannot be fooled by id reuse.
+        if memo is not None and memo[0] is current_global:
+            return iter(memo[1])
+
+        stack_objects = sorted(chain(current_global, node), reverse=True)
+        objects = tuple(x[1] for x in stack_objects)
+        # Nodes are shared across threads via copy_context(); the tuple is
+        # built in full before this single store, so concurrent readers see
+        # either None or a complete memo.
+        node.merged = (current_global, objects)
+
+        # A live chain deliberately keeps a memo per node — each is the warm
+        # cache for a depth the context may pop back to — but memos computed
+        # against an older application stack would pin its popped handlers
+        # until then; drop those now.  Clearing races benignly with
+        # concurrent stores.
+        ancestor = node.parent
+        while ancestor is not None:
+            memo = ancestor.merged
+            if memo is not None and memo[0] is not current_global:
+                ancestor.merged = None
+            ancestor = ancestor.parent
 
-        with self._lock:
-            stack = self._context_stack.get()
-            objects = self._cache.get(stack)
-            if objects is None:
-                if len(self._cache) >= _MAX_CONTEXT_OBJECT_CACHE:
-                    self._cache.clear()
-                stack_objects = sorted(
-                    chain(
-                        self._global,
-                        stack,
-                    ),
-                    reverse=True,
-                )
-                objects = [x[1] for x in stack_objects]
-                self._cache[stack] = objects
         return iter(objects)
 
     def push_context(self, obj: T) -> None:
-        item = (self._stackop(), obj)
-        stack = self._context_stack.get()
-        self._context_stack.set(FrozenSequence((*stack, item)))
+        node = self._context_stack.get()
+        self._context_stack.set(_StackState((self._stackop(), obj), node))
 
     def pop_context(self) -> T:
-        stack = self._context_stack.get()
-        assert stack, "no objects on stack"
-        *remaining, poppped = stack
-        self._context_stack.set(FrozenSequence(remaining))
-        return poppped[1]
+        node = self._context_stack.get()
+        if node.parent is None:
+            raise AssertionError("no objects on stack")
+        assert node.item is not None
+        self._context_stack.set(node.parent)
+        return node.item[1]
 
     def push_application(self, obj: T) -> None:
-        item = (self._stackop(), obj)
-        with self._lock:
-            self._global.append(item)
-            self._cache.clear()
+        with self._write_lock:
+            item = (self._stackop(), obj)
+            self._global = (*self._global, item)
+            # Best-effort lifetime hygiene, not needed for correctness: the
+            # root memo may reference the just-replaced global, and the root
+            # never dies, so drop the memo now rather than waiting for a
+            # future empty-stack iteration to replace it.
+            self._root.merged = None
 
     def pop_application(self) -> T:
-        with self._lock:
-            assert self._global, "no objects on application stack"
-            popped = self._global.pop()
-            self._cache.clear()
-            return popped[1]
+        with self._write_lock:
+            current = self._global
+            if not current:
+                raise AssertionError("no objects on application stack")
+            self._global = current[:-1]
+            # See push_application: keep the immortal root memo from pinning
+            # the popped handler.
+            self._root.merged = None
+            return current[-1][1]

src/logbook/_fallback.py

@@ -8,6 +8,7 @@ rust-version = { workspace = true }
 license = { workspace = true }
 
 [dependencies]
+arc-swap = { workspace = true }
 pyo3 = { workspace = true }
 
 [features]