dataapp_crawler 0.1.0__py3-none-any.whl

dataapp_crawler/sessions.py

@@ -0,0 +1,735 @@
+"""Bounded, run-local anonymous acquisition sessions.
+
+HTTP cookies and browser context state persist independently inside one slot by
+default.  An explicitly enabled cookie bridge can reconcile cookies only; browser
+storage state and identity headers remain engine-local.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import math
+import sys
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine
+from contextlib import asynccontextmanager
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any
+from uuid import uuid4
+
+import httpx
+
+from .cookies import CookieState, cookies_for_browser, validate_shareable_cookies
+from .errors import ConfigurationError, FetchError
+from .http import HttpSession
+from .models import Identity, Limits
+
+CloseResource = Callable[[Any], Awaitable[None]]
+
+
+@dataclass(frozen=True)
+class SessionInfo:
+    """Safe metadata passed to factories; ``health`` is a lower-is-better error score."""
+
+    id: str
+    health: int
+    use_count: int
+    age: float
+
+
+Factory = Callable[[SessionInfo], Any]
+
+
+class Session:
+    """The current exclusive session lease exposed through ``context.session``."""
+
+    __slots__ = ("_slot", "_active")
+
+    def __init__(self, slot: _Slot) -> None:
+        self._slot: _Slot | None = slot
+        self._active = True
+
+    def _check(self) -> None:
+        if not self._active or self._slot is None:
+            raise RuntimeError("Session lease is no longer active")
+
+    def _invalidate(self) -> None:
+        self._active = False
+        self._slot = None
+
+    @property
+    def active(self) -> bool:
+        """Whether this view still owns its exclusive lease."""
+        return self._active
+
+    @property
+    def id(self) -> str:
+        self._check()
+        assert self._slot is not None
+        return self._slot.id
+
+    @property
+    def health(self) -> int:
+        """Current error score; zero is healthiest and ``mark_good`` lowers it."""
+        self._check()
+        assert self._slot is not None
+        return self._slot.error_score
+
+    @property
+    def use_count(self) -> int:
+        self._check()
+        assert self._slot is not None
+        return self._slot.use_count
+
+    @property
+    def age(self) -> float:
+        self._check()
+        assert self._slot is not None
+        return max(0.0, time.monotonic() - self._slot.created_at)
+
+    def mark_good(self) -> None:
+        self._check()
+        assert self._slot is not None
+        self._slot.error_score = max(0, self._slot.error_score - 1)
+
+    def mark_bad(self) -> None:
+        self._check()
+        assert self._slot is not None
+        self._slot.error_score += 1
+
+    def retire(self) -> None:
+        self._check()
+        assert self._slot is not None
+        self._slot.retire = True
+
+    def __repr__(self) -> str:
+        if not self._active:
+            return "Session(inactive)"
+        return (
+            f"Session(id={self.id!r}, health={self.health}, "
+            f"use_count={self.use_count}, age={self.age:.3f})"
+        )
+
+
+@dataclass
+class _Slot:
+    id: str
+    affinity: tuple[str, str]
+    origin: str
+    identity: Identity
+    proxy: Any
+    http: HttpSession
+    browser_provider: Any
+    created_at: float = field(default_factory=time.monotonic)
+    last_used: float = field(default_factory=time.monotonic)
+    use_count: int = 0
+    error_score: int = 0
+    leased: bool = False
+    retire: bool = False
+    browser: Any = None
+    browser_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    cookie_state: CookieState | None = None
+
+    @property
+    def info(self) -> SessionInfo:
+        return SessionInfo(
+            self.id,
+            self.error_score,
+            self.use_count,
+            max(0.0, time.monotonic() - self.created_at),
+        )
+
+
+class SessionLease:
+    """Internal resources for one exclusive pool lease."""
+
+    __slots__ = ("_slot", "session")
+
+    def __init__(self, slot: _Slot) -> None:
+        self._slot = slot
+        self.session = Session(slot)
+
+    @property
+    def http(self) -> HttpSession:
+        return self._slot.http
+
+    @property
+    def identity(self) -> Identity:
+        return self._slot.identity
+
+    @property
+    def proxy(self) -> Any:
+        return self._slot.proxy
+
+    async def sync_to_http(self) -> None:
+        """Apply canonical cookies before HTTP acquisition when sharing is enabled."""
+        state = self._slot.cookie_state
+        if state is not None:
+            self._slot.http.replace_cookies(state.snapshot())
+
+    async def sync_from_http(self) -> None:
+        """Commit HTTP response cookie mutations, including deletions."""
+        state = self._slot.cookie_state
+        if state is not None:
+            state.replace(self._slot.http.export_cookies())
+
+    async def sync_to_browser(self, handle: Any) -> None:
+        """Apply canonical cookies to an isolated browser context."""
+        state = self._slot.cookie_state
+        if state is None:
+            return
+        context = self._cookie_context(handle)
+        existing = await context.cookies()
+        validate_shareable_cookies(existing)
+        await context.clear_cookies()
+        cookies = cookies_for_browser(state.snapshot())
+        if cookies:
+            await context.add_cookies(cookies)
+
+    async def sync_from_browser(self, handle: Any) -> None:
+        """Commit browser cookie mutations after process code has finished."""
+        state = self._slot.cookie_state
+        if state is None:
+            return
+        context = self._cookie_context(handle)
+        state.replace(await context.cookies())
+
+    @staticmethod
+    def _cookie_context(handle: Any) -> Any:
+        if getattr(handle, "owned", None) is not True:
+            raise ConfigurationError("Cookie sharing requires an owned browser context")
+        if "cookies" not in getattr(handle, "capabilities", frozenset()):
+            raise ConfigurationError("Browser Provider does not support cookie sharing")
+        context = getattr(handle, "context", None)
+        if context is None or any(
+            not callable(getattr(context, name, None))
+            for name in ("cookies", "clear_cookies", "add_cookies")
+        ):
+            raise ConfigurationError("Browser context does not support cookie sharing")
+        return context
+
+    async def browser(self, *, required_capabilities: frozenset[str]) -> Any:
+        async with self._slot.browser_lock:
+            if self._slot.browser is None:
+                provider = self._slot.browser_provider
+                if provider is None:
+                    raise ConfigurationError("A browser Provider is required for browser mode")
+                self._slot.browser = await provider.open(
+                    identity=self._slot.identity, proxy=self._slot.proxy
+                )
+            handle = self._slot.browser
+            if "borrowed" in handle.capabilities:
+                self._slot.retire = True
+                failure = ConfigurationError(
+                    "Borrowed browser contexts cannot be isolated as SessionPool slots"
+                )
+                try:
+                    await asyncio.wait_for(handle.close(), timeout=10)
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    failure.add_note("Borrowed browser rejection cleanup also failed")
+                else:
+                    self._slot.browser = None
+                raise failure
+            if required_capabilities - handle.capabilities:
+                raise ConfigurationError(
+                    "Browser Provider does not support required capabilities"
+                )
+            return handle
+
+
+class SessionPool:
+    """A bounded pool of anonymous, origin-affine acquisition generations.
+
+    A pool is single-use and binds to exactly one :class:`Crawler`.  Factories are
+    called once per slot generation with safe :class:`SessionInfo`; a proxy factory
+    may return either a handle or a provider exposing ``acquire()``.  A fixed handle
+    preserves configuration only and is not proof of a physically sticky exit IP.
+    """
+
+    def __init__(
+        self,
+        *,
+        max_size: int = 4,
+        max_age: float = 300,
+        max_usage: int = 100,
+        max_error_score: int = 3,
+        browser_factory: Factory | None = None,
+        proxy_factory: Factory | None = None,
+        share_cookies: bool = False,
+    ) -> None:
+        if isinstance(max_size, bool) or not isinstance(max_size, int) or max_size < 1:
+            raise ConfigurationError("max_size must be a positive integer")
+        if (
+            isinstance(max_usage, bool)
+            or not isinstance(max_usage, int)
+            or max_usage < 1
+        ):
+            raise ConfigurationError("max_usage must be a positive integer")
+        if (
+            isinstance(max_error_score, bool)
+            or not isinstance(max_error_score, int)
+            or max_error_score < 1
+        ):
+            raise ConfigurationError("max_error_score must be a positive integer")
+        if (
+            isinstance(max_age, bool)
+            or not isinstance(max_age, (int, float))
+            or not math.isfinite(max_age)
+            or max_age <= 0
+        ):
+            raise ConfigurationError("max_age must be a finite positive number")
+        if not isinstance(share_cookies, bool):
+            raise ConfigurationError("share_cookies must be a boolean")
+        self.max_size = max_size
+        self.max_age = float(max_age)
+        self.max_usage = max_usage
+        self.max_error_score = max_error_score
+        self.browser_factory = browser_factory
+        self.proxy_factory = proxy_factory
+        self.share_cookies = share_cookies
+        self._condition = asyncio.Condition()
+        self._slots: list[_Slot] = []
+        self._retired: list[_Slot] = []
+        self._closing: set[int] = set()
+        self._orphaned: list[Any] = []
+        self._creating = 0
+        self._pending_affinities: set[tuple[str, str]] = set()
+        self._bound = False
+        self._closed = False
+        self._limits: Limits | None = None
+        self._identity: Identity | None = None
+        self._transport: httpx.AsyncBaseTransport | None = None
+        self._browser_provider: Any = None
+        self._close_resource: CloseResource | None = None
+
+    def __repr__(self) -> str:
+        return (
+            f"SessionPool(max_size={self.max_size}, max_age={self.max_age}, "
+            f"max_usage={self.max_usage}, max_error_score={self.max_error_score}, "
+            f"share_cookies={self.share_cookies})"
+        )
+
+    def bind(
+        self,
+        *,
+        limits: Limits,
+        identity: Identity,
+        transport: httpx.AsyncBaseTransport | None,
+        browser_provider: Any,
+        close_resource: CloseResource,
+    ) -> None:
+        if self._bound:
+            raise ConfigurationError("A SessionPool cannot be shared by multiple Crawlers")
+        self._bound = True
+        self._limits = limits
+        self._identity = identity
+        self._transport = transport
+        self._browser_provider = browser_provider
+        self._close_resource = close_resource
+
+    @asynccontextmanager
+    async def lease(self, affinity: tuple[str, str], *, origin: str) -> AsyncIterator[SessionLease]:
+        slot = await self._acquire(affinity, origin=origin)
+        lease = SessionLease(slot)
+        try:
+            yield lease
+        finally:
+            primary_error = sys.exc_info()[1]
+            lease.session._invalidate()
+            release = asyncio.create_task(self._release(slot))
+            try:
+                await asyncio.shield(release)
+            except asyncio.CancelledError as cancellation:
+                pending_cleanup_error: BaseException | None = None
+                try:
+                    if not release.done():
+                        await asyncio.shield(release)
+                    release.result()
+                except BaseException as caught_cleanup_error:
+                    pending_cleanup_error = caught_cleanup_error
+                if pending_cleanup_error is not None:
+                    if primary_error is None:
+                        raise pending_cleanup_error from None
+                    primary_error.add_note("Session resource cleanup also failed")
+                elif primary_error is None:
+                    raise cancellation from None
+            except BaseException:
+                if primary_error is None:
+                    raise
+                primary_error.add_note("Session resource cleanup also failed")
+
+    async def _acquire(self, affinity: tuple[str, str], *, origin: str) -> _Slot:
+        while True:
+            retired: _Slot | None = None
+            create = False
+            async with self._condition:
+                if self._closed:
+                    raise RuntimeError("SessionPool is closed")
+                matching = next((slot for slot in self._slots if slot.affinity == affinity), None)
+                if matching is not None:
+                    if matching.origin != origin:
+                        raise ConfigurationError(
+                            "One explicit session_key cannot be shared across origins"
+                        )
+                    if not matching.leased:
+                        if self._expired(matching):
+                            self._slots.remove(matching)
+                            self._retired.append(matching)
+                            self._pending_affinities.add(affinity)
+                            retired = matching
+                        else:
+                            matching.leased = True
+                            return matching
+                elif affinity in self._pending_affinities:
+                    await self._condition.wait()
+                    continue
+                elif len(self._slots) + self._creating < self.max_size:
+                    self._creating += 1
+                    self._pending_affinities.add(affinity)
+                    create = True
+                else:
+                    idle = [slot for slot in self._slots if not slot.leased]
+                    if idle:
+                        retired = min(idle, key=lambda slot: slot.last_used)
+                        self._slots.remove(retired)
+                        self._retired.append(retired)
+                        self._creating += 1
+                        self._pending_affinities.add(affinity)
+                        create = True
+                if retired is None and not create:
+                    await self._condition.wait()
+                    continue
+            if retired is not None:
+                try:
+                    await self._close_tracked(retired)
+                except BaseException as error:
+                    if create:
+                        cancellation, cleanup_error = await self._complete_cleanup(
+                            self._release_creation_reservation(affinity)
+                        )
+                        if cleanup_error is not None:
+                            error.add_note("Session creation reservation cleanup also failed")
+                        if cancellation is not None and not isinstance(
+                            error, asyncio.CancelledError
+                        ):
+                            error.add_note("Session cleanup was also cancelled")
+                    else:
+                        cancellation, cleanup_error = await self._complete_cleanup(
+                            self._discard_pending_affinity(affinity)
+                        )
+                        if cleanup_error is not None:
+                            error.add_note("Session affinity cleanup also failed")
+                        if cancellation is not None and not isinstance(
+                            error, asyncio.CancelledError
+                        ):
+                            error.add_note("Session cleanup was also cancelled")
+                    raise
+                if not create:
+                    cancellation, cleanup_error = await self._complete_cleanup(
+                        self._discard_pending_affinity(affinity)
+                    )
+                    if cleanup_error is not None:
+                        raise cleanup_error from None
+                    if cancellation is not None:
+                        raise cancellation from None
+                    continue
+            if create:
+                try:
+                    slot = await self._create_slot(affinity, origin=origin)
+                except BaseException as error:
+                    cancellation, cleanup_error = await self._complete_cleanup(
+                        self._release_creation_reservation(affinity)
+                    )
+                    if cleanup_error is not None:
+                        error.add_note("Session creation reservation cleanup also failed")
+                    if cancellation is not None and not isinstance(error, asyncio.CancelledError):
+                        error.add_note("Session cleanup was also cancelled")
+                    raise
+                try:
+                    async with self._condition:
+                        if self._closed:
+                            close_new = True
+                        else:
+                            self._creating -= 1
+                            self._pending_affinities.discard(affinity)
+                            slot.leased = True
+                            self._slots.append(slot)
+                            close_new = False
+                        self._condition.notify_all()
+                except BaseException as error:
+                    cancellation, cleanup_error = await self._complete_cleanup(
+                        self._rollback_created_slot(slot, affinity)
+                    )
+                    if cleanup_error is not None:
+                        error.add_note("Created session rollback also failed")
+                    if cancellation is not None and not isinstance(error, asyncio.CancelledError):
+                        error.add_note("Session rollback was also cancelled")
+                    raise
+                if close_new:
+                    cancellation, cleanup_error = await self._complete_cleanup(
+                        self._rollback_created_slot(slot, affinity)
+                    )
+                    if cleanup_error is not None:
+                        raise cleanup_error from None
+                    if cancellation is not None:
+                        raise cancellation from None
+                    raise RuntimeError("SessionPool closed while creating a session")
+                return slot
+
+    async def _release_creation_reservation(self, affinity: tuple[str, str]) -> None:
+        async with self._condition:
+            assert self._creating > 0
+            self._creating -= 1
+            self._pending_affinities.discard(affinity)
+            self._condition.notify_all()
+
+    async def _discard_pending_affinity(self, affinity: tuple[str, str]) -> None:
+        async with self._condition:
+            self._pending_affinities.discard(affinity)
+            self._condition.notify_all()
+
+    async def _rollback_created_slot(
+        self, slot: _Slot, affinity: tuple[str, str]
+    ) -> None:
+        close_error: BaseException | None = None
+        try:
+            await self._close_slot(slot)
+        except BaseException as error:
+            close_error = error
+        async with self._condition:
+            if close_error is not None and slot not in self._retired:
+                self._retired.append(slot)
+            assert self._creating > 0
+            self._creating -= 1
+            self._pending_affinities.discard(affinity)
+            self._condition.notify_all()
+        if close_error is not None:
+            raise close_error from None
+
+    @staticmethod
+    async def _complete_cleanup(
+        cleanup: Coroutine[Any, Any, None],
+    ) -> tuple[asyncio.CancelledError | None, BaseException | None]:
+        task: asyncio.Task[None] = asyncio.create_task(cleanup)
+        cancellation: asyncio.CancelledError | None = None
+        while not task.done():
+            try:
+                await asyncio.shield(task)
+            except asyncio.CancelledError as error:
+                cancellation = cancellation or error
+            except BaseException:
+                # The completed task's exception is classified below so caller
+                # cancellation and cleanup failure can both be preserved.
+                pass
+        try:
+            task.result()
+        except BaseException as error:
+            return cancellation, error
+        return cancellation, None
+
+    async def _create_slot(self, affinity: tuple[str, str], *, origin: str) -> _Slot:
+        assert self._limits is not None and self._identity is not None
+        session_id = uuid4().hex
+        info = SessionInfo(session_id, 0, 0, 0.0)
+        proxy = None
+        browser_provider = self._browser_provider
+        try:
+            if self.proxy_factory is not None:
+                proxy = await _resolve(self.proxy_factory(info))
+                acquire = getattr(proxy, "acquire", None)
+                if callable(acquire):
+                    proxy = await _resolve(acquire())
+            if self.browser_factory is not None:
+                browser_provider = await _resolve(self.browser_factory(info))
+            self._validate_browser_provider(browser_provider)
+            identity = deepcopy(self._identity)
+            if self.share_cookies:
+                validate_shareable_cookies(
+                    [*(identity.storage_state or {}).get("cookies", []), *identity.cookies]
+                )
+            http = HttpSession(
+                self._limits,
+                identity,
+                proxy_url=proxy.http_url() if proxy is not None else None,
+                transport=self._transport,
+                close_transport=self._transport is None,
+                reject_partitioned=self.share_cookies,
+            )
+        except BaseException as error:
+            cleanup_failed = False
+            if proxy is not None:
+                cancellation, cleanup_error = await self._complete_cleanup(
+                    self._close_one(proxy)
+                )
+                if cleanup_error is not None:
+                    cleanup_failed = True
+                    _, orphan_error = await self._complete_cleanup(
+                        self._track_orphan(proxy)
+                    )
+                    cleanup_failed = orphan_error is not None or cleanup_failed
+                if cancellation is not None and not isinstance(error, asyncio.CancelledError):
+                    error.add_note("Partial session cleanup was also cancelled")
+            if not isinstance(error, Exception):
+                if cleanup_failed:
+                    error.add_note("Partial session resource cleanup also failed")
+                raise
+            failure = ConfigurationError("Session resource creation failed")
+            if cleanup_failed:
+                failure.add_note("Partial session resource cleanup also failed")
+            raise failure from None
+        return _Slot(
+            session_id,
+            affinity,
+            origin,
+            identity,
+            proxy,
+            http,
+            browser_provider,
+            cookie_state=CookieState.from_jar(http.client.cookies.jar)
+            if self.share_cookies
+            else None,
+        )
+
+    async def _release(self, slot: _Slot) -> None:
+        close = False
+        async with self._condition:
+            slot.use_count += 1
+            slot.last_used = time.monotonic()
+            slot.leased = False
+            if self._closed or self._expired(slot):
+                if slot in self._slots:
+                    self._slots.remove(slot)
+                if slot not in self._retired:
+                    self._retired.append(slot)
+                close = True
+            self._condition.notify_all()
+        if close:
+            await self._close_tracked(slot)
+
+    def _expired(self, slot: _Slot) -> bool:
+        return (
+            slot.retire
+            or time.monotonic() - slot.created_at >= self.max_age
+            or slot.use_count >= self.max_usage
+            or slot.error_score >= self.max_error_score
+        )
+
+    async def retire_idle(self) -> None:
+        """Retire all currently idle generations; active leases finish normally."""
+        async with self._condition:
+            idle = [slot for slot in self._slots if not slot.leased]
+            for slot in idle:
+                self._slots.remove(slot)
+                self._retired.append(slot)
+            self._condition.notify_all()
+        errors = await asyncio.gather(
+            *(self._close_tracked(slot) for slot in idle), return_exceptions=True
+        )
+        if any(isinstance(error, BaseException) for error in errors):
+            raise FetchError("One or more session resources failed to close") from None
+
+    async def close(self) -> None:
+        async with self._condition:
+            self._closed = True
+            self._condition.notify_all()
+            while (
+                any(slot.leased for slot in self._slots)
+                or self._creating
+                or self._closing
+            ):
+                await self._condition.wait()
+            for slot in self._slots:
+                if slot not in self._retired:
+                    self._retired.append(slot)
+            slots = list(self._retired)
+            orphaned = list(self._orphaned)
+            self._slots = []
+            self._condition.notify_all()
+        errors = await asyncio.gather(
+            *(self._close_tracked(slot) for slot in slots),
+            *(self._close_orphan(resource) for resource in orphaned),
+            return_exceptions=True,
+        )
+        if any(isinstance(error, BaseException) for error in errors):
+            raise FetchError("One or more session resources failed to close") from None
+
+    async def _close_slot(self, slot: _Slot) -> None:
+        errors: list[BaseException] = []
+        for resource in (slot.browser, slot.http, slot.proxy):
+            try:
+                await self._close_one(resource)
+            except BaseException as error:
+                errors.append(error)
+        if errors:
+            raise FetchError("Session resource cleanup failed") from None
+
+    async def _close_tracked(self, slot: _Slot) -> None:
+        async with self._condition:
+            while id(slot) in self._closing:
+                await self._condition.wait()
+            if slot not in self._retired:
+                return
+            self._closing.add(id(slot))
+        cancellation, close_error = await self._complete_cleanup(self._close_slot(slot))
+        finalize_cancellation, finalize_error = await self._complete_cleanup(
+            self._finish_tracked_close(slot, succeeded=close_error is None)
+        )
+        cancellation = cancellation or finalize_cancellation
+        if finalize_error is not None:
+            raise finalize_error from None
+        if close_error is not None:
+            raise close_error from None
+        if cancellation is not None:
+            raise cancellation from None
+
+    async def _finish_tracked_close(self, slot: _Slot, *, succeeded: bool) -> None:
+        async with self._condition:
+            if succeeded and slot in self._retired:
+                self._retired.remove(slot)
+            self._closing.discard(id(slot))
+            self._condition.notify_all()
+
+    async def _track_orphan(self, resource: Any) -> None:
+        async with self._condition:
+            self._orphaned.append(resource)
+
+    def _validate_browser_provider(self, provider: Any) -> None:
+        if provider is None:
+            return
+        from .providers.browser import CDPBrowser
+
+        if isinstance(provider, CDPBrowser):
+            raise ConfigurationError(
+                "Borrowed CDP browsers cannot be isolated as SessionPool slots"
+            )
+        if (
+            getattr(provider, "_dataapp_crawler_provider", None) == "relay"
+            and getattr(provider, "_user_data_dir", None) is not None
+        ):
+            raise ConfigurationError(
+                "Persistent Relay user_data_dir is unavailable in SessionPool mode"
+            )
+
+    async def _close_one(self, resource: Any) -> None:
+        if resource is None:
+            return
+        if self._close_resource is not None:
+            await self._close_resource(resource)
+            return
+        closer = getattr(resource, "close", None) or getattr(resource, "aclose", None)
+        if not callable(closer):
+            raise TypeError("Session resource exposes no close method")
+        await _resolve(closer())
+
+    async def _close_orphan(self, resource: Any) -> None:
+        await self._close_one(resource)
+        async with self._condition:
+            self._orphaned = [item for item in self._orphaned if item is not resource]
+
+
+async def _resolve(value: Any) -> Any:
+    return await value if inspect.isawaitable(value) else value