lorebase-keeper 0.2.3__py3-none-any.whl

lorebase_keeper/ratelimit/limiter.py

@@ -0,0 +1,149 @@
+"""Rate limiter: sorts requests into classes, takes from the principal's bucket, counts refusals.
+
+A principal has one bucket per class (`reads`, `writes`, `searches`, `context`, `admin`), keyed
+by the API key id or, for users, the principal id. A Domain that overrides the bucket of a
+class gets a bucket of its own for it.
+
+Failed authentications have buckets of their own (`auth_failures`), in a separate store so
+that unauthenticated traffic cannot push principals' buckets out: one per client address (IPv6
+by network) and one per presented API key prefix.
+"""
+
+from __future__ import annotations
+
+import ipaddress
+import time
+from collections import OrderedDict, deque
+from collections.abc import Callable, Collection, Sequence
+
+from opentelemetry import metrics
+from opentelemetry.metrics import MeterProvider
+
+from lorebase_keeper.client_address import IPAddress
+from lorebase_keeper.ratelimit.bucket import BucketSpec, Decision, InMemoryRateStore, RateStore
+from lorebase_keeper.ratelimit.config import RateLimitSettings, RouteClass
+
+AUTH_FAILURES = "auth_failures"
+
+SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS"})
+
+
+class RateLimiter:
+    """Applies `RateLimitSettings` on top of a `RateStore`."""
+
+    def __init__(
+        self,
+        settings: RateLimitSettings,
+        store: RateStore,
+        *,
+        failure_store: RateStore | None = None,
+        meter_provider: MeterProvider | None = None,
+    ) -> None:
+        self._settings = settings
+        self._store = store
+        self._failure_store = failure_store or InMemoryRateStore(settings.max_buckets)
+        self._failure_spec = BucketSpec.from_settings(settings.auth_failures)
+        meter = (meter_provider or metrics.get_meter_provider()).get_meter("lorebase.keeper")
+        self._limited = meter.create_counter(
+            "lorebase.http.rate_limited",
+            unit="{request}",
+            description="Requests refused with 429 because a rate limit was reached.",
+        )
+
+    def classify(self, method: str, route_path: str, tags: Collection[str] = ()) -> RouteClass:
+        """The class of a request, from its method, route template and the route's tags."""
+        routes = self._settings.routes
+        if not set(tags).isdisjoint(routes.admin_tags):
+            return RouteClass.ADMIN
+        if method in SAFE_METHODS:
+            return RouteClass.READS
+        if method == "POST":
+            if route_path.endswith(routes.search_suffixes):
+                return RouteClass.SEARCHES
+            if route_path.endswith(routes.context_suffixes):
+                return RouteClass.CONTEXT
+            if route_path.endswith(routes.read_post_suffixes):
+                return RouteClass.READS
+        return RouteClass.WRITES
+
+    async def check(self, who: str, route_class: RouteClass, domain_id: str | None) -> Decision:
+        """Take one token from the bucket of `who` for this class (and Domain override)."""
+        bucket = self._settings.bucket(route_class, domain_id)
+        key = f"{who}|{route_class.value}"
+        if self._settings.has_override(route_class, domain_id):
+            key += f"|domain:{domain_id}"
+        decision = await self._store.take(key, BucketSpec.from_settings(bucket))
+        if not decision.allowed:
+            self._limited.add(1, {"class": route_class.value})
+        return decision
+
+    def failure_keys(self, address: IPAddress | None, key_prefix: str | None) -> tuple[str, ...]:
+        """The failed-authentication buckets of a request: its client address (an IPv6
+        address by its network) and, unless turned off, its API key prefix, each when known."""
+        keys: list[str] = []
+        if isinstance(address, ipaddress.IPv6Address):
+            length = self._settings.auth_failures.ipv6_prefix_length
+            keys.append(f"addr:{ipaddress.ip_network(f'{address}/{length}', strict=False)}")
+        elif address is not None:
+            keys.append(f"addr:{address}")
+        if key_prefix is not None and self._settings.auth_failures.per_key_prefix_enabled:
+            keys.append(f"prefix:{key_prefix}")
+        return tuple(f"{AUTH_FAILURES}|{key}" for key in keys)
+
+    async def failures_exhausted(self, keys: Sequence[str]) -> Decision | None:
+        """The refusal (the longest wait) when one of the buckets has no failure left, else
+        None. Takes nothing."""
+        refused: Decision | None = None
+        for key in keys:
+            decision = await self._failure_store.peek(key, self._failure_spec)
+            if not decision.allowed and (
+                refused is None or decision.retry_after_s > refused.retry_after_s
+            ):
+                refused = decision
+        if refused is not None:
+            self._limited.add(1, {"class": AUTH_FAILURES})
+        return refused
+
+    async def record_failure(self, keys: Sequence[str]) -> None:
+        """Take one token from each bucket for a failed authentication."""
+        for key in keys:
+            await self._failure_store.take(key, self._failure_spec)
+
+
+class ThrottleWatch:
+    """Counts the refusals of each principal in a sliding window.
+
+    `record` returns the count when it reaches the threshold and starts counting again, so a
+    principal that keeps hammering the API produces one report per `threshold` refusals.
+    """
+
+    def __init__(
+        self,
+        threshold: int,
+        window_s: float,
+        max_tracked: int,
+        *,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self._threshold = threshold
+        self._window_s = window_s
+        self._max = max_tracked
+        self._clock = clock
+        self._seen: OrderedDict[str, deque[float]] = OrderedDict()
+
+    def record(self, who: str) -> int | None:
+        """Note a refusal of `who`; the number of refusals in the window when it is time to
+        report them."""
+        if self._threshold <= 0:
+            return None
+        now = self._clock()
+        times = self._seen.pop(who, None) or deque()
+        while times and now - times[0] > self._window_s:
+            times.popleft()
+        times.append(now)
+        if len(times) >= self._threshold:
+            return len(times)
+        self._seen[who] = times
+        while len(self._seen) > self._max:
+            self._seen.popitem(last=False)
+        return None

lorebase_keeper/ratelimit/quotas.py

@@ -0,0 +1,204 @@
+"""Daily quotas of API keys: requests and context-pack tokens per UTC day.
+
+Counting is done in memory and written to `api_key_usage` at most every
+`quota.flush_interval_s` per key. A write adds this process's use to the stored row and returns
+the total of every process, so several Keeper processes converge on the same figure and a quota
+holds across restarts. Between writes a process enforces the quota on what it knows (the
+stored total as of its last write, plus its own use since), which lets a key exceed a quota by
+about one interval of traffic. A store failure never refuses a request: the quota is enforced
+on what is known and the write is retried.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from datetime import UTC, date, datetime, timedelta
+from datetime import time as day_time
+from typing import Literal
+
+from opentelemetry import metrics
+from opentelemetry.metrics import MeterProvider
+
+from lorebase_keeper.auth.ports import AuthStore
+from lorebase_keeper.ratelimit.config import QuotaSettings
+
+logger = logging.getLogger("lorebase.keeper.quota")
+
+QuotaKind = Literal["requests", "context_tokens"]
+
+
+@dataclass(frozen=True, slots=True)
+class QuotaSnapshot:
+    """What a key has used today and may still use. A limit of `None` is no limit."""
+
+    day: date
+    resets_at: datetime
+    requests_used: int
+    requests_limit: int | None
+    tokens_used: int
+    tokens_limit: int | None
+
+    def exceeded(self) -> QuotaKind | None:
+        """The quota that is used up, if any."""
+        if self.requests_limit is not None and self.requests_used >= self.requests_limit:
+            return "requests"
+        if self.tokens_limit is not None and self.tokens_used >= self.tokens_limit:
+            return "context_tokens"
+        return None
+
+
+@dataclass
+class _Entry:
+    day: date
+    stored_requests: int = 0
+    stored_tokens: int = 0
+    pending_requests: int = 0
+    pending_tokens: int = 0
+    loaded: bool = False
+    flushing: bool = False
+    last_flush: float = field(default=float("-inf"))
+
+
+class QuotaTracker:
+    """Counts the use of API keys against `QuotaSettings`."""
+
+    def __init__(
+        self,
+        settings: QuotaSettings,
+        store: AuthStore,
+        *,
+        clock: Callable[[], datetime] = lambda: datetime.now(UTC),
+        monotonic: Callable[[], float] = time.monotonic,
+        meter_provider: MeterProvider | None = None,
+    ) -> None:
+        self._settings = settings
+        self._store = store
+        self._clock = clock
+        self._monotonic = monotonic
+        self._entries: dict[str, _Entry] = {}
+        meter = (meter_provider or metrics.get_meter_provider()).get_meter("lorebase.keeper")
+        self._exceeded = meter.create_counter(
+            "lorebase.http.quota_exceeded",
+            unit="{request}",
+            description="Requests refused with 429 because a daily quota was used up.",
+        )
+
+    # Use --------------------------------------------------------------------------------
+
+    async def admit(self, key_id: str) -> tuple[QuotaKind | None, QuotaSnapshot]:
+        """Count one request of the key unless a quota is used up.
+
+        Returns the quota that refuses the request (nothing is counted then) and the state of
+        the key.
+        """
+        entry = await self._entry(key_id)
+        snapshot = self._snapshot(key_id, entry)
+        refused = snapshot.exceeded()
+        if refused is not None:
+            self._exceeded.add(1, {"quota": refused})
+            return refused, snapshot
+        entry.pending_requests += 1
+        await self._flush_when_due(key_id, entry)
+        return None, self._snapshot(key_id, entry)
+
+    async def check_tokens(self, key_id: str) -> QuotaSnapshot | None:
+        """The key's state when its token quota is used up, otherwise `None`."""
+        entry = await self._entry(key_id)
+        snapshot = self._snapshot(key_id, entry)
+        if snapshot.tokens_limit is not None and snapshot.tokens_used >= snapshot.tokens_limit:
+            self._exceeded.add(1, {"quota": "context_tokens"})
+            return snapshot
+        return None
+
+    async def charge_tokens(self, key_id: str, tokens: int) -> None:
+        """Count context-pack tokens delivered to the key."""
+        if tokens <= 0:
+            return
+        entry = await self._entry(key_id)
+        entry.pending_tokens += tokens
+        await self._flush_when_due(key_id, entry)
+
+    async def snapshot(self, key_id: str) -> QuotaSnapshot:
+        """The key's use today, current across processes."""
+        entry = await self._entry(key_id)
+        await self._flush(key_id, entry)
+        return self._snapshot(key_id, entry)
+
+    async def aclose(self) -> None:
+        """Write what has not been written yet."""
+        for key_id, entry in list(self._entries.items()):
+            await self._flush(key_id, entry)
+
+    # Internals --------------------------------------------------------------------------
+
+    def _today(self) -> date:
+        return self._clock().astimezone(UTC).date()
+
+    def _snapshot(self, key_id: str, entry: _Entry) -> QuotaSnapshot:
+        return QuotaSnapshot(
+            day=entry.day,
+            resets_at=datetime.combine(entry.day + timedelta(days=1), day_time.min, tzinfo=UTC),
+            requests_used=entry.stored_requests + entry.pending_requests,
+            requests_limit=self._settings.requests_limit(key_id),
+            tokens_used=entry.stored_tokens + entry.pending_tokens,
+            tokens_limit=self._settings.tokens_limit(key_id),
+        )
+
+    async def _entry(self, key_id: str) -> _Entry:
+        """The key's counters for today, read from the store the first time."""
+        today = self._today()
+        entry = self._entries.get(key_id)
+        if entry is not None and entry.day != today:
+            await self._flush(key_id, entry)  # what the old day counted goes to the old day
+            entry = None
+        if entry is None:
+            entry = self._entries[key_id] = _Entry(day=today)
+        if not entry.loaded:
+            await self._refresh(key_id, entry)
+        return entry
+
+    async def _refresh(self, key_id: str, entry: _Entry) -> None:
+        try:
+            async with self._store.unit() as unit:
+                usage = await unit.usage.get(key_id, entry.day)
+        except Exception:
+            logger.warning("cannot read the usage of an api key", exc_info=True)
+            return
+        entry.stored_requests, entry.stored_tokens = usage.requests, usage.context_tokens
+        entry.loaded = True
+
+    async def _flush_when_due(self, key_id: str, entry: _Entry) -> None:
+        if self._monotonic() - entry.last_flush >= self._settings.flush_interval_s:
+            await self._flush(key_id, entry)
+
+    async def _flush(self, key_id: str, entry: _Entry) -> None:
+        """Add the pending use to the stored row and learn the total of all processes."""
+        if entry.flushing:
+            return
+        entry.flushing = True
+        entry.last_flush = self._monotonic()
+        requests, tokens = entry.pending_requests, entry.pending_tokens
+        try:
+            async with self._store.unit() as unit:
+                if requests or tokens:
+                    total = await unit.usage.add(
+                        key_id,
+                        entry.day,
+                        requests=requests,
+                        context_tokens=tokens,
+                        at=self._clock(),
+                    )
+                else:
+                    total = await unit.usage.get(key_id, entry.day)
+        except Exception:
+            logger.warning("cannot write the usage of an api key", exc_info=True)
+            return
+        finally:
+            entry.flushing = False
+        entry.pending_requests -= requests
+        entry.pending_tokens -= tokens
+        entry.stored_requests, entry.stored_tokens = total.requests, total.context_tokens
+        entry.loaded = True

lorebase_keeper/ratelimit/shared.py

@@ -0,0 +1,242 @@
+"""Rate limit buckets shared by every Keeper, in Postgres (`keeper.limits.rate.backend:
+postgres`).
+
+`PostgresRateStore` implements `RateStore` over `lorebase_store.ratelimit`. Requests are served
+from tokens leased from the shared bucket (`postgres.lease_tokens`, at most a quarter of the
+bucket), so most requests cost no query; a request that finds its lease used up or expired
+(`postgres.lease_ttl_s`) takes a new one, one query per bucket at a time. Leases come out of
+the shared bucket, so the Keepers together never exceed a limit; tokens of an expired lease
+are lost, which only makes it stricter.
+
+Failed-authentication buckets lease one token per failure and are checked before every
+authentication from a reading cached for `postgres.check_cache_s`. They fail closed: when the
+database cannot be read in `postgres.auth_failures_timeout_s`, a request presenting a
+credential is answered 429, unless a reading of its bucket at most `postgres.stale_reading_s`
+old still allows it (so an outage refuses unknown clients, not every client). Other buckets
+follow `postgres.on_error`: `local` (the default) enforces the limits in this process until
+the database answers again, `allow` lets requests through, `deny` refuses them.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import math
+import time
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+from opentelemetry import metrics
+from opentelemetry.metrics import MeterProvider
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from lorebase_keeper.ratelimit.bucket import BucketSpec, BucketState, Decision, RateStore, take
+from lorebase_keeper.ratelimit.config import PostgresRateSettings
+from lorebase_store import ratelimit as buckets
+
+logger = logging.getLogger("lorebase.keeper.limits")
+
+LEASE_SHARE = 4
+"""A lease is at most 1/LEASE_SHARE of its bucket, so several Keepers can hold one at once."""
+PRUNE_ROUNDS = 10
+"""Batches one prune deletes at most."""
+
+
+@dataclass(slots=True)
+class _Local:
+    """What this process knows of one shared bucket."""
+
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    leased: float = 0.0
+    expires: float = float("-inf")
+    shared: float = 0.0
+    """Tokens left in the shared bucket at the last lease."""
+    reading: buckets.BucketReading | None = None
+    read_at: float = float("-inf")
+
+
+def _decision(spec: BucketSpec, tokens: float, allowed: bool, cost: float) -> Decision:
+    """A decision with headers computed from the tokens known to be available."""
+    tokens = max(0.0, min(float(spec.capacity), tokens))
+    retry = 0 if allowed else max(1, math.ceil((cost - tokens) / spec.refill_per_s))
+    return Decision(
+        allowed=allowed,
+        limit=spec.capacity,
+        remaining=math.floor(tokens),
+        reset_s=math.ceil((spec.capacity - tokens) / spec.refill_per_s),
+        retry_after_s=retry,
+    )
+
+
+class PostgresRateStore:
+    """`RateStore` over the shared `rate_limit_buckets` table (see the module doc).
+
+    `fail_closed` refuses when the database fails (failed-authentication buckets); otherwise
+    `settings.on_error` decides, with `fallback` the store of `local`. `lease_tokens` overrides
+    the lease size (1 for failed-authentication buckets). `prune` makes this store delete the
+    rows of full buckets now and then.
+    """
+
+    def __init__(
+        self,
+        sessions: async_sessionmaker[AsyncSession],
+        settings: PostgresRateSettings,
+        *,
+        max_buckets: int,
+        fallback: RateStore | None = None,
+        fail_closed: bool = False,
+        lease_tokens: int | None = None,
+        prune: bool = False,
+        label: str = "principals",
+        clock: Callable[[], float] = time.monotonic,
+        meter_provider: MeterProvider | None = None,
+    ) -> None:
+        if not fail_closed and settings.on_error == "local" and fallback is None:
+            raise ValueError("on_error local needs a fallback store")
+        self._sessions = sessions
+        self._settings = settings
+        self._max = max_buckets
+        self._fallback = fallback
+        self._fail_closed = fail_closed
+        self._timeout = settings.auth_failures_timeout_s if fail_closed else settings.timeout_s
+        self._lease_tokens = lease_tokens or settings.lease_tokens
+        self._prune = prune
+        self._label = label
+        self._clock = clock
+        self._local: OrderedDict[str, _Local] = OrderedDict()
+        self._last_prune = clock()
+        self._pruning: asyncio.Task[None] | None = None
+        meter = (meter_provider or metrics.get_meter_provider()).get_meter("lorebase.keeper")
+        self._queries = meter.create_counter(
+            "lorebase.ratelimit.shared.queries",
+            unit="{query}",
+            description="Queries to the shared rate limit buckets, by store and outcome.",
+        )
+
+    def _entry(self, key: str) -> _Local:
+        entry = self._local.get(key)
+        if entry is None:
+            entry = self._local[key] = _Local()
+        self._local.move_to_end(key)
+        while len(self._local) > self._max:
+            self._local.popitem(last=False)
+        return entry
+
+    def _lease_size(self, spec: BucketSpec) -> int:
+        return max(1, min(self._lease_tokens, spec.capacity // LEASE_SHARE))
+
+    def _from_lease(self, entry: _Local, cost: float) -> bool:
+        if entry.expires > self._clock() and entry.leased >= cost:
+            entry.leased -= cost
+            return True
+        return False
+
+    async def take(self, key: str, spec: BucketSpec, cost: float = 1.0) -> Decision:
+        entry = self._entry(key)
+        if self._from_lease(entry, cost):
+            return _decision(spec, entry.leased + entry.shared, True, cost)
+        async with entry.lock:
+            if self._from_lease(entry, cost):  # another request leased meanwhile
+                return _decision(spec, entry.leased + entry.shared, True, cost)
+            self._maybe_prune()
+            try:
+                async with asyncio.timeout(self._timeout):
+                    async with self._sessions() as session, session.begin():
+                        lease = await buckets.lease(
+                            session,
+                            key,
+                            capacity=spec.capacity,
+                            refill_per_s=spec.refill_per_s,
+                            lease=self._lease_size(spec),
+                            cost=cost,
+                        )
+            except Exception as error:
+                return await self._failed(error, "take", key, spec, cost)
+            self._queries.add(1, {"store": self._label, "outcome": "ok"})
+            entry.shared = lease.remaining
+            entry.reading = buckets.BucketReading(lease.remaining, lease.now, lease.now)
+            entry.read_at = self._clock()
+            if lease.granted < cost:
+                entry.leased, entry.expires = 0.0, float("-inf")
+                return _decision(spec, lease.remaining, False, cost)
+            entry.leased = lease.granted - cost
+            entry.expires = self._clock() + self._settings.lease_ttl_s
+            return _decision(spec, entry.leased + entry.shared, True, cost)
+
+    async def peek(self, key: str, spec: BucketSpec, cost: float = 1.0) -> Decision:
+        entry = self._entry(key)
+        now = self._clock()
+        if entry.reading is None or now - entry.read_at > self._settings.check_cache_s:
+            try:
+                async with asyncio.timeout(self._timeout):
+                    async with self._sessions() as session, session.begin():
+                        reading = await buckets.read(session, key)
+            except Exception as error:
+                stale = entry.reading is not None and (
+                    now - entry.read_at <= self._settings.stale_reading_s
+                )
+                if not stale:
+                    return await self._failed(error, "peek", key, spec, cost)
+                # A recent reading of this bucket still decides: only clients the database
+                # never saw recently are refused while it does not answer.
+                self._queries.add(1, {"store": self._label, "outcome": "stale"})
+            else:
+                self._queries.add(1, {"store": self._label, "outcome": "ok"})
+                entry.reading, entry.read_at = reading, self._clock()
+        known = entry.reading
+        assert known is not None
+        state = None if known.tokens is None else BucketState(known.tokens, known.updated)
+        _, decision = take(state, spec, known.now + (self._clock() - entry.read_at), cost)
+        return decision
+
+    async def _failed(
+        self, error: Exception, operation: str, key: str, spec: BucketSpec, cost: float
+    ) -> Decision:
+        """The decision when the shared buckets cannot be reached."""
+        reason = "timeout" if isinstance(error, TimeoutError) else type(error).__name__
+        self._queries.add(1, {"store": self._label, "outcome": reason})
+        logger.warning(
+            "shared rate limit buckets unavailable",
+            exc_info=not isinstance(error, TimeoutError),
+            extra={"store": self._label, "reason": reason},
+        )
+        mode = "deny" if self._fail_closed else self._settings.on_error
+        if mode == "local" and self._fallback is not None:
+            if operation == "take":
+                return await self._fallback.take(key, spec, cost)
+            return await self._fallback.peek(key, spec, cost)
+        if mode == "allow":
+            return _decision(spec, float(spec.capacity), True, cost)
+        retry = self._settings.error_retry_after_s
+        return Decision(False, spec.capacity, 0, retry, retry)
+
+    def _maybe_prune(self) -> None:
+        if not self._prune or self._pruning is not None:
+            return
+        if self._clock() - self._last_prune < self._settings.prune_interval_s:
+            return
+        self._last_prune = self._clock()
+        self._pruning = asyncio.create_task(self._prune_once(), name="ratelimit-prune")
+
+    async def _prune_once(self) -> None:
+        try:
+            batch = self._settings.prune_batch
+            async with asyncio.timeout(max(self._settings.timeout_s, 5.0) * PRUNE_ROUNDS):
+                for _ in range(PRUNE_ROUNDS):  # until a batch comes back short
+                    async with self._sessions() as session, session.begin():
+                        if await buckets.prune(session, limit=batch) < batch:
+                            break
+        except Exception:
+            logger.warning("cannot prune the shared rate limit buckets", exc_info=True)
+        finally:
+            self._pruning = None
+
+    async def aclose(self) -> None:
+        """Stop a running prune."""
+        task, self._pruning = self._pruning, None
+        if task is not None:
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError, Exception):
+                await task