Ensure heapq use is safe under freethreading

Author: mikeshardmind

src/async_utils/__init__.py

@@ -9,7 +9,7 @@
 __author__ = "Michael Hall"
 __license__ = "Apache-2.0"
 __copyright__ = "Copyright 2020-Present Michael Hall"
-__version__ = "2025.06.05b"
+__version__ = "2025.06.07b"
 
 import os
 import sys

src/async_utils/_qs.py

@@ -20,6 +20,10 @@
 import threading
 from collections import deque
 from collections.abc import Generator
+
+# PYUPDATE: 3.14 release + 3.14 minimum: reaudit
+# heapq methods are not threadsafe pre 3.14
+# see: GH: cpython 135036
 from heapq import heappop, heappush
 
 from . import _typings as t
@@ -566,7 +570,7 @@ def _get(self, /) -> T:
 class PriorityQueue[T](BaseQueue[T]):
     """A thread-safe queue with both sync and async access methods."""
 
-    __slots__ = ("_data",)
+    __slots__ = ("_data", "_lock")
 
     def __init_subclass__(cls) -> t.Never:
         msg = "Don't subclass this"
@@ -577,6 +581,8 @@ def __init_subclass__(cls) -> t.Never:
     def __init__(self, /, maxsize: int | None = None) -> None:
         super().__init__(maxsize)
         self._data: list[T] = []
+        # heapq not threadsafe till 3.14
+        self._lock = threading.RLock()
 
     def _qsize(self, /) -> int:
         return len(self._data)
@@ -585,7 +591,9 @@ def _items(self, /) -> list[T]:
         return self._data.copy()
 
     def _put(self, item: T, /) -> None:
-        heappush(self._data, item)
+        with self._lock:
+            heappush(self._data, item)
 
     def _get(self, /) -> T:
-        return heappop(self._data)
+        with self._lock:
+            return heappop(self._data)

src/async_utils/lockout.py

@@ -16,7 +16,12 @@
 
 import asyncio
 import concurrent.futures as cf
+
+# PYUPDATE: 3.14 release + 3.14 minimum: reaudit
+# heapq methods are not threadsafe pre 3.14
+# see: GH: cpython 135036
 import heapq
+import threading
 import time
 from collections import deque
 from functools import partial
@@ -76,6 +81,7 @@ def __repr__(self) -> str:
 
     def __init__(self) -> None:
         self._lockouts: list[float] = []
+        self._internal_lock: threading.RLock = threading.RLock()
 
     def lockout_for(self, seconds: float, /) -> None:
         """Lock a resource for an amount of time."""
@@ -86,9 +92,13 @@ async def __aenter__(self) -> None:
             now = time.monotonic()
             # There must not be an async context switch between here
             # and replacing the lockout when lockout is in the future
-            ts = heapq.heappop(self._lockouts)
-            if (sleep_for := ts - now) > 0:
-                heapq.heappush(self._lockouts, ts)
+            # PYUPDATE: The lock here can be removed at 3.14 minimum
+            with self._internal_lock:
+                ts = heapq.heappop(self._lockouts)
+                if (sleep_for := ts - now) > 0:
+                    heapq.heappush(self._lockouts, ts)
+
+            if sleep_for > 0:
                 await asyncio.sleep(sleep_for)
 
     async def __aexit__(self, *_dont_care: object) -> None:

src/async_utils/lru.py

@@ -16,8 +16,12 @@
 
 from __future__ import annotations
 
+# PYUPDATE: 3.14 release + 3.14 minimum: reaudit
+# heapq methods are not threadsafe pre 3.14
+# see: GH: cpython 135036
 import heapq
 import math
+import threading
 import time
 
 from . import _typings as t
@@ -191,7 +195,14 @@ class TTLLRU[K, V]:
         Getting items does not refresh their ttl.
     """
 
-    __slots__ = ("_cache", "_expirations", "_maxsize", "_smooth", "_ttl")
+    __slots__ = (
+        "_cache",
+        "_expirations",
+        "_internal_lock",
+        "_maxsize",
+        "_smooth",
+        "_ttl",
+    )
 
     def __init_subclass__(cls) -> t.Never:
         msg = "Don't subclass this"
@@ -205,29 +216,31 @@ def __init__(self, maxsize: int, ttl: float) -> None:
         self._ttl: float = ttl
         self._expirations: list[tuple[float, K]] = []
         self._smooth: int = max(int(math.log2(maxsize // 2)), 1)
+        self._internal_lock: threading.RLock = threading.RLock()
 
     def _remove_some_expired(self) -> None:
         """Remove some number of expired entries."""
         now = time.monotonic()
         tr = max((len(self._expirations) - self._maxsize) >> self._smooth, 2)
 
         while self._expirations and tr > 0:
-            try:
-                ts, k = heapq.heappop(self._expirations)
-            except IndexError:
-                continue
-
-            if ts < now:
-                tr -= 1
+            with self._internal_lock:
                 try:
-                    ts, _v = self._cache[k]
-                except KeyError:
+                    ts, k = heapq.heappop(self._expirations)
+                except IndexError:
                     continue
+
                 if ts < now:
-                    self._cache.pop(k, None)
-            else:
-                heapq.heappush(self._expirations, (ts, k))
-                break
+                    tr -= 1
+                    try:
+                        ts, _v = self._cache[k]
+                    except KeyError:
+                        continue
+                    if ts < now:
+                        self._cache.pop(k, None)
+                else:
+                    heapq.heappush(self._expirations, (ts, k))
+                    break
 
     # needed because in the absence of __iter__ or __contains__ python will
     # attempt iteration/containment checks via __getitem__
@@ -251,7 +264,8 @@ def __getitem__(self, key: K, /) -> V:
 
     def __setitem__(self, key: K, value: V, /) -> None:
         ts = time.monotonic() + self._ttl
-        heapq.heappush(self._expirations, (ts, key))
+        with self._internal_lock:
+            heapq.heappush(self._expirations, (ts, key))
         self._cache[key] = (ts, value)
         self._remove_some_expired()
         if len(self._cache) > self._maxsize:

src/async_utils/priority_sem.py

@@ -17,6 +17,11 @@
 
 import asyncio
 import contextvars
+
+# PYUPDATE: 3.14 release + 3.14 minimum: reaudit
+# heapq methods are not threadsafe pre 3.14
+# see: GH: cpython 135036
+# not currently affected, needs audit when making these loop-agnostic
 import heapq
 import threading
 from collections.abc import Generator
@@ -110,6 +115,7 @@ def __init_subclass__(cls) -> t.Never:
 
     __final__ = True
 
+    # When making this loop-agnostic later, heapq isn't threadsafe
     _loop: asyncio.AbstractEventLoop | None = None
 
     def __init__(self, value: int = 1) -> None: