diff --git a/matrix/api.py b/matrix/api.py index 34a5c57..43074bf 100644 --- a/matrix/api.py +++ b/matrix/api.py @@ -1,10 +1,12 @@ -from typing import Awaitable, TypeVar +import asyncio +from typing import Awaitable, Iterable, TypeVar from nio import ErrorResponse, Response from matrix.errors import MatrixError T = TypeVar("T", bound=Response) +R = TypeVar("R") async def matrix_call(coro: Awaitable[T], /, *, error_message: str) -> T: @@ -33,3 +35,35 @@ async def matrix_call(coro: Awaitable[T], /, *, error_message: str) -> T: raise MatrixError(f"{error_message}: {response}") return response + + +async def bounded_gather( + coros: Iterable[Awaitable[R]], + /, + *, + max_concurrent: int = 8, +) -> list[R]: + """Run `coros` concurrently, limited to `max_concurrent` at a time. + + Behaves like `asyncio.gather`, but never runs more than `max_concurrent` + awaitables at once. Useful for fanning out many API calls (e.g. + broadcasting to many rooms) without overwhelming a shared rate limit by + firing every request at once. The first exception raised by any + coroutine cancels the rest and propagates immediately, same as + `asyncio.gather`'s default behavior. + + ## Example + + ```python + messages = await bounded_gather( + (room.send(content) for room in rooms), max_concurrent=8 + ) + ``` + """ + semaphore = asyncio.Semaphore(max_concurrent) + + async def _run(coro: Awaitable[R]) -> R: + async with semaphore: + return await coro + + return await asyncio.gather(*[_run(coro) for coro in coros]) diff --git a/tests/test_api.py b/tests/test_api.py index fa1ed7f..1a2ad8c 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,8 +1,11 @@ +import asyncio +import inspect + import pytest from nio import RoomSendResponse, RoomSendError from matrix.errors import MatrixError -from matrix.api import matrix_call +from matrix.api import bounded_gather, matrix_call @pytest.mark.asyncio @@ -41,3 +44,55 @@ def test_matrix_call_requires_keyword_error_message__expect_type_error(): def test_matrix_call_requires_positional_coro__expect_type_error(): with pytest.raises(TypeError): matrix_call(coro=None, error_message="Failed to send message") + + +@pytest.mark.asyncio +async def test_bounded_gather_with_success__expect_results_in_order(): + async def value(n): + return n + + results = await bounded_gather([value(1), value(2), value(3)]) + + assert results == [1, 2, 3] + + +@pytest.mark.asyncio +async def test_bounded_gather_limits_concurrency__expect_never_exceeds_max_concurrent(): + active = 0 + peak = 0 + + async def task(): + nonlocal active, peak + active += 1 + peak = max(peak, active) + await asyncio.sleep(0.01) + active -= 1 + return "done" + + results = await bounded_gather([task() for _ in range(10)], max_concurrent=3) + + assert results == ["done"] * 10 + assert peak == 3 + + +@pytest.mark.asyncio +async def test_bounded_gather_with_exception__expect_it_propagates(): + async def ok(): + return "ok" + + async def fail(): + raise ValueError("boom") + + with pytest.raises(ValueError, match="boom"): + await bounded_gather([ok(), fail(), ok()]) + + +def test_bounded_gather_default_max_concurrent__expect_eight(): + default = inspect.signature(bounded_gather).parameters["max_concurrent"].default + + assert default == 8 + + +def test_bounded_gather_requires_positional_coros__expect_type_error(): + with pytest.raises(TypeError): + bounded_gather(coros=[])