✨ add set_thread_limit to concurrency module

fastapi/concurrency.py

@@ -63,6 +63,29 @@ async def _run_in_threadpool_with_overflow(
     return await _starlette_run_in_threadpool(func, *args, **kwargs)
 
 
+def set_thread_limit(limit: int = 40, anti_deadlock_reserve: int = 1) -> None:
+    """
+    Set the maximum number of threads that can be used by the thread pool.
+
+    This is a global setting that affects all calls to `run_in_threadpool` and
+    `iterate_in_threadpool`.
+    """
+    if not isinstance(limit, int):
+        raise TypeError("Thread limit must be an integer.")
+
+    if not isinstance(anti_deadlock_reserve, int):
+        raise TypeError("Anti deadlock reserve must be an integer.")
+
+    if limit < 2:
+        raise ValueError("Thread limit must be at least 2.")
+
+    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
+
+
 @asynccontextmanager
 async def contextmanager_in_threadpool(
     cm: AbstractContextManager[_T],

tests/test_concurrency.py

@@ -1,5 +1,7 @@
 from collections.abc import Iterator
+from typing import Any
 
+import anyio
 import pytest
 from fastapi import concurrency
 
@@ -28,3 +30,37 @@ async def test_iterate_in_threadpool(_reset_capacity_limiter: None) -> None:
         result.append(item)
 
     assert result == [0, 1, 2, 3, 4]
+
+
+@pytest.mark.parametrize(
+    "limit, anti_deadlock_reserve, error_kind",
+    [
+        ("not an int", 2, TypeError),
+        (10, "not an int", TypeError),
+        (1, 0, ValueError),
+        (10, 0, ValueError),
+        (10, 9, ValueError),
+    ],
+)
+def test_set_thread_limit_invalid_args(
+    limit: Any, anti_deadlock_reserve: Any, error_kind: type[Exception]
+) -> None:
+    with pytest.raises(error_kind):
+        concurrency.set_thread_limit(limit, anti_deadlock_reserve)
+
+
+@pytest.mark.anyio
+async def test_set_thread_limit(_reset_capacity_limiter: None) -> None:
+    original_total_tokens = (
+        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
+        )