diff --git a/aio_request/deadline_provider.py b/aio_request/deadline_provider.py index 8ba53eb..0facf64 100644 --- a/aio_request/deadline_provider.py +++ b/aio_request/deadline_provider.py @@ -5,7 +5,7 @@ DeadlineProvider = Callable[[Deadline, int, int], Deadline] -def split_deadline_between_attempts() -> DeadlineProvider: +def split_deadline_between_attempts(split_factor: int | None = None) -> DeadlineProvider: """ Split deadline between attempts. @@ -18,15 +18,19 @@ def split_deadline_between_attempts() -> DeadlineProvider: the last one has received the remaining 8 seconds due to redistribution. """ + if split_factor is not None and split_factor < 2: + raise ValueError("max_split should be greater or equal to 2") + def __provider(deadline: Deadline, attempt: int, attempts_count: int) -> Deadline: if deadline.expired: return deadline - - attempts_left = attempts_count - attempt - if attempts_left == 0: - raise ValueError("no attempts left") - - return deadline / attempts_left + if split_factor is None: + effective_split_factor = attempts_count - attempt + else: + effective_split_factor = min(split_factor, attempts_count) - attempt + if effective_split_factor <= 1: + return deadline + return deadline / effective_split_factor return __provider diff --git a/tests/test_deadline_provider.py b/tests/test_deadline_provider.py index c583fbb..9fcefbd 100644 --- a/tests/test_deadline_provider.py +++ b/tests/test_deadline_provider.py @@ -10,11 +10,29 @@ async def test_split_deadline_between_attempt(): attempt_deadline = provider(deadline, 0, 3) assert 0.3 <= attempt_deadline.timeout <= 0.34 - await asyncio.sleep(attempt_deadline.timeout) + await asyncio.sleep(0.33) attempt_deadline = provider(deadline, 1, 3) assert 0.3 <= attempt_deadline.timeout <= 0.34 - await asyncio.sleep(attempt_deadline.timeout) + await asyncio.sleep(0.33) + + attempt_deadline = provider(deadline, 2, 3) + assert 0.3 <= attempt_deadline.timeout <= 0.34 + + +async def test_split_deadline_between_attempt_with_split_factor(): + provider = aio_request.split_deadline_between_attempts(split_factor=2) + deadline = aio_request.Deadline.from_timeout(1) + + attempt_deadline = provider(deadline, 0, 3) + assert 0.45 <= attempt_deadline.timeout <= 0.5 + + await asyncio.sleep(0.33) + + attempt_deadline = provider(deadline, 1, 3) + assert 0.6 <= attempt_deadline.timeout <= 0.67 + + await asyncio.sleep(0.33) attempt_deadline = provider(deadline, 2, 3) assert 0.3 <= attempt_deadline.timeout <= 0.34 @@ -35,3 +53,21 @@ async def test_split_deadline_between_attempts_fast_attempt_failure(): attempt_deadline = provider(deadline, 2, 3) assert 0.75 <= attempt_deadline.timeout <= 0.8 + + +async def test_split_deadline_between_attempts_fast_attempt_failure_with_split_factor(): + provider = aio_request.split_deadline_between_attempts(split_factor=2) + deadline = aio_request.Deadline.from_timeout(1) + + attempt_deadline = provider(deadline, 0, 3) + assert 0.45 <= attempt_deadline.timeout <= 0.5 + + await asyncio.sleep(0.1) # fast attempt failure + + attempt_deadline = provider(deadline, 1, 3) + assert 0.85 <= attempt_deadline.timeout <= 0.9 + + await asyncio.sleep(0.1) # fast attempt failure + + attempt_deadline = provider(deadline, 2, 3) + assert 0.75 <= attempt_deadline.timeout <= 0.8