hedge-python 0.1.0__py3-none-any.whl

hedge/sketch/_windowed.py

@@ -0,0 +1,115 @@
+"""WindowedSketch: sliding-window quantile estimation over DDSketch pairs."""
+
+from __future__ import annotations
+
+import asyncio
+import math
+import threading
+
+from hedge.sketch._ddsketch import DDSketch
+
+_DEFAULT_WINDOW_DURATION = 30.0  # seconds
+
+
+class WindowedSketch:
+    """Maintains a sliding window over two DDSketches that rotate periodically.
+
+    Quantile queries merge both sketches, giving a window that spans
+    1x to 2x the configured duration. Add always writes to the current sketch.
+
+    The rotation scheme::
+
+        t=0:  current=A, previous=empty
+        t=30: current=B, previous=A       (A covers [0,30))
+        t=60: current=C, previous=B       (A is dropped)
+
+    This class is thread-safe and can be used from both sync and async code.
+    For async rotation, call ``start_async()`` after creating the sketch.
+
+    Args:
+        relative_accuracy: DDSketch relative accuracy (default: 0.01).
+        window_duration: Rotation interval in seconds (default: 30.0).
+    """
+
+    def __init__(
+        self,
+        relative_accuracy: float = 0.01,
+        window_duration: float = _DEFAULT_WINDOW_DURATION,
+    ) -> None:
+        if window_duration <= 0:
+            window_duration = _DEFAULT_WINDOW_DURATION
+        self._relative_accuracy = relative_accuracy
+        self._window_duration = window_duration
+        self._lock = threading.RLock()
+        self._current = DDSketch(relative_accuracy)
+        self._previous = DDSketch(relative_accuracy)
+        # Sync rotation
+        self._stop_event = threading.Event()
+        self._thread: threading.Thread | None = None
+        # Async rotation
+        self._async_task: asyncio.Task[None] | None = None
+
+    def start(self) -> None:
+        """Start background rotation thread (for sync usage)."""
+        if self._thread is not None:
+            return
+        self._stop_event.clear()
+        self._thread = threading.Thread(target=self._rotate_loop, daemon=True)
+        self._thread.start()
+
+    def start_async(self) -> None:
+        """Start background rotation as an asyncio task (for async usage).
+
+        Must be called from within a running event loop.
+        """
+        if self._async_task is not None:
+            return
+        self._async_task = asyncio.ensure_future(self._rotate_loop_async())
+
+    def stop(self) -> None:
+        """Stop the background rotation thread and wait for it to exit."""
+        if self._thread is not None:
+            self._stop_event.set()
+            self._thread.join()
+            self._thread = None
+        if self._async_task is not None:
+            self._async_task.cancel()
+            self._async_task = None
+
+    def add(self, value: float) -> None:
+        """Record a latency sample (in seconds) to the current sketch."""
+        with self._lock:
+            self._current.add(value)
+
+    def quantile(self, q: float) -> float:
+        """Return the estimated quantile q in [0, 1] over the sliding window.
+
+        Returns math.nan if no data has been recorded.
+        """
+        with self._lock:
+            if self._current.count == 0 and self._previous.count == 0:
+                return math.nan
+            merged = DDSketch(self._relative_accuracy)
+            merged.merge(self._previous)
+            merged.merge(self._current)
+            return merged.quantile(q)
+
+    def rotate(self) -> None:
+        """Manually rotate the window. Mostly useful for testing."""
+        with self._lock:
+            self._previous = self._current
+            self._current = DDSketch(self._relative_accuracy)
+
+    def _rotate_loop(self) -> None:
+        """Background thread rotation loop."""
+        while not self._stop_event.wait(self._window_duration):
+            self.rotate()
+
+    async def _rotate_loop_async(self) -> None:
+        """Async rotation loop."""
+        try:
+            while True:
+                await asyncio.sleep(self._window_duration)
+                self.rotate()
+        except asyncio.CancelledError:
+            return

hedge/transport/__init__.py

@@ -0,0 +1,24 @@
+"""Framework-specific hedge transports."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from hedge.transport._aiohttp import HedgedAiohttpSession
+    from hedge.transport._httpx import HedgedHttpxTransport
+
+
+def __getattr__(name: str):  # type: ignore[no-untyped-def]
+    if name == "HedgedHttpxTransport":
+        from hedge.transport._httpx import HedgedHttpxTransport
+
+        return HedgedHttpxTransport
+    if name == "HedgedAiohttpSession":
+        from hedge.transport._aiohttp import HedgedAiohttpSession
+
+        return HedgedAiohttpSession
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["HedgedHttpxTransport", "HedgedAiohttpSession"]

hedge/transport/_aiohttp.py

@@ -0,0 +1,135 @@
+"""Hedged session wrapper for aiohttp.ClientSession.
+
+Usage::
+
+    from hedge import HedgeConfig
+    from hedge.transport import HedgedAiohttpSession
+
+    async with HedgedAiohttpSession(config=HedgeConfig()) as session:
+        resp = await session.get("https://api.example.com/data")
+        data = await resp.json()
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+try:
+    import aiohttp
+except ImportError as exc:
+    raise ImportError(
+        "aiohttp is required for HedgedAiohttpSession. "
+        "Install it with: pip install hedge-python[aiohttp]"
+    ) from exc
+
+from hedge._options import HedgeConfig
+from hedge.transport._base import HedgeScheduler
+
+if TYPE_CHECKING:
+    from hedge._stats import Stats
+
+
+class HedgedAiohttpSession:
+    """An aiohttp session wrapper that adds adaptive hedged requests.
+
+    Wraps an ``aiohttp.ClientSession`` and races a backup request when the
+    primary exceeds its estimated latency percentile.
+
+    Args:
+        config: Hedge configuration. Defaults to ``HedgeConfig()``.
+        **session_kwargs: Additional keyword arguments passed to
+            ``aiohttp.ClientSession()``.
+    """
+
+    def __init__(
+        self,
+        config: HedgeConfig | None = None,
+        **session_kwargs: Any,
+    ) -> None:
+        self._config = config or HedgeConfig()
+        self._session_kwargs = session_kwargs
+        self._session: aiohttp.ClientSession | None = None
+        self._scheduler = HedgeScheduler(self._config)
+
+    def _get_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession(**self._session_kwargs)
+        return self._session
+
+    @property
+    def stats(self) -> Stats:
+        """Access the live Stats object."""
+        return self._scheduler.stats
+
+    async def _request(
+        self,
+        method: str,
+        url: str,
+        **kwargs: Any,
+    ) -> aiohttp.ClientResponse:
+        """Perform a hedged request."""
+        parsed = urlparse(url)
+        host = parsed.netloc or parsed.hostname or url
+        sketch = self._scheduler.sketch_for(host)
+
+        can_hedge = method.upper() in ("GET", "HEAD", "OPTIONS")
+        session = self._get_session()
+
+        async def do_request() -> aiohttp.ClientResponse:
+            return await session.request(method, url, **kwargs)
+
+        def record_latency(response: aiohttp.ClientResponse, elapsed: float) -> None:
+            sketch.add(elapsed)
+
+        return await self._scheduler.execute_with_hedge(
+            host=host,
+            primary_func=do_request,
+            hedge_func=do_request,
+            record_latency=record_latency,
+            can_hedge=can_hedge,
+        )
+
+    async def get(self, url: str, **kwargs: Any) -> aiohttp.ClientResponse:
+        """Perform a hedged GET request."""
+        return await self._request("GET", url, **kwargs)
+
+    async def post(self, url: str, **kwargs: Any) -> aiohttp.ClientResponse:
+        """Perform a POST request (no hedging; body cannot be safely replayed)."""
+        return await self._request("POST", url, **kwargs)
+
+    async def put(self, url: str, **kwargs: Any) -> aiohttp.ClientResponse:
+        """Perform a PUT request (no hedging)."""
+        return await self._request("PUT", url, **kwargs)
+
+    async def delete(self, url: str, **kwargs: Any) -> aiohttp.ClientResponse:
+        """Perform a DELETE request (no hedging)."""
+        return await self._request("DELETE", url, **kwargs)
+
+    async def head(self, url: str, **kwargs: Any) -> aiohttp.ClientResponse:
+        """Perform a hedged HEAD request."""
+        return await self._request("HEAD", url, **kwargs)
+
+    async def options(self, url: str, **kwargs: Any) -> aiohttp.ClientResponse:
+        """Perform a hedged OPTIONS request."""
+        return await self._request("OPTIONS", url, **kwargs)
+
+    async def close(self) -> None:
+        """Close the session and stop background tasks."""
+        await self._scheduler.close()
+        if self._session and not self._session.closed:
+            await self._session.close()
+
+    async def __aenter__(self) -> HedgedAiohttpSession:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.close()

hedge/transport/_base.py

@@ -0,0 +1,159 @@
+"""Shared async hedge scheduling logic used by all transports."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import math
+import time
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
+
+from hedge._stats import Stats
+from hedge.budget import TokenBucket
+from hedge.sketch import WindowedSketch
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Coroutine
+
+    from hedge._options import HedgeConfig
+
+T = TypeVar("T")
+
+
+class HedgeScheduler:
+    """Core async hedge scheduling shared across frameworks.
+
+    Manages per-host sketches, request counters, token bucket budget,
+    and the race-then-cancel logic.
+
+    This is an internal building block; users should use the framework-specific
+    transports (httpx, aiohttp) or interceptors (gRPC).
+    """
+
+    def __init__(self, config: HedgeConfig) -> None:
+        self.config = config
+        self.stats = config.stats or Stats()
+        self.budget = TokenBucket(config.budget_percent, config.estimated_rps)
+        self._sketches: dict[str, WindowedSketch] = {}
+        self._counters: dict[str, int] = defaultdict(int)
+        self._lock = asyncio.Lock()
+
+    def sketch_for(self, host: str) -> WindowedSketch:
+        """Get or create a WindowedSketch for the given host."""
+        if host not in self._sketches:
+            sketch = WindowedSketch(
+                relative_accuracy=0.01,
+                window_duration=self.config.window_duration,
+            )
+            sketch.start_async()
+            self._sketches[host] = sketch
+        return self._sketches[host]
+
+    async def increment_counter(self, host: str) -> int:
+        """Atomically increment and return the request counter for a host."""
+        async with self._lock:
+            self._counters[host] += 1
+            return self._counters[host]
+
+    def compute_hedge_delay(self, host: str, request_number: int) -> float:
+        """Compute the hedge delay in seconds for a given host and request number."""
+        if request_number <= self.config.warmup_requests:
+            self.stats.increment_warmup()
+            delay = self.config.warmup_delay
+        else:
+            sketch = self.sketch_for(host)
+            estimate = sketch.quantile(self.config.percentile)
+            delay = estimate if estimate > 0 and not math.isnan(estimate) else self.config.warmup_delay
+
+        return max(delay, self.config.min_delay)
+
+    async def execute_with_hedge(
+        self,
+        host: str,
+        primary_func: Callable[[], Awaitable[T]],
+        hedge_func: Callable[[], Awaitable[T]],
+        record_latency: Callable[[T, float], None],
+        can_hedge: bool = True,
+    ) -> T:
+        """Execute primary request with hedge racing logic.
+
+        Args:
+            host: Target host key for per-host sketch tracking.
+            primary_func: Async callable that performs the primary request.
+            hedge_func: Async callable that performs the hedge request.
+            record_latency: Callback to record latency to the sketch.
+            can_hedge: Whether the request is safe to hedge (idempotent).
+
+        Returns:
+            The result from whichever request finishes first.
+        """
+        self.stats.increment_total()
+
+        request_number = await self.increment_counter(host)
+        hedge_delay = self.compute_hedge_delay(host, request_number)
+        start = time.monotonic()
+
+        # Launch primary. ``primary_func()`` returns ``Awaitable[T]`` in the
+        # signature, but in practice transports always supply ``async def``
+        # callables (i.e. coroutine functions). ``asyncio.create_task`` only
+        # accepts coroutines, so cast for the type checker.
+        primary_coro = cast("Coroutine[Any, Any, T]", primary_func())
+        primary_task: asyncio.Task[T] = asyncio.create_task(primary_coro)
+
+        # Wait for hedge delay
+        done, _ = await asyncio.wait({primary_task}, timeout=hedge_delay)
+        if done:
+            result: T = primary_task.result()
+            elapsed = time.monotonic() - start
+            record_latency(result, elapsed)
+            return result
+
+        # Check if hedging is allowed
+        if not can_hedge:
+            result = await primary_task
+            elapsed = time.monotonic() - start
+            record_latency(result, elapsed)
+            return result
+
+        # Check budget
+        if not self.budget.try_acquire():
+            self.stats.increment_budget_exhausted()
+            result = await primary_task
+            elapsed = time.monotonic() - start
+            record_latency(result, elapsed)
+            return result
+
+        # Launch hedge
+        self.stats.increment_hedged()
+        hedge_coro = cast("Coroutine[Any, Any, T]", hedge_func())
+        hedge_task: asyncio.Task[T] = asyncio.create_task(hedge_coro)
+
+        done, pending = await asyncio.wait(
+            {primary_task, hedge_task},
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
+        winner_task = done.pop()
+        elapsed = time.monotonic() - start
+
+        # Cancel losers
+        for task in pending:
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError, Exception):
+                await task
+
+        if winner_task is primary_task:
+            self.stats.increment_primary_wins()
+        else:
+            self.stats.increment_hedge_wins()
+
+        result = winner_task.result()
+        record_latency(result, elapsed)
+        return result
+
+    async def close(self) -> None:
+        """Stop all background sketch rotation tasks."""
+        for sketch in self._sketches.values():
+            sketch.stop()
+        self._sketches.clear()

hedge/transport/_httpx.py

@@ -0,0 +1,84 @@
+"""Hedged transport for httpx.AsyncClient.
+
+Usage::
+
+    import httpx
+    from hedge import HedgeConfig
+    from hedge.transport import HedgedHttpxTransport
+
+    transport = HedgedHttpxTransport(config=HedgeConfig())
+    async with httpx.AsyncClient(transport=transport) as client:
+        resp = await client.get("https://api.example.com/data")
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+try:
+    import httpx
+except ImportError as exc:
+    raise ImportError(
+        "httpx is required for HedgedHttpxTransport. "
+        "Install it with: pip install hedge-python[httpx]"
+    ) from exc
+
+from hedge._options import HedgeConfig
+from hedge.transport._base import HedgeScheduler
+
+if TYPE_CHECKING:
+    from hedge._stats import Stats
+
+
+class HedgedHttpxTransport(httpx.AsyncBaseTransport):
+    """An httpx async transport that adds adaptive hedged requests.
+
+    Wraps an inner transport (default: ``httpx.AsyncHTTPTransport``) and
+    races a backup request when the primary exceeds its estimated latency
+    percentile.
+
+    Args:
+        inner: The underlying transport to wrap. Defaults to a new
+            ``httpx.AsyncHTTPTransport()``.
+        config: Hedge configuration. Defaults to ``HedgeConfig()``.
+    """
+
+    def __init__(
+        self,
+        inner: httpx.AsyncBaseTransport | None = None,
+        config: HedgeConfig | None = None,
+    ) -> None:
+        self._inner = inner or httpx.AsyncHTTPTransport()
+        self._config = config or HedgeConfig()
+        self._scheduler = HedgeScheduler(self._config)
+
+    @property
+    def stats(self) -> Stats:
+        """Access the live Stats object."""
+        return self._scheduler.stats
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        """Handle an outgoing request with adaptive hedging."""
+        host = str(request.url.host)
+        sketch = self._scheduler.sketch_for(host)
+
+        can_hedge = request.method.upper() in ("GET", "HEAD", "OPTIONS")
+
+        async def do_request() -> httpx.Response:
+            return await self._inner.handle_async_request(request)
+
+        def record_latency(response: httpx.Response, elapsed: float) -> None:
+            sketch.add(elapsed)
+
+        return await self._scheduler.execute_with_hedge(
+            host=host,
+            primary_func=do_request,
+            hedge_func=do_request,
+            record_latency=record_latency,
+            can_hedge=can_hedge,
+        )
+
+    async def aclose(self) -> None:
+        """Close the transport and stop background tasks."""
+        await self._scheduler.close()
+        await self._inner.aclose()