🐛 fix deadlock with dependencies

fastapi/concurrency.py

@@ -1,18 +1,67 @@
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable, Iterable
 from contextlib import AbstractContextManager
 from contextlib import asynccontextmanager as asynccontextmanager
-from typing import TypeVar
+from typing import ParamSpec, TypeVar
 
 import anyio.to_thread
 from anyio import CapacityLimiter
-from starlette.concurrency import iterate_in_threadpool as iterate_in_threadpool  # noqa
-from starlette.concurrency import run_in_threadpool as run_in_threadpool  # noqa
-from starlette.concurrency import (  # noqa
-    run_until_first_complete as run_until_first_complete,
+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,
 )
 
+_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
+_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:
+        global_limit = anyio.to_thread.current_default_thread_limiter().total_tokens
+        _anti_deadlock_capacity_limiter = CapacityLimiter(global_limit - 1)
+    return _anti_deadlock_capacity_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 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)
+
+
+# NOTE: a separate function is required only because mypy dislikes trying to add
+# a boolean flag along side the param spec
+async def _run_in_threadpool_with_overflow(
+    func: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
+) -> _T:
+    """Run a function in the thread pool, allowing it to use the overflow threads.
+
+    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)
+
 
 @asynccontextmanager
 async def contextmanager_in_threadpool(

fastapi/dependencies/utils.py

@@ -53,6 +53,7 @@ from fastapi.background import BackgroundTasks
 from fastapi.concurrency import (
     asynccontextmanager,
     contextmanager_in_threadpool,
+    run_in_threadpool,
 )
 from fastapi.dependencies.models import Dependant
 from fastapi.exceptions import DependencyScopeError
@@ -63,7 +64,6 @@ from fastapi.utils import create_model_field, get_path_param_names
 from pydantic import BaseModel, Json
 from pydantic.fields import FieldInfo
 from starlette.background import BackgroundTasks as StarletteBackgroundTasks
-from starlette.concurrency import run_in_threadpool
 from starlette.datastructures import (
     FormData,
     Headers,

fastapi/routing.py

@@ -38,6 +38,11 @@ from fastapi._compat import (
     Undefined,
     lenient_issubclass,
 )
+from fastapi.concurrency import (
+    _run_in_threadpool_with_overflow,
+    iterate_in_threadpool,
+    run_in_threadpool,
+)
 from fastapi.datastructures import Default, DefaultPlaceholder
 from fastapi.dependencies.models import Dependant
 from fastapi.dependencies.utils import (
@@ -75,7 +80,6 @@ from fastapi.utils import (
 from starlette import routing
 from starlette._exception_handler import wrap_app_handling_exceptions
 from starlette._utils import is_async_callable
-from starlette.concurrency import iterate_in_threadpool, run_in_threadpool
 from starlette.datastructures import FormData
 from starlette.exceptions import HTTPException
 from starlette.requests import Request
@@ -292,7 +296,7 @@ async def serialize_response(
         if is_coroutine:
             value, errors = field.validate(response_content, {}, loc=("response",))
         else:
-            value, errors = await run_in_threadpool(
+            value, errors = await _run_in_threadpool_with_overflow(
                 field.validate, response_content, {}, loc=("response",)
             )
         if errors:

tests/test_concurrency.py

@@ -0,0 +1,30 @@
+from collections.abc import Iterator
+
+import pytest
+from fastapi import concurrency
+
+
+@pytest.fixture
+def _reset_capacity_limiter() -> Iterator[None]:
+    # Reset the capacity limiter before each test to ensure a clean slate
+    concurrency._anti_deadlock_capacity_limiter = None
+    yield
+    concurrency._anti_deadlock_capacity_limiter = None
+
+
+@pytest.mark.anyio
+async def test_run_in_threadpool(_reset_capacity_limiter: None) -> None:
+    def blocking_function(x: int, y: int) -> int:
+        return x + y
+
+    result = await concurrency.run_in_threadpool(blocking_function, 1, y=2)
+    assert result == 3
+
+
+@pytest.mark.anyio
+async def test_iterate_in_threadpool(_reset_capacity_limiter: None) -> None:
+    result = []
+    async for item in concurrency.iterate_in_threadpool(range(5)):
+        result.append(item)
+
+    assert result == [0, 1, 2, 3, 4]

tests/test_depends_deadlock.py

@@ -0,0 +1,56 @@
+import asyncio
+import threading
+import time
+from collections.abc import Iterator
+
+from fastapi import Depends, FastAPI
+from httpx import ASGITransport, AsyncClient
+from pydantic import BaseModel
+
+# Mutex, acting as our "connection pool" for a database for example
+mutex_pool = threading.Lock()
+
+
+# Simulate releaasing a pooled resource in the teardown of a Depends,
+# which in reality is usually a database connection or similar.
+def release_resource() -> Iterator[None]:
+    try:
+        time.sleep(0.001)
+        yield
+    finally:
+        time.sleep(0.001)
+        mutex_pool.release()
+
+
+app = FastAPI()
+
+
+class Item(BaseModel):
+    name: str
+    id: int
+
+
+# An endpoint that uses Depends for resource management and also includes
+# a response_model definition would previously deadlock in the validation
+# of the model and the cleanup of the Depends
+@app.get("/deadlock", response_model=Item)
+def get_deadlock(dep: None = Depends(release_resource)):
+    mutex_pool.acquire()
+    return Item(name="foo", id=1)
+
+
+# Fire off 100 requests in parallel(ish) in order to create contention
+# over the shared resource (simulating a fastapi server that interacts with
+# a database connection pool).
+def test_depends_deadlock():
+    async def make_request(client: AsyncClient):
+        await client.get("/deadlock")
+
+    async def run_requests():
+        async with AsyncClient(
+            transport=ASGITransport(app=app), base_url="http://testserver"
+        ) as aclient:
+            tasks = [make_request(aclient) for _ in range(100)]
+            await asyncio.gather(*tasks)
+
+    asyncio.run(run_requests())