Python 데코레이터로 async 재시도·타임아웃 공용화

import asyncio
import functools
import random
import time
from dataclasses import dataclass
from typing import Awaitable, Callable, Optional, Tuple, Type


@dataclass(frozen=True)
class RetryPolicy:
    attempts: int = 3
    timeout_s: float = 2.0
    base_delay_s: float = 0.2
    max_delay_s: float = 2.0
    backoff: float = 2.0
    jitter_s: float = 0.1
    retry_on: Tuple[Type[BaseException], ...] = (TimeoutError, OSError)
    # 주의: CancelledError는 절대 retry_on에 넣지 않습니다.


def async_retry_timeout(
    policy: RetryPolicy,
    *,
    on_retry: Optional[Callable[[int, BaseException, float], None]] = None,
) -> Callable[[Callable[..., Awaitable]], Callable[..., Awaitable]]:
    """async 함수에 재시도 + 타임아웃을 공용화하는 데코레이터.

    on_retry: (attempt_index, exception, next_delay_s) 콜백
    """

    def decorator(fn: Callable[..., Awaitable]):
        @functools.wraps(fn)
        async def wrapper(*args, **kwargs):
            last_exc: Optional[BaseException] = None

            for i in range(1, policy.attempts + 1):
                try:
                    # 각 시도마다 timeout 적용
                    return await asyncio.wait_for(
                        fn(*args, **kwargs),
                        timeout=policy.timeout_s,
                    )
                except asyncio.CancelledError:
                    # 상위 취소 전파는 무조건 보장
                    raise
                except Exception as exc:
                    last_exc = exc

                    should_retry = isinstance(exc, policy.retry_on) and i < policy.attempts
                    if not should_retry:
                        raise

                    # 지수 백오프 + 지터
                    exp = policy.base_delay_s * (policy.backoff ** (i - 1))
                    delay = min(policy.max_delay_s, exp) + random.uniform(0, policy.jitter_s)

                    if on_retry:
                        on_retry(i, exc, delay)

                    await asyncio.sleep(delay)

            # 논리상 도달하지 않지만, 타입 안정성 차원에서 남김
            assert last_exc is not None
            raise last_exc

        return wrapper

    return decorator

import aiohttp

policy = RetryPolicy(
    attempts=4,
    timeout_s=1.5,
    retry_on=(TimeoutError, aiohttp.ClientConnectionError, OSError),
)


def log_retry(i: int, exc: BaseException, delay_s: float) -> None:
    print(f"retry={i} exc={type(exc).__name__} delay_s={delay_s:.2f}")


@async_retry_timeout(policy, on_retry=log_retry)
async def fetch_text(url: str) -> str:
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            resp.raise_for_status()
            return await resp.text()


async def main():
    text = await fetch_text("https://example.com")
    print(text[:80])

policy = RetryPolicy(
    attempts=3,
    timeout_s=2.0,
    retry_on=(asyncio.TimeoutError, TimeoutError, OSError),
)

from typing import Protocol


class RetryIf(Protocol):
    def __call__(self, exc: BaseException) -> bool: ...


def async_retry_timeout2(
    policy: RetryPolicy,
    *,
    retry_if: Optional[RetryIf] = None,
):
    def decorator(fn):
        @functools.wraps(fn)
        async def wrapper(*args, **kwargs):
            for i in range(1, policy.attempts + 1):
                try:
                    return await asyncio.wait_for(fn(*args, **kwargs), timeout=policy.timeout_s)
                except asyncio.CancelledError:
                    raise
                except Exception as exc:
                    type_ok = isinstance(exc, policy.retry_on)
                    cond_ok = retry_if(exc) if retry_if else True
                    if not (type_ok and cond_ok and i < policy.attempts):
                        raise
                    exp = policy.base_delay_s * (policy.backoff ** (i - 1))
                    delay = min(policy.max_delay_s, exp) + random.uniform(0, policy.jitter_s)
                    await asyncio.sleep(delay)

        return wrapper

    return decorator

async def do_work():
    return await fetch_text("https://example.com")


async def main():
    # 전체 작업은 5초 내로 끝내기
    result = await asyncio.wait_for(do_work(), timeout=5.0)
    print(result[:50])

import logging

logger = logging.getLogger("retry")


def on_retry(i: int, exc: BaseException, delay_s: float) -> None:
    logger.warning(
        "retrying",
        extra={
            "attempt": i,
            "exc_type": type(exc).__name__,
            "delay_s": delay_s,
        },
    )

from typing import Callable, Awaitable


def async_retry_timeout_testable(
    policy: RetryPolicy,
    *,
    sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
):
    def decorator(fn):
        @functools.wraps(fn)
        async def wrapper(*args, **kwargs):
            for i in range(1, policy.attempts + 1):
                try:
                    return await asyncio.wait_for(fn(*args, **kwargs), timeout=policy.timeout_s)
                except asyncio.CancelledError:
                    raise
                except Exception as exc:
                    if not (isinstance(exc, policy.retry_on) and i < policy.attempts):
                        raise
                    exp = policy.base_delay_s * (policy.backoff ** (i - 1))
                    delay = min(policy.max_delay_s, exp)
                    await sleep(delay)

        return wrapper

    return decorator

import pytest


async def noop_sleep(_: float) -> None:
    return None


@pytest.mark.asyncio
async def test_retry_succeeds_after_two_failures():
    calls = {"n": 0}

    policy = RetryPolicy(attempts=3, timeout_s=0.5, retry_on=(ValueError,))

    @async_retry_timeout_testable(policy, sleep=noop_sleep)
    async def flaky():
        calls["n"] += 1
        if calls["n"] < 3:
            raise ValueError("temporary")
        return "ok"

    assert await flaky() == "ok"
    assert calls["n"] == 3