async def throttled(
*,
throttler: containers.Throttler,
delays: Iterable[float],
wakeup: Optional[Union[asyncio.Event, primitives.DaemonStopper]] = None,
logger: Union[logging.Logger, logging.LoggerAdapter],
errors: Union[Type[BaseException], Tuple[Type[BaseException], ...]] = Exception,
) -> AsyncGenerator[bool, None]:
"""
A helper to throttle any arbitrary operation.
"""
# The 1st sleep: if throttling is already active, but was interrupted by a queue replenishment.
# It is needed to properly process the latest known event after the successful sleep.
if throttler.active_until is not None:
remaining_time = throttler.active_until - time.monotonic()
unslept_time = await sleep_or_wait(remaining_time, wakeup=wakeup)
if unslept_time is None:
logger.info("Throttling is over. Switching back to normal operations.")
throttler.active_until = None
# Run only if throttling either is not active initially, or has just finished sleeping.
should_run = throttler.active_until is None
try:
yield should_run
except Exception as e:
# If it is not an error-of-interest, escalate normally. BaseExceptions are escalated always.
if not isinstance(e, errors):
raise
# If the code does not follow the recommendation to not run, escalate.
if not should_run:
raise
# Activate throttling if not yet active, or reuse the active sequence of delays.
if throttler.source_of_delays is None:
throttler.source_of_delays = iter(delays)
# Choose a delay. If there are none, avoid throttling at all.
throttle_delay = next(throttler.source_of_delays, throttler.last_used_delay)
if throttle_delay is not None:
throttler.last_used_delay = throttle_delay
throttler.active_until = time.monotonic() + throttle_delay
logger.exception(f"Throttling for {throttle_delay} seconds due to an unexpected error:")
else:
# Reset the throttling. Release the iterator to keep the memory free during normal run.
if should_run:
throttler.source_of_delays = throttler.last_used_delay = None
# The 2nd sleep: if throttling has been just activated (i.e. there was a fresh error).
# It is needed to have better logging/sleeping without workers exiting for "no events".
if throttler.active_until is not None and should_run:
remaining_time = throttler.active_until - time.monotonic()
unslept_time = await sleep_or_wait(remaining_time, wakeup=wakeup)
if unslept_time is None:
throttler.active_until = None
logger.info("Throttling is over. Switching back to normal operations.")
async def test_activates_on_expected_errors(exc_cls, kwargs):
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[123], **kwargs):
raise exc_cls()
assert throttler.source_of_delays is not None
assert throttler.last_used_delay is not None
async def test_remains_inactive_on_success():
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[123]):
pass
assert throttler.source_of_delays is None
assert throttler.last_used_delay is None
async def test_recommends_running_initially():
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger,
delays=[123]) as should_run:
remembered_should_run = should_run
assert remembered_should_run is True
async def test_continuation_when_overdue(clock, sleep):
wakeup = asyncio.Event()
logger = logging.getLogger()
throttler = Throttler()
clock.return_value = 1000 # simulated "now"
sleep.return_value = 55 # simulated sleep time left
async with throttled(throttler=throttler,
logger=logger,
delays=[123, 234],
wakeup=wakeup):
raise Exception()
sleep.reset_mock()
clock.return_value = 2000 # simulated "now"
sleep.return_value = None # simulated sleep time left
async with throttled(throttler=throttler,
logger=logger,
delays=[...],
wakeup=wakeup):
raise Exception()
assert throttler.last_used_delay == 234
assert throttler.source_of_delays is not None
assert throttler.active_until is None # means: no sleep time is left
assert sleep.mock_calls == [
call(123 - 1000, wakeup=wakeup),
call(234, wakeup=wakeup)
]
async def test_escalates_unexpected_errors(exc_cls, kwargs):
logger = logging.getLogger()
throttler = Throttler()
with pytest.raises(exc_cls):
async with throttled(throttler=throttler,
logger=logger,
delays=[123],
**kwargs):
raise exc_cls()
async def test_logging_when_deactivates_immediately(caplog):
caplog.set_level(0)
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[123]):
raise Exception()
assert caplog.messages == [
"Throttling for 123 seconds due to an unexpected error:",
"Throttling is over. Switching back to normal operations.",
]
async def test_sleeps_for_the_first_delay_when_inactive(sleep):
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[123, 234]):
raise Exception()
assert throttler.last_used_delay == 123
assert throttler.source_of_delays is not None
assert next(throttler.source_of_delays) == 234
assert throttler.active_until is None # means: no sleep time left
assert sleep.mock_calls == [call(123, wakeup=None)]
async def test_recommends_running_immediately_after_continued(sleep):
logger = logging.getLogger()
throttler = Throttler()
sleep.return_value = 33 # simulated sleep time left
async with throttled(throttler=throttler, logger=logger, delays=[123]):
raise Exception()
sleep.return_value = None # simulated sleep time left
async with throttled(throttler=throttler, logger=logger, delays=[...]) as should_run:
remembered_should_run = should_run
assert remembered_should_run is True
async def test_recommends_skipping_immediately_after_interrupted_error(sleep):
logger = logging.getLogger()
throttler = Throttler()
sleep.return_value = 33 # simulated sleep time left
async with throttled(throttler=throttler, logger=logger, delays=[123]):
raise Exception()
sleep.return_value = 33 # simulated sleep time left
async with throttled(throttler=throttler, logger=logger, delays=[...]) as should_run:
remembered_should_run = should_run
assert remembered_should_run is False
async def test_skips_on_no_delays(sleep):
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[]):
raise Exception()
assert throttler.last_used_delay is None
assert throttler.source_of_delays is not None
assert next(throttler.source_of_delays, 999) == 999
assert throttler.active_until is None # means: no sleep time left
assert sleep.mock_calls == []
async def test_interruption(clock, sleep):
wakeup = asyncio.Event()
logger = logging.getLogger()
throttler = Throttler()
clock.return_value = 1000 # simulated "now"
sleep.return_value = 55 # simulated sleep time left
async with throttled(throttler=throttler, logger=logger, delays=[123, 234], wakeup=wakeup):
raise Exception()
assert throttler.last_used_delay == 123
assert throttler.source_of_delays is not None
assert throttler.active_until == 1123 # means: some sleep time is left
assert sleep.mock_calls == [call(123, wakeup=wakeup)]
async def test_resets_on_success(sleep):
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[123]):
raise Exception()
sleep.reset_mock()
async with throttled(throttler=throttler, logger=logger, delays=[...]):
pass
assert throttler.last_used_delay is None
assert throttler.source_of_delays is None
assert throttler.active_until is None
assert sleep.mock_calls == []
async def test_logging_when_deactivates_on_reentry(sleep, caplog):
caplog.set_level(0)
logger = logging.getLogger()
throttler = Throttler()
sleep.return_value = 55 # simulated sleep time left
async with throttled(throttler=throttler, logger=logger, delays=[123]):
raise Exception()
sleep.return_value = None # simulated sleep time left
async with throttled(throttler=throttler, logger=logger, delays=[...]):
pass
assert caplog.messages == [
"Throttling for 123 seconds due to an unexpected error:",
"Throttling is over. Switching back to normal operations.",
]
async def test_renews_on_repeated_failure(sleep):
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[123]):
raise Exception()
async with throttled(throttler=throttler, logger=logger, delays=[...]):
pass
sleep.reset_mock()
async with throttled(throttler=throttler, logger=logger, delays=[234]):
raise Exception()
assert throttler.last_used_delay is 234
assert throttler.source_of_delays is not None
assert throttler.active_until is None
assert sleep.mock_calls == [call(234, wakeup=None)]
async def test_sleeps_for_the_last_known_delay_when_depleted(sleep):
logger = logging.getLogger()
throttler = Throttler()
async with throttled(throttler=throttler, logger=logger, delays=[123, 234]):
raise Exception()
async with throttled(throttler=throttler, logger=logger, delays=[...]):
raise Exception()
sleep.reset_mock()
async with throttled(throttler=throttler, logger=logger, delays=[...]):
raise Exception()
assert throttler.last_used_delay == 234
assert throttler.source_of_delays is not None
assert next(throttler.source_of_delays, 999) == 999
assert throttler.active_until is None # means: no sleep time left
assert sleep.mock_calls == [call(234, wakeup=None)]
