perf(impit): optimize client and connection reuse across Sessions

pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "async-timeout>=5.0.1",
     "cachetools>=5.5.0",
     "colorama>=0.4.0",
-    "impit>=0.8.0",
+    "impit>=0.13.2",
     "more-itertools>=10.2.0",
     "protego>=0.5.0",
     "psutil>=6.0.0",

src/crawlee/http_clients/_impit.py

@@ -2,13 +2,15 @@
 
 import asyncio
 from contextlib import asynccontextmanager
+from http import HTTPStatus
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, TypedDict
+from time import monotonic
+from typing import TYPE_CHECKING, Any
 
-from cachetools import LRUCache
-from impit import AsyncClient, Browser, HTTPError, Response, TimeoutException, TransportError
+from impit import AsyncClient, Browser, HTTPError, Response, TimeoutException, TooManyRedirects, TransportError
 from impit import ProxyError as ImpitProxyError
 from typing_extensions import override
+from yarl import URL
 
 from crawlee._types import HttpHeaders
 from crawlee._utils.blocked import ROTATE_PROXY_ERRORS
@@ -20,7 +22,6 @@
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator
     from datetime import timedelta
-    from http.cookiejar import CookieJar
 
     from crawlee import Request
     from crawlee._types import HttpMethod, HttpPayload
@@ -30,12 +31,49 @@
 
 logger = getLogger(__name__)
 
+_REDIRECT_STATUS_CODES = frozenset(
+    {
+        HTTPStatus.MOVED_PERMANENTLY,
+        HTTPStatus.FOUND,
+        HTTPStatus.SEE_OTHER,
+        HTTPStatus.TEMPORARY_REDIRECT,
+        HTTPStatus.PERMANENT_REDIRECT,
+    }
+)
 
-class _ClientCacheEntry(TypedDict):
-    """Type definition for client cache entries."""
+# Status codes that redirect a `POST` as a `GET`. `HTTPStatus.SEE_OTHER` does so for any method but `GET` and `HEAD`.
+_MOVED_STATUS_CODES = frozenset({HTTPStatus.MOVED_PERMANENTLY, HTTPStatus.FOUND})
 
-    client: AsyncClient
-    cookie_jar: CookieJar | None
+_HTTP_SCHEMES = frozenset({'http', 'https'})
+
+# Headers scoped to a single origin, dropped as soon as a redirect leaves it.
+_CROSS_ORIGIN_HEADERS = frozenset({'authorization', 'cookie', 'proxy-authorization'})
+
+_REQUEST_BODY_HEADERS = frozenset(
+    {'content-encoding', 'content-language', 'content-location', 'content-type', 'content-length'}
+)
+
+
+def _is_cross_origin(url: URL, next_url: URL) -> bool:
+    """Check whether a redirect from `url` to `next_url` leaves the origin.
+
+    Origins are compared as strings, because `yarl` considers an explicitly written default port different from
+    an omitted one.
+    """
+    return str(url.origin()) != str(next_url.origin())
+
+
+def _redirect_method(status_code: int, method: str) -> str:
+    """Resolve the method of a redirected request, following the `HTTP-redirect fetch` algorithm.
+
+    See https://fetch.spec.whatwg.org/#http-redirect-fetch.
+    """
+    if (status_code in _MOVED_STATUS_CODES and method == 'POST') or (
+        status_code == HTTPStatus.SEE_OTHER and method not in {'GET', 'HEAD'}
+    ):
+        return 'GET'
+
+    return method
 
 
 class _ImpitResponse:
@@ -96,6 +134,8 @@ def __init__(
         http3: bool = False,
         verify: bool = True,
         browser: Browser | None = 'firefox',
+        follow_redirects: bool = True,
+        max_redirects: int = 20,
         **async_client_kwargs: Any,
     ) -> None:
         """Initialize a new instance.
@@ -105,6 +145,8 @@ def __init__(
             http3: Whether to enable HTTP/3 support.
             verify: SSL certificates used to verify the identity of requested hosts.
             browser: Browser to impersonate.
+            follow_redirects: Whether to follow HTTP redirects.
+            max_redirects: Maximum number of redirects to follow before raising `impit.TooManyRedirects`.
             async_client_kwargs: Additional keyword arguments for `impit.AsyncClient`.
         """
         super().__init__(
@@ -113,10 +155,12 @@ def __init__(
         self._http3 = http3
         self._verify = verify
         self._browser = browser
+        self._follow_redirects = follow_redirects
+        self._max_redirects = max_redirects
 
         self._async_client_kwargs = async_client_kwargs
 
-        self._client_by_proxy_url = LRUCache[str | None, _ClientCacheEntry](maxsize=10)
+        self._client_by_proxy_url = dict[str | None, AsyncClient]()
 
     @override
     async def crawl(
@@ -128,15 +172,15 @@ async def crawl(
         statistics: Statistics | None = None,
         timeout: timedelta | None = None,
     ) -> HttpCrawlingResult:
-        client = self._get_client(proxy_info.url if proxy_info else None, session.cookies.jar if session else None)
-
         try:
-            response = await client.request(
-                url=request.url,
+            response = await self._request_with_redirects(
                 method=request.method,
-                content=request.payload,
-                headers=dict(request.headers) if request.headers else None,
-                timeout=timeout.total_seconds() if timeout else None,
+                url=request.url,
+                headers=dict(request.headers) if request.headers else {},
+                payload=request.payload,
+                session=session,
+                proxy_info=proxy_info,
+                timeout=timeout,
             )
         except TimeoutException as exc:
             raise asyncio.TimeoutError from exc
@@ -169,15 +213,15 @@ async def send_request(
         if isinstance(headers, dict) or headers is None:
             headers = HttpHeaders(headers or {})
 
-        client = self._get_client(proxy_info.url if proxy_info else None, session.cookies.jar if session else None)
-
         try:
-            response = await client.request(
+            response = await self._request_with_redirects(
                 method=method,
                 url=url,
-                content=payload,
-                headers=dict(headers) if headers else None,
-                timeout=timeout.total_seconds() if timeout else None,
+                headers=dict(headers),
+                payload=payload,
+                session=session,
+                proxy_info=proxy_info,
+                timeout=timeout,
             )
         except TimeoutException as exc:
             raise asyncio.TimeoutError from exc
@@ -203,15 +247,18 @@ async def stream(
     ) -> AsyncGenerator[HttpResponse]:
         validate_http_url(url)
 
-        client = self._get_client(proxy_info.url if proxy_info else None, session.cookies.jar if session else None)
+        if isinstance(headers, dict) or headers is None:
+            headers = HttpHeaders(headers or {})
 
         try:
-            response = await client.request(
+            response = await self._request_with_redirects(
                 method=method,
                 url=url,
-                content=payload,
-                headers=dict(headers) if headers else None,
-                timeout=timeout.total_seconds() if timeout else None,
+                headers=dict(headers),
+                payload=payload,
+                session=session,
+                proxy_info=proxy_info,
+                timeout=timeout,
                 stream=True,
             )
         except TimeoutException as exc:
@@ -222,34 +269,124 @@ async def stream(
         finally:
             response.close()
 
-    def _get_client(self, proxy_url: str | None, cookie_jar: CookieJar | None) -> AsyncClient:
+    async def _request_with_redirects(
+        self,
+        *,
+        method: str,
+        url: str,
+        headers: dict[str, str],
+        payload: HttpPayload | None,
+        session: Session | None,
+        proxy_info: ProxyInfo | None,
+        timeout: timedelta | None,
+        stream: bool = False,
+    ) -> Response:
+        """Perform a request, following redirects one hop at a time.
+
+        Redirects are resolved here instead of by `impit`, so that cookies of the given session are attached to and
+        collected from every single hop. A client is therefore never bound to a session and can be shared by all
+        of them.
+
+        Args:
+            method: The HTTP method to use.
+            url: The URL to send the request to.
+            headers: The headers to include in the request.
+            payload: The data to be sent as the request body.
+            session: The session whose cookies are sent and updated.
+            proxy_info: The information about the proxy to be used.
+            timeout: Maximum time allowed to process the request.
+            stream: Whether the body of the final response should be streamed.
+
+        Raises:
+            TooManyRedirects: If the number of redirects exceeds `max_redirects`.
+
+        Returns:
+            The final response of the redirect chain.
+        """
+        client = self._get_client(proxy_info.url if proxy_info else None)
+        # `encoded=True` keeps the URL byte for byte and safe `%2F` in query.
+        current_url = URL(url, encoded=True)
+        content = payload
+        # The timeout bounds the whole chain rather than each hop, which is how `impit` treats it as well.
+        deadline = monotonic() + timeout.total_seconds() if timeout else None
+
+        for _ in range(self._max_redirects + 1):
+            remaining = deadline - monotonic() if deadline is not None else None
+            if remaining is not None and remaining <= 0:
+                raise asyncio.TimeoutError
+
+            # Rebuilt per hop, so that the `Cookie` header never reaches a URL its cookies do not match. A header
+            # passed by the caller wins over the session, which is how `impit` treats its own cookie jar.
+            request_headers = dict(headers)
+            if (
+                session
+                and 'cookie' not in request_headers
+                and (cookie_string := session.cookies.get_cookie_string(str(current_url)))
+            ):
+                request_headers['cookie'] = cookie_string
+
+            response = await client.request(
+                method=method,
+                url=str(current_url),
+                content=content,
+                headers=request_headers or None,
+                timeout=remaining,
+                stream=stream,
+            )
+
+            if session and self._persist_cookies_per_session:
+                session.cookies.extract_cookies_from_headers(str(current_url), response.headers.get_list('set-cookie'))
+
+            if not self._follow_redirects or response.status_code not in _REDIRECT_STATUS_CODES:
+                return response
+
+            location = response.headers.get('location')
+            if not location:
+                return response
+
+            next_url = current_url.join(URL(location, encoded=True))
+            if next_url.scheme not in _HTTP_SCHEMES:
+                return response
+
+            next_method = _redirect_method(response.status_code, method)
+            if next_method != method:
+                method = next_method
+                content = None
+                headers = {key: value for key, value in headers.items() if key not in _REQUEST_BODY_HEADERS}
+
+            if _is_cross_origin(current_url, next_url):
+                headers = {key: value for key, value in headers.items() if key not in _CROSS_ORIGIN_HEADERS}
+
+            if stream:
+                response.close()
+
+            current_url = next_url
+
+        raise TooManyRedirects(f'Exceeded the limit of {self._max_redirects} redirects while requesting {url}.')
+
+    def _get_client(self, proxy_url: str | None) -> AsyncClient:
         """Retrieve or create an HTTP client for the given proxy URL.
 
         If a client for the specified proxy URL does not exist, create and store a new one.
         """
-        cached_data = self._client_by_proxy_url.get(proxy_url)
-        if cached_data:
-            client = cached_data['client']
-            client_cookie_jar = cached_data['cookie_jar']
-            if client_cookie_jar is cookie_jar:
-                # If the cookie jar matches, return the existing client.
-                return client
+        if client := self._client_by_proxy_url.get(proxy_url):
+            return client
 
         # Prepare a default kwargs for the new client.
         kwargs: dict[str, Any] = {
             'proxy': proxy_url,
             'http3': self._http3,
             'verify': self._verify,
-            'follow_redirects': True,
             'browser': self._browser,
         }
 
         # Update the default kwargs with any additional user-provided kwargs.
         kwargs.update(self._async_client_kwargs)
 
-        client = AsyncClient(**kwargs, cookie_jar=cookie_jar)
+        # Redirects are followed hop by hop by `_request_with_redirects`.
+        client = AsyncClient(**kwargs, follow_redirects=False)
 
-        self._client_by_proxy_url[proxy_url] = _ClientCacheEntry(client=client, cookie_jar=cookie_jar)
+        self._client_by_proxy_url[proxy_url] = client
 
         return client
 

src/crawlee/sessions/_cookies.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 from copy import deepcopy
+from email.message import Message
 from http.cookiejar import Cookie, CookieJar
 from typing import TYPE_CHECKING, Any, Literal
+from urllib.request import Request as UrlRequest
 
 from typing_extensions import NotRequired, Required, TypedDict
 
@@ -13,6 +15,18 @@
     from typing import TypeGuard
 
 
+class _SetCookieResponse:
+    """Minimal response adapter exposing `Set-Cookie` headers to `CookieJar.extract_cookies`."""
+
+    def __init__(self, set_cookie_headers: list[str]) -> None:
+        self._message = Message()
+        for header in set_cookie_headers:
+            self._message['Set-Cookie'] = header
+
+    def info(self) -> Message:
+        return self._message
+
+
 @docs_group('Session management')
 class CookieParam(TypedDict, total=False):
     """Dictionary representation of cookies for `SessionCookies.set` method."""
@@ -217,6 +231,36 @@ def set_cookies(self, cookie_dicts: list[CookieParam]) -> None:
             self.set(**cookie_dict)
         self._jar.clear_expired_cookies()
 
+    def get_cookie_string(self, url: str) -> str:
+        """Build the value of the `Cookie` header for the given URL.
+
+        Only cookies matching the domain, path and security requirements of the URL are included.
+
+        Args:
+            url: The URL the header is built for.
+
+        Returns:
+            The `Cookie` header value, or an empty string if no stored cookie matches the URL.
+        """
+        # `UrlRequest` is only used as a carrier of the URL and headers for the jar, it never opens a connection.
+        url_request = UrlRequest(url)  # noqa: S310
+        self._jar.add_cookie_header(url_request)
+        return url_request.get_header('Cookie', '')
+
+    def extract_cookies_from_headers(self, url: str, set_cookie_headers: list[str]) -> None:
+        """Store cookies from the raw `Set-Cookie` headers of a response.
+
+        Attributes omitted from a header, such as domain and path, are derived from the URL. It must therefore be
+        the URL that produced the given headers, not the URL the request started from.
+
+        Args:
+            url: The URL of the response carrying the headers.
+            set_cookie_headers: Raw values of the `Set-Cookie` response headers.
+        """
+        response = _SetCookieResponse(set_cookie_headers)
+        self._jar.extract_cookies(response, UrlRequest(url))  # noqa: S310 # ty: ignore[invalid-argument-type]
+        self._jar.clear_expired_cookies()
+
     def get_cookies_as_playwright_format(self) -> list[PlaywrightCookieParam]:
         """Get cookies in playwright format."""
         return [self._to_playwright(cookie) for cookie in self.get_cookies_as_dicts()]