hedge-python 0.1.0__py3-none-any.whl

hedge/__init__.py

@@ -0,0 +1,17 @@
+"""Adaptive hedged request library for Python.
+
+Learns per-host latency distributions using DDSketch, fires a backup request
+when the primary exceeds its estimated p90, and caps hedge rate with a token
+bucket to prevent load amplification during outages.
+"""
+
+from hedge._options import HedgeConfig
+from hedge._stats import Stats, StatsSnapshot
+
+__all__ = [
+    "HedgeConfig",
+    "Stats",
+    "StatsSnapshot",
+]
+
+__version__ = "0.1.0"

hedge/_options.py

@@ -0,0 +1,35 @@
+"""Configuration options for hedge transports and interceptors."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from hedge._stats import Stats
+
+
+@dataclass
+class HedgeConfig:
+    """Configuration for hedge behavior.
+
+    Attributes:
+        percentile: Sketch quantile used as hedge trigger (default: 0.90).
+        max_hedges: Maximum concurrent hedge requests per call (default: 1).
+        budget_percent: Max hedge rate as percent of total traffic (default: 10.0).
+        estimated_rps: Expected requests per second; sets token bucket capacity (default: 100.0).
+        min_delay: Floor on the hedge delay in seconds (default: 0.001).
+        warmup_requests: Number of initial requests using fixed delay (default: 20).
+        warmup_delay: Fixed hedge delay during warmup in seconds (default: 0.01).
+        window_duration: Sketch window rotation interval in seconds (default: 30.0).
+    """
+
+    percentile: float = 0.90
+    max_hedges: int = 1
+    budget_percent: float = 10.0
+    estimated_rps: float = 100.0
+    min_delay: float = 0.001
+    warmup_requests: int = 20
+    warmup_delay: float = 0.01
+    window_duration: float = 30.0
+    stats: Stats | None = field(default=None, repr=False)

hedge/_stats.py

@@ -0,0 +1,75 @@
+"""Thread-safe statistics for hedge operations."""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+
+
+class Stats:
+    """Thread-safe counters for hedge operations.
+
+    All fields use a lock for atomic updates and are safe to read concurrently.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self.total_requests: int = 0
+        self.hedged_requests: int = 0
+        self.hedge_wins: int = 0
+        self.primary_wins: int = 0
+        self.budget_exhausted: int = 0
+        self.warmup_requests: int = 0
+
+    def _increment(self, field: str, value: int = 1) -> None:
+        with self._lock:
+            setattr(self, field, getattr(self, field) + value)
+
+    def increment_total(self) -> None:
+        self._increment("total_requests")
+
+    def increment_hedged(self) -> None:
+        self._increment("hedged_requests")
+
+    def increment_hedge_wins(self) -> None:
+        self._increment("hedge_wins")
+
+    def increment_primary_wins(self) -> None:
+        self._increment("primary_wins")
+
+    def increment_budget_exhausted(self) -> None:
+        self._increment("budget_exhausted")
+
+    def increment_warmup(self) -> None:
+        self._increment("warmup_requests")
+
+    def snapshot(self) -> StatsSnapshot:
+        """Take a consistent point-in-time copy of all counters."""
+        with self._lock:
+            return StatsSnapshot(
+                total_requests=self.total_requests,
+                hedged_requests=self.hedged_requests,
+                hedge_wins=self.hedge_wins,
+                primary_wins=self.primary_wins,
+                budget_exhausted=self.budget_exhausted,
+                warmup_requests=self.warmup_requests,
+            )
+
+    def hedge_rate(self) -> float:
+        """Return hedged_requests / total_requests, or 0.0 if no requests."""
+        with self._lock:
+            if self.total_requests == 0:
+                return 0.0
+            return self.hedged_requests / self.total_requests
+
+
+@dataclass(frozen=True)
+class StatsSnapshot:
+    """Immutable point-in-time snapshot of Stats."""
+
+    total_requests: int
+    hedged_requests: int
+    hedge_wins: int
+    primary_wins: int
+    budget_exhausted: int
+    warmup_requests: int

hedge/budget/__init__.py

@@ -0,0 +1,5 @@
+"""Token bucket rate limiter for hedge request budget control."""
+
+from hedge.budget._token_bucket import TokenBucket
+
+__all__ = ["TokenBucket"]

hedge/budget/_token_bucket.py

@@ -0,0 +1,52 @@
+"""Token bucket algorithm for controlling hedge request rate."""
+
+from __future__ import annotations
+
+import threading
+import time
+
+
+class TokenBucket:
+    """Controls hedge request rate using a token bucket algorithm.
+
+    The bucket refills at ``estimated_rps * budget_percent / 100`` tokens per
+    second. During genuine outages the bucket drains and hedging stops,
+    preventing the load-doubling spiral that would deepen the incident.
+
+    Args:
+        budget_percent: Max hedge rate as percent of total traffic.
+        estimated_rps: Expected requests per second.
+    """
+
+    def __init__(self, budget_percent: float = 10.0, estimated_rps: float = 100.0) -> None:
+        self._lock = threading.Lock()
+        self._budget_percent = budget_percent
+        self._rate = estimated_rps * (budget_percent / 100.0)
+        self._max_burst = max(self._rate * 2, 1.0)
+        self._tokens = self._max_burst
+        self._last_refill = time.monotonic()
+
+    def try_acquire(self) -> bool:
+        """Return True if a hedge token is available, False if over budget."""
+        with self._lock:
+            now = time.monotonic()
+            elapsed = now - self._last_refill
+            self._last_refill = now
+
+            self._tokens += elapsed * self._rate
+            if self._tokens > self._max_burst:
+                self._tokens = self._max_burst
+
+            if self._tokens < 1.0:
+                return False
+            self._tokens -= 1.0
+            return True
+
+    def set_rps(self, rps: float) -> None:
+        """Update the hedge rate as traffic changes, preserving the budget percentage."""
+        with self._lock:
+            self._rate = rps * (self._budget_percent / 100.0)
+            max_burst = max(self._rate * 2, 1.0)
+            self._max_burst = max_burst
+            if self._tokens > self._max_burst:
+                self._tokens = self._max_burst

hedge/interceptor/__init__.py

@@ -0,0 +1,23 @@
+"""gRPC interceptors with adaptive hedging."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from hedge.interceptor._grpc import HedgedServerStreamInterceptor, HedgedUnaryInterceptor
+
+
+def __getattr__(name: str):  # type: ignore[no-untyped-def]
+    if name == "HedgedUnaryInterceptor":
+        from hedge.interceptor._grpc import HedgedUnaryInterceptor
+
+        return HedgedUnaryInterceptor
+    if name == "HedgedServerStreamInterceptor":
+        from hedge.interceptor._grpc import HedgedServerStreamInterceptor
+
+        return HedgedServerStreamInterceptor
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["HedgedUnaryInterceptor", "HedgedServerStreamInterceptor"]

hedge/interceptor/_grpc.py

@@ -0,0 +1,302 @@
+"""Hedged gRPC client interceptors.
+
+Supports both Unary-Unary and Unary-Stream (server streaming) RPCs.
+
+Usage::
+
+    from hedge import HedgeConfig
+    from hedge.interceptor import HedgedUnaryInterceptor
+
+    channel = grpc.aio.insecure_channel("localhost:50051")
+    hedged_channel = grpc.aio.insecure_channel(
+        "localhost:50051",
+        interceptors=[HedgedUnaryInterceptor(config=HedgeConfig())],
+    )
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import math
+import time
+from collections import defaultdict
+from typing import Any, Callable
+
+try:
+    import grpc  # type: ignore[import-untyped]
+    import grpc.aio  # type: ignore[import-untyped]
+except ImportError as exc:
+    raise ImportError(
+        "grpcio is required for gRPC interceptors. "
+        "Install it with: pip install hedge-python[grpc]"
+    ) from exc
+
+from hedge._options import HedgeConfig
+from hedge._stats import Stats
+from hedge.budget import TokenBucket
+from hedge.sketch import WindowedSketch
+
+
+class HedgedUnaryInterceptor(grpc.aio.UnaryUnaryClientInterceptor):  # type: ignore[misc]
+    """gRPC Unary-Unary client interceptor with adaptive hedging.
+
+    Learns per-target latency distributions and fires a backup RPC when the
+    primary exceeds its estimated percentile threshold.
+
+    Args:
+        config: Hedge configuration. Defaults to ``HedgeConfig()``.
+    """
+
+    def __init__(self, config: HedgeConfig | None = None) -> None:
+        self._config = config or HedgeConfig()
+        self.stats = self._config.stats or Stats()
+        self._budget = TokenBucket(self._config.budget_percent, self._config.estimated_rps)
+        self._sketches: dict[str, WindowedSketch] = {}
+        self._counters: dict[str, int] = defaultdict(int)
+        self._lock = asyncio.Lock()
+
+    def _sketch_for(self, target: str) -> WindowedSketch:
+        if target not in self._sketches:
+            sketch = WindowedSketch(0.01, self._config.window_duration)
+            sketch.start_async()
+            self._sketches[target] = sketch
+        return self._sketches[target]
+
+    async def _increment_counter(self, target: str) -> int:
+        async with self._lock:
+            self._counters[target] += 1
+            return self._counters[target]
+
+    async def intercept_unary_unary(
+        self,
+        continuation: Callable[..., Any],
+        client_call_details: grpc.aio.ClientCallDetails,
+        request: Any,
+    ) -> Any:
+        """Intercept a unary-unary RPC with adaptive hedging."""
+        self.stats.increment_total()
+
+        target = client_call_details.method
+        sketch = self._sketch_for(target)
+        request_number = await self._increment_counter(target)
+
+        if request_number <= self._config.warmup_requests:
+            self.stats.increment_warmup()
+            hedge_delay = self._config.warmup_delay
+        else:
+            estimate = sketch.quantile(self._config.percentile)
+            hedge_delay = estimate if estimate > 0 and not math.isnan(estimate) else self._config.warmup_delay
+
+        hedge_delay = max(hedge_delay, self._config.min_delay)
+        start = time.monotonic()
+
+        # We track the live Call object for each in-flight attempt so we can
+        # cancel the underlying RPC (not just the asyncio task) when a loser
+        # is dropped. ``continuation`` returns a Call object almost
+        # immediately; the actual RTT is spent in ``await call``. We must
+        # combine both steps inside the task, otherwise the task completes
+        # instantly and the hedge timer never fires.
+        call_holder: dict[asyncio.Task[Any], Any] = {}
+
+        async def invoke() -> Any:
+            call = await continuation(client_call_details, request)
+            # Record the Call for the currently running task so the loser
+            # branch can cancel the RPC on the wire.
+            current = asyncio.current_task()
+            if current is not None:
+                call_holder[current] = call
+            return await call
+
+        # Launch primary
+        primary_task = asyncio.create_task(invoke())
+
+        # Wait hedge_delay
+        done, _ = await asyncio.wait({primary_task}, timeout=hedge_delay)
+        if done:
+            response = primary_task.result()
+            sketch.add(time.monotonic() - start)
+            return response
+
+        # Budget check
+        if not self._budget.try_acquire():
+            self.stats.increment_budget_exhausted()
+            response = await primary_task
+            sketch.add(time.monotonic() - start)
+            return response
+
+        # Launch hedge
+        self.stats.increment_hedged()
+        hedge_task = asyncio.create_task(invoke())
+
+        done, pending = await asyncio.wait(
+            {primary_task, hedge_task},
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
+        winner_task = done.pop()
+        elapsed = time.monotonic() - start
+
+        # Cancel losers: cancel the underlying gRPC Call first (best-effort),
+        # then await the task so we don't leak it.
+        for task in pending:
+            loser_call = call_holder.get(task)
+            if loser_call is not None:
+                with contextlib.suppress(Exception):
+                    loser_call.cancel()
+            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()
+
+        response = winner_task.result()
+        sketch.add(elapsed)
+        return response
+
+
+class HedgedServerStreamInterceptor(grpc.aio.UnaryStreamClientInterceptor):  # type: ignore[misc]
+    """gRPC Unary-Stream (server streaming) client interceptor with adaptive hedging.
+
+    Uses time-to-first-message (TTFM) as the hedge signal: if the primary
+    stream does not yield its first message within the estimated percentile
+    latency, a backup stream is started. Whichever stream yields first wins;
+    the loser is cancelled.
+
+    Args:
+        config: Hedge configuration. Defaults to ``HedgeConfig()``.
+    """
+
+    def __init__(self, config: HedgeConfig | None = None) -> None:
+        self._config = config or HedgeConfig()
+        self.stats = self._config.stats or Stats()
+        self._budget = TokenBucket(self._config.budget_percent, self._config.estimated_rps)
+        self._sketches: dict[str, WindowedSketch] = {}
+        self._counters: dict[str, int] = defaultdict(int)
+        self._lock = asyncio.Lock()
+
+    def _sketch_for(self, target: str) -> WindowedSketch:
+        if target not in self._sketches:
+            sketch = WindowedSketch(0.01, self._config.window_duration)
+            sketch.start_async()
+            self._sketches[target] = sketch
+        return self._sketches[target]
+
+    async def _increment_counter(self, target: str) -> int:
+        async with self._lock:
+            self._counters[target] += 1
+            return self._counters[target]
+
+    async def intercept_unary_stream(
+        self,
+        continuation: Callable[..., Any],
+        client_call_details: grpc.aio.ClientCallDetails,
+        request: Any,
+    ) -> Any:
+        """Intercept a unary-stream RPC with adaptive hedging based on TTFM."""
+        self.stats.increment_total()
+
+        target = client_call_details.method
+        sketch = self._sketch_for(target)
+        request_number = await self._increment_counter(target)
+
+        if request_number <= self._config.warmup_requests:
+            self.stats.increment_warmup()
+            hedge_delay = self._config.warmup_delay
+        else:
+            estimate = sketch.quantile(self._config.percentile)
+            hedge_delay = estimate if estimate > 0 and not math.isnan(estimate) else self._config.warmup_delay
+
+        hedge_delay = max(hedge_delay, self._config.min_delay)
+        start = time.monotonic()
+
+        # Combine ``await continuation`` (which may itself include connection
+        # setup / header round-trips) and ``call.read()`` (TTFM) in one task,
+        # so the timer reflects the real time-to-first-message.
+        call_holder: dict[asyncio.Task[Any], Any] = {}
+
+        async def invoke_and_read_first() -> tuple[Any, Any]:
+            call = await continuation(client_call_details, request)
+            current = asyncio.current_task()
+            if current is not None:
+                call_holder[current] = call
+            first_msg = await call.read()
+            return call, first_msg
+
+        # Launch primary
+        primary_task = asyncio.create_task(invoke_and_read_first())
+
+        done, _ = await asyncio.wait({primary_task}, timeout=hedge_delay)
+        if done:
+            primary_call, first_msg = primary_task.result()
+            sketch.add(time.monotonic() - start)
+            return _PrependedStream(first_msg, primary_call)
+
+        # Budget check
+        if not self._budget.try_acquire():
+            self.stats.increment_budget_exhausted()
+            primary_call, first_msg = await primary_task
+            sketch.add(time.monotonic() - start)
+            return _PrependedStream(first_msg, primary_call)
+
+        # Launch hedge
+        self.stats.increment_hedged()
+        hedge_task = asyncio.create_task(invoke_and_read_first())
+
+        done, pending = await asyncio.wait(
+            {primary_task, hedge_task},
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
+        winner_task = done.pop()
+        winner_call, first_msg = winner_task.result()
+        elapsed = time.monotonic() - start
+        sketch.add(elapsed)
+
+        if winner_task is primary_task:
+            self.stats.increment_primary_wins()
+        else:
+            self.stats.increment_hedge_wins()
+
+        # Cancel loser: cancel the underlying gRPC Call first (best-effort),
+        # then await the task so we don't leak it.
+        for task in pending:
+            loser_call = call_holder.get(task)
+            if loser_call is not None:
+                with contextlib.suppress(Exception):
+                    loser_call.cancel()
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError, Exception):
+                await task
+
+        return _PrependedStream(first_msg, winner_call)
+
+
+class _PrependedStream:
+    """Wraps a gRPC stream, prepending an already-read first message.
+
+    This allows the interceptor to consume the first message for TTFM
+    measurement while still presenting a complete stream to the caller.
+    """
+
+    def __init__(self, first_msg: Any, call: Any) -> None:
+        self._first_msg = first_msg
+        self._call = call
+        self._first_yielded = False
+
+    def __aiter__(self) -> _PrependedStream:
+        return self
+
+    async def __anext__(self) -> Any:
+        if not self._first_yielded:
+            self._first_yielded = True
+            if self._first_msg is grpc.aio.EOF:
+                raise StopAsyncIteration
+            return self._first_msg
+        msg = await self._call.read()
+        if msg is grpc.aio.EOF:
+            raise StopAsyncIteration
+        return msg

hedge/sketch/__init__.py

@@ -0,0 +1,6 @@
+"""DDSketch-based streaming quantile estimation with sliding windows."""
+
+from hedge.sketch._ddsketch import DDSketch
+from hedge.sketch._windowed import WindowedSketch
+
+__all__ = ["DDSketch", "WindowedSketch"]

hedge/sketch/_ddsketch.py

@@ -0,0 +1,173 @@
+"""DDSketch streaming quantile estimator.
+
+Based on Masson et al., "DDSketch: A fast and fully-mergeable quantile sketch
+with relative-error guarantees", VLDB 2019.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import defaultdict
+
+
+class _LogMapping:
+    """Maps values to bucket indices using logarithmic scaling.
+
+    For a given relative accuracy alpha, gamma = (1+alpha)/(1-alpha).
+    A positive value x maps to bucket index ceil(ln(x) / ln(gamma)).
+
+    The guarantee: any value in bucket i is within a factor of gamma^0.5 of
+    the bucket's representative value, giving relative error <= alpha.
+    """
+
+    __slots__ = ("gamma", "multiplier")
+
+    def __init__(self, relative_accuracy: float) -> None:
+        self.gamma = (1 + relative_accuracy) / (1 - relative_accuracy)
+        self.multiplier = 1.0 / math.log(self.gamma)
+
+    def index(self, value: float) -> int:
+        """Return the bucket index for a strictly positive value."""
+        return math.ceil(math.log(value) * self.multiplier)
+
+    def value(self, index: int) -> float:
+        """Return the representative value (geometric midpoint) for a bucket."""
+        return math.exp((index - 0.5) / self.multiplier)
+
+
+class _Store:
+    """Sparse map of bucket indices to cumulative counts."""
+
+    __slots__ = ("bins", "count")
+
+    def __init__(self) -> None:
+        self.bins: dict[int, float] = defaultdict(float)
+        self.count: float = 0.0
+
+    def add(self, index: int) -> None:
+        self.bins[index] += 1.0
+        self.count += 1.0
+
+    def merge(self, other: _Store) -> None:
+        for idx, cnt in other.bins.items():
+            self.bins[idx] += cnt
+        self.count += other.count
+
+    def reset(self) -> None:
+        self.bins = defaultdict(float)
+        self.count = 0.0
+
+
+class DDSketch:
+    """Streaming quantile sketch with relative-error guarantees.
+
+    Positive and negative values are stored in separate sparse bucket maps.
+    Zero values are counted separately. Min and max are tracked exactly.
+
+    Property: for any quantile q, the returned estimate satisfies
+        |estimate - true_value| / |true_value| <= relative_accuracy
+
+    Args:
+        relative_accuracy: Target relative accuracy. 0.01 means estimates
+            are within +/-1% of the true value. Must be in (0, 1).
+    """
+
+    def __init__(self, relative_accuracy: float = 0.01) -> None:
+        if not (0 < relative_accuracy < 1):
+            raise ValueError("relative_accuracy must be in (0, 1)")
+        self._mapping = _LogMapping(relative_accuracy)
+        self._positive = _Store()
+        self._negative = _Store()
+        self._zero_count: float = 0.0
+        self._count: int = 0
+        self._min: float = math.inf
+        self._max: float = -math.inf
+
+    @property
+    def count(self) -> int:
+        """Total number of values added."""
+        return self._count
+
+    def add(self, value: float) -> None:
+        """Record a single value. O(1) per insert.
+
+        NaN and infinite values are silently ignored.
+        """
+        if math.isnan(value) or math.isinf(value):
+            return
+
+        if value > 0:
+            self._positive.add(self._mapping.index(value))
+        elif value < 0:
+            self._negative.add(self._mapping.index(-value))
+        else:
+            self._zero_count += 1.0
+
+        self._count += 1
+        if value < self._min:
+            self._min = value
+        if value > self._max:
+            self._max = value
+
+    def quantile(self, q: float) -> float:
+        """Return the estimated value at quantile q in [0, 1].
+
+        Returns math.nan if the sketch is empty.
+
+        The estimate satisfies the relative-error guarantee:
+            |estimate - true_value| / |true_value| <= relative_accuracy
+        """
+        if self._count == 0:
+            return math.nan
+        if q <= 0:
+            return self._min
+        if q >= 1:
+            return self._max
+
+        rank: float = float(math.ceil(q * self._count))
+
+        # Negative values: iterate descending (most negative -> least negative)
+        if self._negative.count > 0:
+            cumulative = 0.0
+            for idx in sorted(self._negative.bins.keys(), reverse=True):
+                cumulative += self._negative.bins[idx]
+                if cumulative >= rank:
+                    return -self._mapping.value(idx)
+            rank -= self._negative.count
+
+        # Zero values
+        if self._zero_count > 0:
+            rank -= self._zero_count
+            if rank <= 0:
+                return 0.0
+
+        # Positive values: iterate ascending (least positive -> most positive)
+        if self._positive.count > 0:
+            cumulative = 0.0
+            for idx in sorted(self._positive.bins.keys()):
+                cumulative += self._positive.bins[idx]
+                if cumulative >= rank:
+                    return self._mapping.value(idx)
+
+        return self._max
+
+    def merge(self, other: DDSketch) -> None:
+        """Combine other into self. The merge is exact: no additional error accumulates."""
+        self._positive.merge(other._positive)
+        self._negative.merge(other._negative)
+        self._zero_count += other._zero_count
+        self._count += other._count
+        if other._count > 0:
+            if other._min < self._min:
+                self._min = other._min
+            if other._max > self._max:
+                self._max = other._max
+
+    def reset(self) -> None:
+        """Clear all state. Used for windowed / tumbling-window decay."""
+        self._positive.reset()
+        self._negative.reset()
+        self._zero_count = 0.0
+        self._count = 0
+        self._min = math.inf
+        self._max = -math.inf