_throttling_request_manager.py

from __future__ import annotations

import asyncio
import contextlib
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from logging import getLogger
from typing import TYPE_CHECKING, Generic, Protocol, TypeVar

from typing_extensions import override
from yarl import URL

from crawlee._service_locator import ServiceLocator
from crawlee._service_locator import service_locator as global_service_locator
from crawlee._utils.docs import docs_group
from crawlee.request_loaders._request_manager import RequestManager

if TYPE_CHECKING:
    from collections.abc import Sequence

    from crawlee._request import Request
    from crawlee.configuration import Configuration
    from crawlee.storage_clients import StorageClient
    from crawlee.storage_clients.models import ProcessedRequest

logger = getLogger(__name__)

TRequestManager = TypeVar('TRequestManager', bound=RequestManager)

_NEVER_THROTTLED = datetime.min.replace(tzinfo=timezone.utc)
"""Sentinel `throttled_until` value meaning the domain has no active backoff."""


@docs_group('Request loaders')
class ThrottlingRequestManager(RequestManager, Generic[TRequestManager]):
    """A request manager that wraps another and enforces per-domain delays.

    Requests for explicitly configured domains are routed into dedicated sub-managers. A request added through this
    manager lives in exactly one of them, which keeps deduplication within a single store. A request that reached
    `inner` before its domain was configured stays there, and is fetched and completed against `inner` without the
    domain's delay applied.

    When `fetch_next_request()` is called, it returns requests from the sub-manager whose domain has been waiting the
    longest. If all configured domains are throttled, it falls back to the inner manager for non-throttled domains. If
    the inner manager is also empty and all sub-managers are throttled, it sleeps until the earliest cooldown expires.

    Delay sources:
    - HTTP 429 responses (via `record_domain_delay`)
    - robots.txt crawl-delay directives (via `set_crawl_delay`)

    The class is generic over the wrapped manager type. The first asynchronous operation - adding, fetching,
    completing, counting, purging, or dropping - makes the `request_manager_opener` callback open one sub-manager per
    configured domain, so every sub-manager shares the same `RequestManager` subclass and backing store as `inner`. The
    synchronous delay methods (`record_domain_delay`, `record_success`, `set_crawl_delay`) only touch in-memory state
    and never open anything. The opener must accept `alias`, `storage_client`, and `configuration` keyword arguments (as
    `RequestQueue.open` does) and return the same concrete subclass as `inner`.

    Opening the sub-managers up front also makes requests left over in a persistent store by a previous run visible
    again. With the default `purge_on_start=True` those leftovers are purged at open, so resuming them requires
    `purge_on_start=False`. Aliased stores are not exempt from that purge but named ones are, so a named `inner` keeps
    its requests across a restart while the per-domain stores are emptied.

    ### Usage

    ```python
    from crawlee.crawlers import BasicCrawler
    from crawlee.request_loaders import ThrottlingRequestManager
    from crawlee.storages import RequestQueue

    queue = await RequestQueue.open()
    throttler = ThrottlingRequestManager(
        inner=queue,
        domains=['api.example.com', 'slow-site.org'],
        request_manager_opener=RequestQueue.open,
    )
    crawler = BasicCrawler(request_manager=throttler)
    ```
    """

    def __init__(
        self,
        inner: TRequestManager,
        *,
        domains: Sequence[str],
        request_manager_opener: _RequestManagerOpener[TRequestManager],
        service_locator: ServiceLocator | None = None,
        base_delay: timedelta = timedelta(seconds=2),
        max_delay: timedelta = timedelta(seconds=60),
    ) -> None:
        """Initialize the throttling manager.

        Args:
            inner: The underlying request manager to wrap (typically a `RequestQueue`). Requests for non-throttled
                domains are stored here.
            domains: Explicit list of domain hostnames to throttle. Only requests matching these domains will be routed
                to per-domain sub-managers. Matching is case-insensitive (hostnames are lowercased) and exact: subdomain
                wildcards such as `*.example.com` are not supported — list each subdomain explicitly if needed.
            request_manager_opener: Async callable used to open one sub-manager per configured domain on first use.
                Must accept `alias`, `storage_client`, and `configuration` keyword arguments and return the same
                concrete subclass as `inner` (e.g. `RequestQueue.open` when `inner` is a `RequestQueue`).
            service_locator: Service locator for creating sub-managers. If not provided, defaults to the global service
                locator, ensuring consistency with the crawler's storage backend.
            base_delay: Initial delay after the first 429 response from a domain.
            max_delay: Maximum delay between requests to a rate-limited domain.
        """
        self._inner: TRequestManager = inner
        self._service_locator = service_locator if service_locator is not None else global_service_locator
        self._base_delay = base_delay
        self._max_delay = max_delay
        self._request_manager_opener = request_manager_opener
        self._domain_states: dict[str, _DomainState] = {d.lower(): _DomainState(domain=d.lower()) for d in domains if d}
        self._sub_managers: dict[str, TRequestManager] = {}
        self._sub_managers_ready = False
        self._sub_managers_lock = asyncio.Lock()
        self._in_flight_from_inner: set[tuple[str, str]] = set()
        """`(unique_key, url)` pairs of configured-domain requests handed out by `fetch_next_request` from the inner
        manager. Such a request can live in `inner` if it was added before its domain was listed, and it must be given
        back to the manager it came from. Requests for unconfigured domains need no record, as they route to `inner` by
        default. The URL is part of the key because `unique_key` may be set explicitly and is only unique per store, so
        a key alone could match a same-key request held by a sub-manager.

        The pair identifies a request by value, not by object. One URL can be in flight from both `inner` and its
        sub-manager at once - deduplication is per store, so both may hold it - and the two completions can then be
        routed to each other's manager. Both stores hold the key, so each completion still lands: the cost is a
        duplicate crawl of that URL and a retry that skips the domain's delay, not a stalled queue. Telling the two
        copies apart would take per-request identity, which `Request` cannot offer as it is unhashable and compares
        by value."""
        self._new_work_event = asyncio.Event()
        """Set whenever a request is added or reclaimed. Lets `fetch_next_request` wake from a throttle
        wait early when fresh work appears, instead of sleeping for the full computed cooldown."""

    @property
    def inner(self) -> TRequestManager:
        """The wrapped request manager that stores requests for non-throttled domains."""
        return self._inner

    @override
    async def drop(self) -> None:
        await self._ensure_sub_managers()
        await asyncio.gather(self._inner.drop(), *(sm.drop() for sm in self._sub_managers.values()))
        self._sub_managers.clear()
        self._sub_managers_ready = False
        self._in_flight_from_inner.clear()

    @override
    async def purge(self) -> None:
        """Empty the inner manager and all sub-managers, and reset transient per-domain throttle state.

        The configured domain list and any robots.txt-derived `crawl_delay` are preserved. Only the dynamic backoff
        state (consecutive 429 counter and `throttled_until`) is cleared. Sub-managers stay open; they're just emptied.
        """
        await self._ensure_sub_managers()
        await asyncio.gather(self._inner.purge(), *(sm.purge() for sm in self._sub_managers.values()))
        self._in_flight_from_inner.clear()
        for state in self._domain_states.values():
            state.consecutive_429_count = 0
            state.throttled_until = _NEVER_THROTTLED

    @override
    async def add_request(self, request: str | Request, *, forefront: bool = False) -> ProcessedRequest | None:
        """Add a request, routing it to the appropriate manager.

        Requests for explicitly configured domains are routed directly to their per-domain sub-manager. All other
        requests go to the inner manager.
        """
        await self._ensure_sub_managers()

        url = self._get_url_from_request(request)
        domain = self._extract_domain(url)

        if domain in self._domain_states:
            result = await self._sub_managers[domain].add_request(request, forefront=forefront)
        else:
            result = await self._inner.add_request(request, forefront=forefront)

        self._signal_new_work()
        return result

    @override
    async def add_requests(
        self,
        requests: Sequence[str | Request],
        *,
        forefront: bool = False,
        batch_size: int = 1000,
        wait_time_between_batches: timedelta = timedelta(seconds=1),
        wait_for_all_requests_to_be_added: bool = False,
        wait_for_all_requests_to_be_added_timeout: timedelta | None = None,
    ) -> None:
        """Add multiple requests, routing each to the appropriate manager."""
        await self._ensure_sub_managers()

        inner_requests: list[str | Request] = []
        domain_requests: dict[str, list[str | Request]] = {}

        for request in requests:
            url = self._get_url_from_request(request)
            domain = self._extract_domain(url)

            if domain in self._domain_states:
                domain_requests.setdefault(domain, []).append(request)
            else:
                inner_requests.append(request)

        if inner_requests:
            await self._inner.add_requests(
                inner_requests,
                forefront=forefront,
                batch_size=batch_size,
                wait_time_between_batches=wait_time_between_batches,
                wait_for_all_requests_to_be_added=wait_for_all_requests_to_be_added,
                wait_for_all_requests_to_be_added_timeout=wait_for_all_requests_to_be_added_timeout,
            )

        for domain, reqs in domain_requests.items():
            await self._sub_managers[domain].add_requests(
                reqs,
                forefront=forefront,
                batch_size=batch_size,
                wait_time_between_batches=wait_time_between_batches,
                wait_for_all_requests_to_be_added=wait_for_all_requests_to_be_added,
                wait_for_all_requests_to_be_added_timeout=wait_for_all_requests_to_be_added_timeout,
            )

        if inner_requests or domain_requests:
            self._signal_new_work()

    @override
    async def fetch_next_request(self) -> Request | None:
        """Fetch the next request, respecting per-domain delays.

        Sub-managers are checked in order of longest-overdue domain first (sorted by `throttled_until` ascending). If
        all configured domains are throttled, falls back to the inner manager for non-throttled domains. If the inner
        manager is also empty and all sub-managers are throttled, waits until either the earliest domain becomes
        available or new work is added (whichever comes first).
        """
        await self._ensure_sub_managers()

        while True:
            # Clear the event before checking the queues. Any add/reclaim that races with this iteration will set the
            # event again, so the wait at the end of the loop returns immediately rather than blocking until the
            # throttle expires.
            self._new_work_event.clear()

            now = datetime.now(timezone.utc)
            available_domains = sorted(
                (domain for domain, state in self._domain_states.items() if now >= state.throttled_until),
                key=lambda d: self._domain_states[d].throttled_until,
            )

            for domain in available_domains:
                req = await self._sub_managers[domain].fetch_next_request()
                if req:
                    self._mark_domain_dispatched(domain)
                    return req

            request = await self._inner.fetch_next_request()
            if request is not None:
                if self._extract_domain(request.url) in self._domain_states:
                    self._in_flight_from_inner.add((request.unique_key, request.url))
                return request

            sub_managers_empty = await asyncio.gather(*(sm.is_empty() for sm in self._sub_managers.values()))
            if all(sub_managers_empty):
                return None

            earliest = self._get_earliest_available_time(now)
            sleep_duration = max(
                (earliest - now).total_seconds(),
                0.1,  # Avoid tight loops if a throttle expired during the previous iteration.
            )
            logger.debug(
                f'All configured domains are throttled and inner manager is empty. '
                f'Waiting up to {sleep_duration:.1f}s for earliest domain to become available or new work.'
            )
            await self._wait_for_new_work_or_timeout(sleep_duration)

    @override
    async def reclaim_request(self, request: Request, *, forefront: bool = False) -> ProcessedRequest | None:
        await self._ensure_sub_managers()
        manager = self._fetch_owner(request)
        result = await manager.reclaim_request(request, forefront=forefront)
        self._clear_fetch_owner(request)
        self._signal_new_work()
        return result

    @override
    async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None:
        await self._ensure_sub_managers()
        manager = self._fetch_owner(request)
        result = await manager.mark_request_as_handled(request)
        self._clear_fetch_owner(request)
        self.record_success(request.url)
        return result

    @override
    async def get_handled_count(self) -> int:
        await self._ensure_sub_managers()
        counts = await asyncio.gather(
            self._inner.get_handled_count(), *(sm.get_handled_count() for sm in self._sub_managers.values())
        )
        return sum(counts)

    @override
    async def get_total_count(self) -> int:
        await self._ensure_sub_managers()
        counts = await asyncio.gather(
            self._inner.get_total_count(), *(sm.get_total_count() for sm in self._sub_managers.values())
        )
        return sum(counts)

    @override
    async def is_empty(self) -> bool:
        await self._ensure_sub_managers()
        results = await asyncio.gather(self._inner.is_empty(), *(sm.is_empty() for sm in self._sub_managers.values()))
        return all(results)

    @override
    async def is_finished(self) -> bool:
        await self._ensure_sub_managers()
        results = await asyncio.gather(
            self._inner.is_finished(), *(sm.is_finished() for sm in self._sub_managers.values())
        )
        return all(results)

    def record_domain_delay(self, url: str, *, retry_after: timedelta | None = None) -> bool:
        """Record a 429 Too Many Requests response for the domain of the given URL.

        Increments the consecutive 429 count and calculates the next allowed request time using exponential backoff or
        the `Retry-After` value.

        Args:
            url: The URL that received a 429 response.
            retry_after: Optional delay from the `Retry-After` header. If provided, it takes priority over the
                calculated exponential backoff.

        Returns:
            True if the URL's domain is configured for throttling and the delay was applied; False if the domain is not
            in the configured `domains` list, in which case the call is a no-op.
        """
        state = self._get_domain_state(url)
        if state is None:
            return False

        state.consecutive_429_count += 1
        delay = retry_after if retry_after is not None else self._base_delay * (2 ** (state.consecutive_429_count - 1))
        if delay > self._max_delay:
            source = 'Retry-After header' if retry_after is not None else 'exponential backoff'
            logger.warning(
                f'Capping {source} delay of {delay.total_seconds():.1f}s for domain "{state.domain}" '
                f'to max_delay ({self._max_delay.total_seconds():.1f}s); the domain may continue to rate-limit. '
                f'Consider increasing max_delay if this recurs.'
            )
            delay = self._max_delay
        state.throttled_until = datetime.now(timezone.utc) + delay

        logger.info(
            f'Rate limit (429) detected for domain "{state.domain}" '
            f'(consecutive: {state.consecutive_429_count}, delay: {delay.total_seconds():.1f}s)'
        )
        return True

    def record_success(self, url: str) -> None:
        """Record a successful request, resetting the backoff state for that domain.

        Args:
            url: The URL that received a successful response.
        """
        state = self._get_domain_state(url)
        if state is not None and state.consecutive_429_count > 0:
            logger.debug(f'Resetting rate limit state for domain "{state.domain}" after successful request')
            state.consecutive_429_count = 0

    def set_crawl_delay(self, url: str, delay_seconds: int) -> None:
        """Set the robots.txt crawl-delay for a domain.

        The delay is locked once set so robots.txt re-fetches (e.g. after LRU eviction) can't change the in-flight
        dispatch cadence and cause oscillation mid-crawl. Subsequent calls for the same domain are no-ops.

        Args:
            url: A URL from the domain to throttle.
            delay_seconds: The crawl-delay value in seconds.
        """
        state = self._get_domain_state(url)
        if state is None or state.crawl_delay is not None:
            return

        state.crawl_delay = timedelta(seconds=delay_seconds)
        logger.debug(f'Set crawl-delay for domain "{state.domain}" to {delay_seconds}s')

    @staticmethod
    def _extract_domain(url: str) -> str:
        """Extract the domain (hostname) from a URL."""
        return URL(url).host or ''

    @staticmethod
    def _get_url_from_request(request: str | Request) -> str:
        """Extract URL string from a request that may be a string or Request object."""
        return request if isinstance(request, str) else request.url

    def _get_domain_state(self, url: str) -> _DomainState | None:
        """Look up the per-domain state for the given URL, if the domain is configured."""
        domain = self._extract_domain(url)
        return self._domain_states.get(domain) if domain else None

    async def _open_sub_manager(self, domain: str) -> None:
        """Open the sub-manager for a single domain using the configured `request_manager_opener`."""
        self._sub_managers[domain] = await self._request_manager_opener(
            alias=f'throttled-{domain}',
            storage_client=self._service_locator.get_storage_client(),
            configuration=self._service_locator.get_configuration(),
        )

    async def _ensure_sub_managers(self) -> None:
        """Open a sub-manager for every configured domain, once.

        Sub-managers that opened before a sibling failed are kept, so a retry after a failure opens only what is
        still missing.
        """
        if self._sub_managers_ready:
            return

        async with self._sub_managers_lock:
            if self._sub_managers_ready:
                return

            # Every attempt has to settle before the lock is released. A propagating error would leave the remaining
            # openers running unawaited, free to write into `_sub_managers` after a retry has already replaced the
            # manager for that domain - stranding whatever the loser of that race holds.
            missing = [domain for domain in self._domain_states if domain not in self._sub_managers]
            results = await asyncio.gather(
                *(self._open_sub_manager(domain) for domain in missing), return_exceptions=True
            )
            for result in results:
                if isinstance(result, BaseException):
                    raise result

            self._sub_managers_ready = True

    def _is_domain_throttled(self, domain: str) -> bool:
        """Check if a domain is currently throttled."""
        state = self._domain_states.get(domain)
        if state is None:
            return False
        return datetime.now(timezone.utc) < state.throttled_until

    def _get_earliest_available_time(self, now: datetime) -> datetime:
        """Get the earliest time any throttled domain becomes available."""
        earliest = now + self._max_delay

        for state in self._domain_states.values():
            if now < state.throttled_until < earliest:
                earliest = state.throttled_until

        return earliest

    def _mark_domain_dispatched(self, domain: str) -> None:
        """Record that a request to this domain was just dispatched.

        If a crawl-delay is configured, push throttled_until forward by that amount.
        """
        state = self._domain_states.get(domain)
        if state is not None and state.crawl_delay is not None:
            state.throttled_until = datetime.now(timezone.utc) + state.crawl_delay

    def _signal_new_work(self) -> None:
        """Wake `fetch_next_request` if it is sleeping inside a throttle wait."""
        self._new_work_event.set()

    def _fetch_owner(self, request: Request) -> TRequestManager:
        """Return the manager the request must be given back to, leaving its in-flight record in place.

        The record is dropped by `_clear_fetch_owner` only once the owning manager has accepted the completion, so a
        completion retried after a transient storage failure still resolves to the same manager.
        """
        if (request.unique_key, request.url) in self._in_flight_from_inner:
            return self._inner
        return self._sub_managers.get(self._extract_domain(request.url), self._inner)

    def _clear_fetch_owner(self, request: Request) -> None:
        """Drop the in-flight record of a request whose completion the owning manager has accepted."""
        self._in_flight_from_inner.discard((request.unique_key, request.url))

    async def _wait_for_new_work_or_timeout(self, timeout: float) -> None:
        """Wait until new work is signaled or `timeout` seconds elapse, whichever comes first.

        The signal is set by `add_request`, `add_requests`, and `reclaim_request`, allowing `fetch_next_request` to wake
        up immediately when fresh work appears during a throttle wait instead of sleeping for the full computed cooldown
        (up to `max_delay`).
        """
        with contextlib.suppress(asyncio.TimeoutError):
            await asyncio.wait_for(self._new_work_event.wait(), timeout=timeout)


class _RequestManagerOpener(Protocol[TRequestManager]):
    """Callable that opens a `RequestManager` instance.

    Matches the keyword-only signature shared by storage `open` classmethods such as `RequestQueue.open`.
    `ThrottlingRequestManager` invokes the opener at sub-manager creation time, so every sub-manager shares the same
    backing type as `inner`.
    """

    async def __call__(
        self,
        *,
        alias: str | None = ...,
        storage_client: StorageClient | None = ...,
        configuration: Configuration | None = ...,
    ) -> TRequestManager: ...


@dataclass
class _DomainState:
    """Tracks delay state for a single domain."""

    domain: str
    """The domain being tracked."""

    throttled_until: datetime = _NEVER_THROTTLED
    """Earliest time the next request to this domain is allowed."""

    consecutive_429_count: int = 0
    """Number of consecutive 429 responses (for exponential backoff)."""

    crawl_delay: timedelta | None = None
    """Minimum interval between requests, used to push `throttled_until` on dispatch."""