Middleware¶

class RetryMiddleware:
    __slots__ = (
        '_backoff_factor',
        '_jitter',
        '_logger',
        '_max_attempts',
        '_max_delay',
        '_max_retry_after',
        '_respect_retry_after',
        '_retry_methods',
        '_status_codes',
    )

    def __init__(
        self,
        *,
        max_attempts: int = 3,
        backoff_factor: float = 0.5,
        max_delay: float = 30.0,
        respect_retry_after: bool = True,
        max_retry_after: float | None = None,
        jitter: bool = True,
        status_codes: Iterable[int] = (408, 429, 500, 502, 503, 504),
        retry_methods: Iterable[str] = (
            'GET',
            'HEAD',
            'OPTIONS',
            'PUT',
            'DELETE',
        ),
        logger: Logger | None = None,
    ) -> None:
        super().__init__()
        self._logger = logger or getLogger('oxidehttp-retry')
        self._max_attempts = max_attempts
        self._backoff_factor = backoff_factor
        self._max_delay = max_delay
        self._respect_retry_after = respect_retry_after
        self._max_retry_after = max_retry_after
        self._jitter = jitter
        self._status_codes = frozenset(status_codes)
        self._retry_methods = frozenset(retry_methods)

        if self._max_attempts < 1:
            exc = 'max_attempts must be at least 1'
            raise ValueError(exc)

        if self._backoff_factor < 0:
            exc = 'backoff_factor cannot be negative'
            raise ValueError(exc)

        if self._max_delay <= 0:
            exc = 'max_delay must be greater than zero'
            raise ValueError(exc)

    @staticmethod
    def _get_jittered_delay(delay: float) -> float:
        if delay <= 0:
            return 0.0

        # 1 ms resolution.
        upper = max(1, int(delay * 1000))
        return randbelow(upper + 1) / 1000

    def _should_retry(
        self, response: Response, method: str, attempt: int
    ) -> bool:
        if attempt >= self._max_attempts:
            return False

        if method not in self._retry_methods:
            return False

        return response.status in self._status_codes

    def _parse_retry_after(self, response: Response) -> float | None:
        value = response.headers.get('Retry-After')

        if value is None:
            return None

        value = value.strip()

        if not value:
            return None

        try:
            return float(int(value))

        except ValueError:
            pass

        try:
            retry_at = parsedate_to_datetime(value).replace(tzinfo=UTC)

        except TypeError, ValueError, OverflowError:
            return None

        delay = (retry_at - datetime.now(UTC)).total_seconds()

        return max(0.0, delay)

    def _get_retry_delay(
        self, attempt: int, response: Response | None = None
    ) -> float:
        if self._respect_retry_after and response is not None:
            retry_after = self._parse_retry_after(response)

            if retry_after is not None:
                if self._max_retry_after is not None:
                    return min(retry_after, self._max_retry_after)

                return retry_after

        exponential = float(self._backoff_factor * (2 ** (attempt - 1)))
        delay = min(exponential, self._max_delay)

        if self._jitter:
            delay = self._get_jittered_delay(delay)

        return delay

    async def run(self, request: Request, next_req: Next) -> Response:
        method = request.method

        for attempt in range(1, self._max_attempts + 1):
            try:
                response = await next_req.run(request)

            except PyreqwestError as e:
                if attempt >= self._max_attempts:
                    raise

                delay = self._get_retry_delay(attempt)
                await sleep(delay)
                self._logger.warning(
                    'Retrying HTTP request',
                    extra={
                        'method': method,
                        'url': str(request.url),
                        'attempt': attempt,
                        'error': str(e),
                        'delay': delay,
                    },
                )
                continue

            if not self._should_retry(response, method, attempt):
                return response

            delay = self._get_retry_delay(attempt, response)
            await sleep(delay)
            self._logger.warning(
                'Retrying HTTP request',
                extra={
                    'method': method,
                    'url': str(request.url),
                    'attempt': attempt,
                    'delay': delay,
                },
            )

        exc = 'Retry loop exited unexpectedly'
        raise RuntimeError(exc)

class RequestIDMiddleware:
    __slots__ = ('_logger', '_request_id_factory', '_request_id_header')

    def __init__(
        self,
        *,
        request_id_header: str | None = None,
        request_id_factory: Callable[[], str] | None = None,
        logger: Logger | None = None,
    ) -> None:
        super().__init__()
        self._logger = logger or getLogger('oxidehttp-requestid')
        self._request_id_header = request_id_header or 'X-Request-ID'
        self._request_id_factory = (
            request_id_factory or self.request_id_factory
        )

    def request_id_factory(self) -> str:
        request_id = str(uuid4())

        self._logger.debug('Generated Request-ID: %s', request_id)

        return request_id

    async def run(self, request: Request, next_req: Next) -> Response:
        extensions = cast('Extensions', request.extensions)
        request_id = extensions.get('request_id')

        if request_id is None:
            request_id = self._request_id_factory()

        request.headers.setdefault(self._request_id_header, request_id)
        response = await next_req.run(request)
        response.headers.setdefault(self._request_id_header, request_id)

        return response

class RateLimitMiddleware:
    __slots__ = ('_backend', '_limit', '_period', '_store')

    def __init__(
        self, store: RateLimitStore, *, limit: int, period: float
    ) -> None:
        super().__init__()

        if limit < 1:
            exc = 'limit must be at least 1'
            raise ValueError(exc)

        if period <= 0:
            exc = 'period must be greater than zero'
            raise ValueError(exc)

        self._store = store
        self._limit = limit
        self._period = period

    def _key(self, request: Request, extension_key: str | None) -> str:
        keys: list[str] = []

        if request.url.host_str is not None:
            keys.append(request.url.host_str)

        if extension_key is not None:
            keys.append(extension_key)

        return ':'.join(keys)

    async def run(self, request: Request, next_req: Next) -> Response:
        extensions = cast('Extensions', request.extensions)

        extension_key = extensions.get('rateLimitKey')
        key = self._key(request, extension_key)

        await self._store.acquire(key, limit=self._limit, period=self._period)

        return await next_req.run(request)

class RateLimitStore(Protocol):
    @abstractmethod
    async def acquire(self, key: str, *, limit: int, period: float) -> float:
        pass

class MemoryRateLimitStore:
    __slots__ = ('_lock', '_windows')

    def __init__(self) -> None:
        self._lock = Lock()
        self._windows: dict[str, tuple[float, int]] = {}

    async def acquire(self, key: str, *, limit: int, period: float) -> float:
        waited = 0.0

        while True:
            async with self._lock:
                now = monotonic()
                window = self._windows.get(key)

                if window is None or now >= window[0]:
                    self._windows[key] = (now + period, 1)
                    return waited

                expires_at, count = window

                if count < limit:
                    self._windows[key] = (expires_at, count + 1)
                    return waited

                delay = expires_at - now

            await sleep(delay)
            waited += delay

class RedisRateLimitStore:
    __slots__ = ('_prefix', '_redis')

    def __init__(
        self, redis: Redis, *, prefix: str = 'oxidehttp:ratelimit:'
    ) -> None:
        self._redis = redis
        self._prefix = prefix

    async def acquire(self, key: str, *, limit: int, period: float) -> float:
        script = """
        local current = redis.call('INCR', KEYS[1])

        if current == 1 then
            redis.call('PEXPIRE', KEYS[1], ARGV[1])
        end

        if current <= tonumber(ARGV[2]) then
            return 0
        end

        local ttl = redis.call('PTTL', KEYS[1])
        return ttl
        """

        redis_key = f'{self._prefix}{key}'
        period_ms = max(1, int(period * 1000))

        while True:
            delay_ms = await self._redis.eval(
                script, 1, redis_key, period_ms, limit
            )

            if delay_ms == 0:
                return 0.0

            await sleep(delay_ms / 1000)