85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
"""Retry helper shared by the Telegram API calls and the file downloads."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import TypeVar
|
|
|
|
import aiohttp
|
|
from telegram.error import BadRequest, Forbidden, InvalidToken, NetworkError, RetryAfter
|
|
|
|
from sticker_downloader.errors import FetchError
|
|
|
|
T = TypeVar("T")
|
|
|
|
#: Transient failures. ``BadRequest`` subclasses ``NetworkError`` in
|
|
#: python-telegram-bot, so it has to be excluded explicitly below — retrying a
|
|
#: "sticker set is invalid" response would just waste three round-trips.
|
|
_RETRIABLE = (NetworkError, aiohttp.ClientError, asyncio.TimeoutError, OSError)
|
|
_NEVER_RETRIABLE = (BadRequest, Forbidden, InvalidToken)
|
|
|
|
|
|
def is_retriable(error: BaseException) -> bool:
|
|
"""Whether another attempt could plausibly succeed."""
|
|
if isinstance(error, RetryAfter):
|
|
return True
|
|
if isinstance(error, _NEVER_RETRIABLE):
|
|
return False
|
|
if isinstance(error, FetchError):
|
|
return error.retriable
|
|
return isinstance(error, _RETRIABLE)
|
|
|
|
|
|
def retry_delay(error: BaseException, attempt: int, base_delay: float) -> float:
|
|
"""Seconds to wait before ``attempt`` (1-based) is retried.
|
|
|
|
Honours Telegram's own ``retry_after`` hint when it sends one, and otherwise
|
|
backs off exponentially.
|
|
"""
|
|
if isinstance(error, RetryAfter):
|
|
retry_after = getattr(error, "retry_after", None)
|
|
if isinstance(retry_after, (int, float)):
|
|
return float(retry_after)
|
|
# Newer python-telegram-bot versions may hand back a timedelta.
|
|
seconds = getattr(retry_after, "total_seconds", None)
|
|
if callable(seconds):
|
|
return float(seconds())
|
|
return base_delay * (2 ** (attempt - 1))
|
|
|
|
|
|
async def with_retries(
|
|
operation: Callable[[], Awaitable[T]],
|
|
*,
|
|
attempts: int = 3,
|
|
base_delay: float = 0.5,
|
|
on_retry: Callable[[BaseException, int, float], None] | None = None,
|
|
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
|
) -> T:
|
|
"""Await ``operation``, retrying transient failures up to ``attempts`` times.
|
|
|
|
Args:
|
|
operation: A zero-argument coroutine function; called afresh per attempt.
|
|
attempts: Total number of tries, including the first one.
|
|
base_delay: Seed for the exponential backoff, in seconds.
|
|
on_retry: Called with ``(error, attempt, delay)`` before each retry.
|
|
sleep: Injectable sleep, so tests do not have to wait.
|
|
|
|
Raises:
|
|
The last error, once the attempts are exhausted or it is not retriable.
|
|
"""
|
|
last_error: BaseException
|
|
for attempt in range(1, attempts + 1):
|
|
try:
|
|
return await operation()
|
|
except Exception as error:
|
|
last_error = error
|
|
if attempt == attempts or not is_retriable(error):
|
|
raise
|
|
delay = retry_delay(error, attempt, base_delay)
|
|
if on_retry is not None:
|
|
on_retry(error, attempt, delay)
|
|
await sleep(delay)
|
|
|
|
raise last_error # pragma: no cover — loop always returns or raises
|