wpr-http 0.2.0__py3-none-any.whl

wpr_http/client.py

@@ -0,0 +1,588 @@
+"""Async HTTP client facade with prepared Request objects and bounded Workers."""
+
+from __future__ import annotations
+
+import asyncio
+import json as jsonlib
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from operator import index as integer_index
+from typing import Any
+from urllib.parse import urlencode, urljoin, urlsplit, urlunsplit
+
+from wpr_runtime_core.registry import resolve_runtime
+from wpr_runtime_core.session import AsyncWorkerSession as AsyncGeckoWorkerSession
+from wpr_runtime_core.session import RuntimeConfig
+
+from .exceptions import (
+    WprTimeoutError,
+    dispatch_timeout_error,
+    pool_timeout_error,
+)
+from .models import Headers, HttpResponse, Request, Timeout
+from .transport import GeckoNetworkTransport
+
+_UNSET = object()
+
+
+@dataclass(eq=False)
+class _WorkerSlot:
+    session: AsyncGeckoWorkerSession
+
+
+def _timeout(value: float | Timeout | None) -> Timeout:
+    if isinstance(value, Timeout):
+        return value
+    return Timeout(total=None if value is None else float(value))
+
+
+def _append_params(
+    url: str,
+    params: Mapping[str, Any] | Sequence[tuple[str, Any]] | None,
+) -> str:
+    if not params:
+        return url
+    query = urlencode(params, doseq=True)
+    parsed = urlsplit(url)
+    merged = "&".join(part for part in (parsed.query, query) if part)
+    return urlunsplit(
+        (parsed.scheme, parsed.netloc, parsed.path, merged, parsed.fragment)
+    )
+
+
+def _merge_headers(*groups: Mapping[str, str] | None) -> dict[str, str]:
+    values = Headers()
+    for group in groups:
+        if not group:
+            continue
+        for name, value in group.items():
+            values[str(name)] = str(value)
+    return dict(values.items())
+
+
+class AsyncWprHttpClient:
+    """Async HTTP client backed by a bounded pool of Gecko Network Workers.
+
+    Defaults use one Gecko cookie/profile session and up to twenty Channels.
+    Increasing the Worker bound adds independent concurrent Worker sessions;
+    each slot has its own Gecko CookieService/profile state.
+    """
+
+    def __init__(
+        self,
+        *,
+        config: RuntimeConfig | None = None,
+        firefox_version: str | int | None = None,
+        runtime_home: str | None = None,
+        proxy: str | None = None,
+        headers: Mapping[str, str] | None = None,
+        base_url: str | None = None,
+        timeout: float | Timeout | None = 30.0,
+        follow_redirects: bool = True,
+        max_workers: int = 1,
+        release_mode: str = "piggyback",
+        ipc_backend: str = "iocp",
+        max_channels_per_worker: int = 20,
+    ) -> None:
+        try:
+            worker_limit = integer_index(max_workers)
+        except TypeError:
+            raise ValueError("max_workers must be a positive integer") from None
+        if isinstance(max_workers, bool) or worker_limit < 1:
+            raise ValueError("max_workers must be a positive integer")
+        self.max_workers = worker_limit
+        if release_mode not in {"immediate", "piggyback"}:
+            raise ValueError("release_mode must be 'immediate' or 'piggyback'")
+        self.release_mode = release_mode
+        if ipc_backend not in {"threaded", "iocp"}:
+            raise ValueError("ipc_backend must be 'threaded' or 'iocp'")
+        self.ipc_backend = ipc_backend
+        if (
+            type(max_channels_per_worker) is not int
+            or not 1 <= max_channels_per_worker <= 64
+        ):
+            raise ValueError(
+                "max_channels_per_worker must be an integer between 1 and 64"
+            )
+        if max_channels_per_worker > 1 and ipc_backend != "iocp":
+            raise ValueError("multiple channels require ipc_backend='iocp'")
+        self.max_channels_per_worker = max_channels_per_worker
+        self._request_capacity = self.max_workers * max_channels_per_worker
+        self._startup_lock = asyncio.Lock()
+        if config is not None and (
+            firefox_version is not None or runtime_home is not None
+        ):
+            raise ValueError(
+                "config is a developer override; use it separately from firefox_version/runtime_home"
+            )
+        self.runtime_manifest = None
+        if firefox_version is not None or runtime_home is not None:
+            self.runtime_manifest = resolve_runtime(
+                firefox_version or "155", runtime_home=runtime_home
+            )
+            config = self.runtime_manifest.config()
+        self.config = config
+        self.proxy = proxy if proxy is not None else (config.proxy if config else None)
+        self.headers = dict(headers or {})
+        self.base_url = base_url
+        self.timeout = _timeout(timeout)
+        self.follow_redirects = bool(follow_redirects)
+        self._slots: list[_WorkerSlot] = []
+        self._available: asyncio.Queue[_WorkerSlot] | None = None
+        self._capacity: asyncio.Semaphore | None = None
+        self._responses: set[HttpResponse] = set()
+        self._transport = GeckoNetworkTransport()
+        self._entered = False
+        self._requests_started = False
+        self._warming_up = False
+
+    async def __aenter__(self) -> AsyncWprHttpClient:  # noqa: PYI034
+        if self._entered:
+            raise RuntimeError("AsyncWprHttpClient is already open")
+        self._available = asyncio.Queue()
+        self._capacity = asyncio.Semaphore(self._request_capacity)
+        self._requests_started = False
+        self._warming_up = False
+        try:
+            first_slot = await self._start_session()
+            for _ in range(self.max_channels_per_worker):
+                self._available.put_nowait(first_slot)
+        except Exception:
+            self._available = None
+            self._capacity = None
+            self._entered = False
+            raise
+        self._entered = True
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        await self.aclose()
+
+    @property
+    def is_closed(self) -> bool:
+        return not self._entered
+
+    @property
+    def transport_stats(self) -> dict[str, Any]:
+        """Snapshot of the currently resident Worker pool, not retired Workers."""
+        calls: dict[str, int] = {}
+        pending_count = pending_bytes = 0
+        for slot in self._slots:
+            stats = slot.session.transport_stats
+            for method, count in stats["rpc_dispatches"].items():
+                calls[method] = calls.get(method, 0) + count
+            pending_count += stats["pending_blob_count"]
+            pending_bytes += stats["pending_blob_bytes"]
+        return {
+            "worker_count": len(self._slots),
+            "max_channels_per_worker": self.max_channels_per_worker,
+            "rpc_dispatches": calls,
+            "pending_blob_count": pending_count,
+            "pending_blob_bytes": pending_bytes,
+        }
+
+    async def flush_releases(self) -> None:
+        """Confirm all queued Worker-side releases; unread responses stay live."""
+        await asyncio.gather(
+            *(slot.session.flush_body_releases() for slot in self._slots)
+        )
+
+    async def aclose(self) -> None:
+        """Drain active sends, release response leases, and close all Workers."""
+        if not self._entered and not self._slots:
+            return
+        self._entered = False
+        capacity = self._capacity
+        if capacity is not None:
+            for _ in range(self._request_capacity):
+                await capacity.acquire()
+        errors: list[Exception] = []
+        for response in tuple(self._responses):
+            try:
+                await response.aclose()
+            except Exception as exc:  # noqa: BLE001 -- continue closing all resources
+                errors.append(exc)
+        for slot in tuple(self._slots):
+            try:
+                await slot.session.__aexit__(None, None, None)
+            except Exception as exc:  # noqa: BLE001 -- continue closing all Workers
+                errors.append(exc)
+        self._slots.clear()
+        self._responses.clear()
+        self._available = None
+        self._capacity = None
+        if errors:
+            raise RuntimeError(
+                f"client cleanup encountered {len(errors)} error(s): {errors[0]}"
+            ) from errors[0]
+
+    async def _start_session(self) -> _WorkerSlot:
+        session = AsyncGeckoWorkerSession(
+            self.config,
+            proxy=self.proxy,
+            network_only=True,
+            defer_body_releases=self.release_mode == "piggyback",
+            ipc_backend=self.ipc_backend,
+            network_channels=self.max_channels_per_worker,
+        )
+        await session.__aenter__()
+        slot = _WorkerSlot(session)
+        self._slots.append(slot)
+        return slot
+
+    async def warmup(self, url: str) -> list[int]:
+        """Start every Worker and make one real GET through each channel slot.
+
+        Call this immediately after entering the client and before normal
+        requests. It performs ``max_workers * max_channels_per_worker`` GETs, reads and releases their
+        bodies, and keeps the Worker sessions and network connections alive.
+        Choose a safe, idempotent URL; response cookies are retained by each
+        Worker's separate profile.
+        """
+        if not self._entered:
+            raise RuntimeError("AsyncWprHttpClient is not open")
+        if self._requests_started:
+            raise RuntimeError("warmup must run before the first client request")
+        if self._warming_up:
+            raise RuntimeError("client warmup is already in progress")
+        self._warming_up = True
+        try:
+            assert self._available is not None
+            idle_slots = []
+            while True:
+                try:
+                    idle_slots.append(self._available.get_nowait())
+                except asyncio.QueueEmpty:
+                    break
+            missing = max(0, self.max_workers - len(self._slots))
+            started = await asyncio.gather(
+                *(self._start_session() for _ in range(missing)),
+                return_exceptions=True,
+            )
+            new_slots = [item for item in started if isinstance(item, _WorkerSlot)]
+            warm_slots = list(dict.fromkeys(idle_slots)) + new_slots
+            startup_errors = [
+                item for item in started if isinstance(item, BaseException)
+            ]
+            if startup_errors:
+                for slot in warm_slots:
+                    for _ in range(self.max_channels_per_worker):
+                        self._available.put_nowait(slot)
+                raise RuntimeError(
+                    f"could not start {len(startup_errors)} Worker(s) for warmup"
+                ) from startup_errors[0]
+
+            request = self.build_request("GET", url)
+
+            async def warm_slot(slot: _WorkerSlot) -> int:
+                response = await self._transport.send(
+                    slot.session,
+                    request=request,
+                    proxy=self.proxy,
+                    timeout=request.timeout,
+                )
+                try:
+                    await response.aread()
+                    return response.status_code
+                finally:
+                    await response.aclose()
+
+            warm_tokens = [
+                slot for slot in warm_slots for _ in range(self.max_channels_per_worker)
+            ]
+            results = await asyncio.gather(
+                *(warm_slot(slot) for slot in warm_tokens), return_exceptions=True
+            )
+            errors = [item for item in results if isinstance(item, BaseException)]
+            stopped_slots = {
+                slot
+                for slot, result in zip(warm_tokens, results, strict=True)
+                if not slot.session.is_open
+                or getattr(result, "worker_stopped", False)
+                or (
+                    isinstance(result, asyncio.CancelledError)
+                    and not slot.session.supports_request_cancel
+                )
+            }
+            reusable_slots = []
+            for slot in warm_slots:
+                if slot in stopped_slots:
+                    await slot.session.abort()
+                    try:
+                        self._slots.remove(slot)
+                    except ValueError:
+                        pass
+                else:
+                    reusable_slots.append(slot)
+            for slot in reusable_slots:
+                for _ in range(self.max_channels_per_worker):
+                    self._available.put_nowait(slot)
+            if errors:
+                raise RuntimeError(
+                    f"warmup GET failed in {len(errors)} Worker(s)"
+                ) from errors[0]
+            return [int(item) for item in results]
+        finally:
+            self._warming_up = False
+
+    def _resolve_url(self, url: str) -> str:
+        return urljoin(self.base_url.rstrip("/") + "/", url) if self.base_url else url
+
+    def build_request(
+        self,
+        method: str,
+        url: str,
+        *,
+        headers: Mapping[str, str] | None = None,
+        content: str | None = None,
+        data: Any = None,
+        json: Any = None,
+        params: Mapping[str, Any] | Sequence[tuple[str, Any]] | None = None,
+        timeout: float | Timeout | None | object = _UNSET,
+        follow_redirects: bool | None = None,
+    ) -> Request:
+        """Merge client defaults and return a prepared, inspectable request."""
+        if sum(value is not None for value in (content, data, json)) > 1:
+            raise ValueError("pass only one of content, data, or json")
+        if isinstance(content, bytes):
+            raise TypeError(
+                "binary request bodies are not supported by the current Worker protocol; pass UTF-8 text"
+            )
+
+        request_url = _append_params(self._resolve_url(url), params)
+        request_headers = _merge_headers(self.headers, headers)
+        request_body = content
+        if json is not None:
+            request_body = jsonlib.dumps(json, separators=(",", ":"))
+            if not any(name.lower() == "content-type" for name in request_headers):
+                request_headers["content-type"] = "application/json"
+        elif data is not None:
+            if isinstance(data, Mapping) or (
+                isinstance(data, Sequence)
+                and not isinstance(data, (str, bytes, bytearray))
+            ):
+                request_body = urlencode(data, doseq=True)
+                if not any(name.lower() == "content-type" for name in request_headers):
+                    request_headers["content-type"] = (
+                        "application/x-www-form-urlencoded"
+                    )
+            elif isinstance(data, str):
+                request_body = data
+            elif isinstance(data, bytes):
+                raise TypeError(
+                    "binary request bodies are not supported by the current Worker protocol"
+                )
+            else:
+                raise TypeError("data must be a mapping, sequence of pairs, or str")
+
+        selected_timeout = self.timeout if timeout is _UNSET else _timeout(timeout)
+        return Request(
+            method=method,
+            url=request_url,
+            headers=request_headers,
+            content=request_body,
+            timeout=selected_timeout,
+            follow_redirects=(
+                self.follow_redirects
+                if follow_redirects is None
+                else bool(follow_redirects)
+            ),
+        )
+
+    async def _acquire_slot(
+        self,
+        request: Request,
+        queued_at: float,
+        total_deadline: float | None,
+    ) -> _WorkerSlot:
+        assert self._capacity is not None and self._available is not None
+        loop = asyncio.get_running_loop()
+        pool_deadline = (
+            queued_at + request.timeout.pool
+            if request.timeout.pool is not None
+            else None
+        )
+        wait_deadlines = [d for d in (pool_deadline, total_deadline) if d is not None]
+        wait_deadline = min(wait_deadlines) if wait_deadlines else None
+        wait_budget = max(0.0, wait_deadline - loop.time()) if wait_deadline else None
+        pool_limit_is_first = pool_deadline is not None and (
+            total_deadline is None or pool_deadline <= total_deadline
+        )
+        try:
+            if self._capacity.locked():
+                if wait_budget is None:
+                    await self._capacity.acquire()
+                else:
+                    await asyncio.wait_for(self._capacity.acquire(), wait_budget)
+            else:
+                await self._capacity.acquire()
+        except asyncio.TimeoutError as exc:
+            elapsed = loop.time() - queued_at
+            if pool_limit_is_first:
+                raise pool_timeout_error(
+                    request.method,
+                    request.url,
+                    request.timeout.pool or elapsed,
+                    request=request,
+                ) from exc
+            raise dispatch_timeout_error(
+                request.method,
+                request.url,
+                request.timeout.total or elapsed,
+                request=request,
+            ) from exc
+
+        try:
+            if not self._entered:
+                raise RuntimeError("AsyncWprHttpClient is closed")
+            if wait_deadline is None:
+                return await self._take_slot()
+            if not hasattr(asyncio, "timeout_at"):
+                return await asyncio.wait_for(
+                    self._take_slot(), max(0.0, wait_deadline - loop.time())
+                )
+            async with asyncio.timeout_at(wait_deadline):
+                return await self._take_slot()
+        except asyncio.TimeoutError as exc:
+            self._capacity.release()
+            elapsed = loop.time() - queued_at
+            if pool_limit_is_first:
+                raise pool_timeout_error(
+                    request.method,
+                    request.url,
+                    request.timeout.pool or elapsed,
+                    request=request,
+                ) from exc
+            raise dispatch_timeout_error(
+                request.method,
+                request.url,
+                request.timeout.total or elapsed,
+                request=request,
+            ) from exc
+        except BaseException:
+            self._capacity.release()
+            raise
+
+    async def _take_slot(self) -> _WorkerSlot:
+        assert self._available is not None
+        while True:
+            async with self._startup_lock:
+                try:
+                    slot = self._available.get_nowait()
+                except asyncio.QueueEmpty:
+                    if len(self._slots) < self.max_workers:
+                        slot = await self._start_session()
+                        for _ in range(self.max_channels_per_worker - 1):
+                            self._available.put_nowait(slot)
+                        return slot
+                else:
+                    if slot.session.is_open:
+                        return slot
+                    self._retire_slot(slot)
+                    continue
+            slot = await self._available.get()
+            if slot.session.is_open:
+                return slot
+            self._retire_slot(slot)
+
+    def _retire_slot(self, slot: _WorkerSlot) -> None:
+        if slot in self._slots:
+            self._slots.remove(slot)
+        if self._available is not None:
+            keep = []
+            while not self._available.empty():
+                token = self._available.get_nowait()
+                if token is not slot:
+                    keep.append(token)
+            for token in keep:
+                self._available.put_nowait(token)
+
+    async def send(self, request: Request) -> HttpResponse:
+        """Send a prepared request and return its response object."""
+        if not self._entered:
+            raise RuntimeError("AsyncWprHttpClient is not open")
+        if not isinstance(request, Request):
+            raise TypeError("request must be a wpr_http.Request")
+        if self._warming_up:
+            raise RuntimeError("normal requests are paused during client warmup")
+        self._requests_started = True
+        loop = asyncio.get_running_loop()
+        queued_at = loop.time()
+        total_deadline = (
+            queued_at + request.timeout.total
+            if request.timeout.total is not None
+            else None
+        )
+        slot = await self._acquire_slot(request, queued_at, total_deadline)
+        pool_wait_ms = (loop.time() - queued_at) * 1000
+        reusable = True
+        try:
+            if not self._entered:
+                raise RuntimeError("AsyncWprHttpClient is closed")
+            active_timeout = request.timeout
+            if total_deadline is not None:
+                remaining = max(0.0, total_deadline - loop.time())
+                if remaining == 0:
+                    raise dispatch_timeout_error(
+                        request.method,
+                        request.url,
+                        request.timeout.total or loop.time() - queued_at,
+                        request=request,
+                    )
+                active_timeout = Timeout(total=remaining, pool=request.timeout.pool)
+            response = await self._transport.send(
+                slot.session,
+                request=request,
+                proxy=self.proxy,
+                timeout=active_timeout,
+            )
+            response._wpr_diagnostics["pool_wait_ms"] = pool_wait_ms
+            self._responses.add(response)
+            response._set_on_close(self._responses.discard)
+            return response
+        except WprTimeoutError as exc:
+            if exc.worker_stopped:
+                reusable = False
+                await slot.session.abort()
+            raise
+        except asyncio.CancelledError:
+            if not slot.session.supports_request_cancel:
+                reusable = False
+                await slot.session.abort()
+            raise
+        finally:
+            reusable = reusable and slot.session.is_open
+            if reusable and self._entered and self._available is not None:
+                self._available.put_nowait(slot)
+            elif not reusable:
+                self._retire_slot(slot)
+            if self._capacity is not None:
+                self._capacity.release()
+
+    async def request(self, method: str, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.send(self.build_request(method, url, **kwargs))
+
+    async def get(self, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.request("GET", url, **kwargs)
+
+    async def post(self, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.request("POST", url, **kwargs)
+
+    async def put(self, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.request("PUT", url, **kwargs)
+
+    async def patch(self, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.request("PATCH", url, **kwargs)
+
+    async def delete(self, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.request("DELETE", url, **kwargs)
+
+    async def head(self, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.request("HEAD", url, **kwargs)
+
+    async def options(self, url: str, **kwargs: Any) -> HttpResponse:
+        return await self.request("OPTIONS", url, **kwargs)
+
+
+class WprHttpClient(AsyncWprHttpClient):
+    """Naming alias; this client remains asynchronous."""