♻️ use global thread limiter instead of anyio default

this is needed to prevent 3rd party libraries from blocking the extra thread we need to prevent deadlocks

fastapi/concurrency.py

@@ -1,53 +1,63 @@
+import functools
 from collections.abc import AsyncGenerator, Callable, Iterable
 from contextlib import AbstractContextManager
 from contextlib import asynccontextmanager as asynccontextmanager
-from typing import ParamSpec, TypeVar
+from typing import Literal, ParamSpec, TypeVar
 
 import anyio.to_thread
 from anyio import CapacityLimiter
 from starlette.concurrency import (
-    iterate_in_threadpool as _starlette_iterate_in_threadpool,
-)
-from starlette.concurrency import run_in_threadpool as _starlette_run_in_threadpool
-from starlette.concurrency import (
-    run_until_first_complete as _starlette_run_until_first_complete,
+    _next,
+    _StopIteration,
+    run_until_first_complete,  # noqa
 )
 
 _P = ParamSpec("_P")
 _T = TypeVar("_T")
 
-# The default threadpool in anyio is 40. This limiter keeps one thread
-# for teardown tasks in order to prevent deadlocks when there is a pool
-# of finite resources (e.g. database connections) which threads will block
-# on trying to acquire.
-# NOTE: we defer instantiation until runtime since we must support anyio's trio backend
+# A pair of limiters keep one thread for teardown tasks in order to prevent deadlocks
+# when there is a pool of finite resources (e.g. database connections) which threads will
+# block on trying to acquire.
+# NOTE: we cannot use anyio's default limiter, since other libraries will not respect the
+# anti-deadlock reserve and may use up all available threads - we maintain our own pool.
+_global_capacity_limiter: CapacityLimiter | None = None
 _anti_deadlock_capacity_limiter: CapacityLimiter | None = None
 
 
-def _get_anti_deadlock_capacity_limiter() -> CapacityLimiter:
-    global _anti_deadlock_capacity_limiter
-    if _anti_deadlock_capacity_limiter is None:
+def _get_capacity_limiter(kind: Literal["global", "anti_deadlock"]) -> CapacityLimiter:
+    global _global_capacity_limiter, _anti_deadlock_capacity_limiter
+    if _global_capacity_limiter is None:
         global_limit = anyio.to_thread.current_default_thread_limiter().total_tokens
+        _global_capacity_limiter = CapacityLimiter(global_limit)
         _anti_deadlock_capacity_limiter = CapacityLimiter(global_limit - 1)
-    return _anti_deadlock_capacity_limiter
+    return (
+        _anti_deadlock_capacity_limiter
+        if kind == "anti_deadlock"
+        else _global_capacity_limiter
+    )
 
 
+# These are vendored from starlette to allow setting our own thread limiter
 async def run_in_threadpool(
     func: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
 ) -> _T:
-    async with _get_anti_deadlock_capacity_limiter():
-        return await _starlette_run_in_threadpool(func, *args, **kwargs)
+    async with _get_capacity_limiter("anti_deadlock"):
+        func = functools.partial(func, *args, **kwargs)
+        return await anyio.to_thread.run_sync(
+            func, limiter=_get_capacity_limiter("global")
+        )
 
 
 async def iterate_in_threadpool(iterator: Iterable[_T]) -> AsyncGenerator[_T, None]:
-    async with _get_anti_deadlock_capacity_limiter():
-        async for item in _starlette_iterate_in_threadpool(iterator):
-            yield item
-
-
-async def run_until_first_complete(*args: tuple[Callable, dict]) -> None:  # type: ignore[type-arg]
-    async with _get_anti_deadlock_capacity_limiter():
-        return await _starlette_run_until_first_complete(*args)
+    async with _get_capacity_limiter("anti_deadlock"):
+        as_iterator = iter(iterator)
+        while True:
+            try:
+                yield await anyio.to_thread.run_sync(
+                    _next, as_iterator, limiter=_get_capacity_limiter("global")
+                )
+            except _StopIteration:
+                break
 
 
 # NOTE: a separate function is required only because mypy dislikes trying to add
@@ -60,7 +70,8 @@ async def _run_in_threadpool_with_overflow(
     Unless you know what you are doing you probably do not want to use this function.
     It has access to the entire thread pool, including the anti-deadlock reserve threads.
     """
-    return await _starlette_run_in_threadpool(func, *args, **kwargs)
+    func = functools.partial(func, *args, **kwargs)
+    return await anyio.to_thread.run_sync(func, limiter=_global_capacity_limiter)
 
 
 def set_thread_limit(limit: int = 40, anti_deadlock_reserve: int = 1) -> None:
@@ -82,8 +93,8 @@ def set_thread_limit(limit: int = 40, anti_deadlock_reserve: int = 1) -> None:
     if not 0 < anti_deadlock_reserve < limit - 1:
         raise ValueError("Anti deadlock reserve must be between 0 and limit - 1.")
 
-    anyio.to_thread.current_default_thread_limiter().total_tokens = limit
-    _get_anti_deadlock_capacity_limiter().total_tokens = limit - anti_deadlock_reserve
+    _get_capacity_limiter("global").total_tokens = limit
+    _get_capacity_limiter("anti_deadlock").total_tokens = limit - anti_deadlock_reserve
 
 
 @asynccontextmanager

tests/test_concurrency.py

@@ -10,8 +10,10 @@ from fastapi import concurrency
 def _reset_capacity_limiter() -> Iterator[None]:
     # Reset the capacity limiter before each test to ensure a clean slate
     concurrency._anti_deadlock_capacity_limiter = None
+    concurrency._global_capacity_limiter = None
     yield
     concurrency._anti_deadlock_capacity_limiter = None
+    concurrency._global_capacity_limiter = None
 
 
 @pytest.mark.anyio
@@ -51,16 +53,16 @@ def test_set_thread_limit_invalid_args(
 
 @pytest.mark.anyio
 async def test_set_thread_limit(_reset_capacity_limiter: None) -> None:
-    original_total_tokens = (
+    original_default_thread_limit = (
         anyio.to_thread.current_default_thread_limiter().total_tokens
     )
 
-    try:
-        concurrency.set_thread_limit(10, anti_deadlock_reserve=2)
-        assert concurrency._anti_deadlock_capacity_limiter.total_tokens == 8
-        assert anyio.to_thread.current_default_thread_limiter().total_tokens == 10
-    finally:
-        # Restore original settings to avoid affecting other tests
-        anyio.to_thread.current_default_thread_limiter().total_tokens = (
-            original_total_tokens
-        )
+    concurrency.set_thread_limit(10, anti_deadlock_reserve=2)
+    assert concurrency._anti_deadlock_capacity_limiter.total_tokens == 8
+    assert concurrency._global_capacity_limiter.total_tokens == 10
+
+    # default anyio token limiter not affected
+    assert (
+        anyio.to_thread.current_default_thread_limiter().total_tokens
+        == original_default_thread_limit
+    )