スライディングウィンドウとバーストクレジットを備えたレートリミッター

"Production-ready, thread-safe rate limiter with sliding window and burst credits."

from __future__ import annotations

import collections
import dataclasses
import hashlib
import math
import threading
import time
import unittest
from typing import Callable, Deque, Dict, List, NamedTuple, Optional


@dataclasses.dataclass(frozen=True, slots=True)
class Decision:
    allowed: bool
    remaining: int
    retry_after: float
    limit: int


class _WindowEntry(NamedTuple):
    timestamp: float
    cost: int


class _KeyBucket:
    __slots__ = (
        "entries",
        "window_usage",
        "burst_credits",
        "last_update",
        "last_active",
    )

    def __init__(self, initial_credits: float, now: float) -> None:
        self.entries: Deque[_WindowEntry] = collections.deque()
        self.window_usage: int = 0
        self.burst_credits: float = float(initial_credits)
        self.last_update: float = now
        self.last_active: float = now

    def clean_window(self, cutoff: float) -> None:
        while self.entries and self.entries[0].timestamp <= cutoff:
            expired = self.entries.popleft()
            self.window_usage -= expired.cost
        if self.window_usage < 0:
            self.window_usage = 0

    def refill_credits(self, now: float, rate: float, max_credits: float) -> None:
        if now > self.last_update:
            if rate > 0.0 and self.burst_credits < max_credits:
                added = (now - self.last_update) * rate
                self.burst_credits = min(max_credits, self.burst_credits + added)
            self.last_update = now

    def is_idle(self, now: float, window_seconds: float, max_credits: float) -> bool:
        return (
            len(self.entries) == 0
            and math.isclose(self.burst_credits, max_credits, abs_tol=1e-9)
            and (now - self.last_active) >= window_seconds
        )


class _Shard:
    __slots__ = ("lock", "buckets", "access_order")

    def __init__(self) -> None:
        self.lock = threading.Lock()
        self.buckets: Dict[str, _KeyBucket] = {}
        self.access_order: collections.OrderedDict[str, None] = collections.OrderedDict()


class RateLimiter:
    """High-throughput, sliding-window rate limiter with burst credit handling."""

    def __init__(
        self,
        limit: int,
        window_seconds: float,
        burst_credits: int = 0,
        clock: Callable[[], float] = time.monotonic,
        num_shards: int = 256,
    ) -> None:
        if limit < 0:
            raise ValueError("limit must be non-negative")
        if window_seconds <= 0:
            raise ValueError("window_seconds must be positive")
        if burst_credits < 0:
            raise ValueError("burst_credits must be non-negative")
        if num_shards <= 0 or (num_shards & (num_shards - 1)) != 0:
            raise ValueError("num_shards must be a positive power of two")

        self.limit = int(limit)
        self.window_seconds = float(window_seconds)
        self.burst_credits = int(burst_credits)
        self.clock = clock
        self._rate = (self.limit / self.window_seconds) if self.limit > 0 else 0.0

        self._num_shards = num_shards
        self._mask = num_shards - 1
        self._shards: List[_Shard] = [_Shard() for _ in range(num_shards)]

    def _get_shard(self, key: str) -> _Shard:
        digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
        index = int.from_bytes(digest, byteorder="little") & self._mask
        return self._shards[index]

    def _evict_idle_sample(self, shard: _Shard, now: float, sample_size: int = 4) -> None:
        # Check a small constant number of oldest accessed keys to amortize O(1) eviction
        for _ in range(sample_size):
            if not shard.access_order:
                break
            oldest_key = next(iter(shard.access_order))
            bucket = shard.buckets.get(oldest_key)
            if bucket is None:
                shard.access_order.pop(oldest_key, None)
                continue
            bucket.clean_window(now - self.window_seconds)
            bucket.refill_credits(now, self._rate, float(self.burst_credits))
            if bucket.is_idle(now, self.window_seconds, float(self.burst_credits)):
                del shard.buckets[oldest_key]
                shard.access_order.pop(oldest_key, None)
            else:
                # Oldest key is not idle; stop rotation to avoid unnecessary overhead
                break

    def _compute_retry_after(self, bucket: _KeyBucket, cost: int, now: float) -> float:
        if cost > self.limit + self.burst_credits:
            return float("inf")
        if cost <= 0:
            return 0.0

        cutoff = now - self.window_seconds
        sim_entries = collections.deque(bucket.entries)
        sim_usage = bucket.window_usage
        sim_credits = bucket.burst_credits
        last_t = now
        rate = self._rate
        max_credits = float(self.burst_credits)

        while True:
            window_avail = max(0, self.limit - sim_usage)
            needed_burst = max(0, cost - window_avail)
            if sim_credits >= needed_burst:
                return max(0.0, last_t - now)

            # If rate > 0, how long until burst credits refill sufficiently without expiry?
            dt_refill = (needed_burst - sim_credits) / rate if rate > 0.0 else float("inf")

            if not sim_entries:
                if rate > 0.0:
                    return max(0.0, (last_t + dt_refill) - now)
                return float("inf")

            next_entry = sim_entries[0]
            entry_expiry_time = next_entry.timestamp + self.window_seconds
            dt_to_expiry = max(0.0, entry_expiry_time - last_t)

            if rate > 0.0 and dt_refill <= dt_to_expiry:
                return max(0.0, (last_t + dt_refill) - now)

            # Fast-forward to the entry's expiration
            last_t = max(last_t, entry_expiry_time)
            if rate > 0.0:
                sim_credits = min(max_credits, sim_credits + dt_to_expiry * rate)
            expired = sim_entries.popleft()
            sim_usage = max(0, sim_usage - expired.cost)

    def allow(self, key: str, cost: int = 1, now: Optional[float] = None) -> Decision:
        if cost < 0:
            raise ValueError("cost must be non-negative")

        current_time = float(self.clock() if now is None else now)
        shard = self._get_shard(key)

        with shard.lock:
            self._evict_idle_sample(shard, current_time)

            bucket = shard.buckets.get(key)
            if bucket is None:
                bucket = _KeyBucket(float(self.burst_credits), current_time)
                shard.buckets[key] = bucket

            shard.access_order[key] = None
            shard.access_order.move_to_end(key)

            # Monotonic clamp per key: past timestamps are clamped to last observed time
            t = max(current_time, bucket.last_active)
            bucket.last_active = t

            bucket.clean_window(t - self.window_seconds)
            bucket.refill_credits(t, self._rate, float(self.burst_credits))

            window_avail = max(0, self.limit - bucket.window_usage)
            credits_avail = bucket.burst_credits
            total_avail = window_avail + int(credits_avail)

            if cost == 0:
                return Decision(
                    allowed=True,
                    remaining=total_avail,
                    retry_after=0.0,
                    limit=self.limit,
                )

            if cost > self.limit + self.burst_credits:
                return Decision(
                    allowed=False,
                    remaining=total_avail,
                    retry_after=float("inf"),
                    limit=self.limit,
                )

            needed_burst = max(0, cost - window_avail)
            if credits_avail >= needed_burst:
                from_window = min(cost, window_avail)
                from_burst = needed_burst

                if from_window > 0:
                    bucket.entries.append(_WindowEntry(timestamp=t, cost=from_window))
                    bucket.window_usage += from_window
                bucket.burst_credits -= from_burst

                new_window_avail = max(0, self.limit - bucket.window_usage)
                new_total_avail = new_window_avail + int(bucket.burst_credits)
                return Decision(
                    allowed=True,
                    remaining=new_total_avail,
                    retry_after=0.0,
                    limit=self.limit,
                )
            else:
                retry_after = self._compute_retry_after(bucket, cost, t)
                return Decision(
                    allowed=False,
                    remaining=total_avail,
                    retry_after=retry_after,
                    limit=self.limit,
                )

    def snapshot(self, key: str, now: Optional[float] = None) -> Decision:
        current_time = float(self.clock() if now is None else now)
        shard = self._get_shard(key)

        with shard.lock:
            bucket = shard.buckets.get(key)
            if bucket is None:
                return Decision(
                    allowed=True,
                    remaining=self.limit + self.burst_credits,
                    retry_after=0.0,
                    limit=self.limit,
                )

            t = max(current_time, bucket.last_active)
            bucket.clean_window(t - self.window_seconds)
            bucket.refill_credits(t, self._rate, float(self.burst_credits))

            window_avail = max(0, self.limit - bucket.window_usage)
            remaining = window_avail + int(bucket.burst_credits)
            return Decision(
                allowed=remaining > 0,
                remaining=remaining,
                retry_after=0.0 if remaining > 0 else self._compute_retry_after(bucket, 1, t),
                limit=self.limit,
            )


"""Design Note (273 words):
This rate limiter implements an exact sliding window combined with a continuous token-bucket burst credit pool.

Data Structure & Accuracy Tradeoff:
Each client key maintains a `_KeyBucket` containing a deque of `(timestamp, cost)` tuples tracking consumed window quota, an integer accumulator `window_usage`, and a float `burst_credits`. An exact sliding window is chosen over approximate bucketing or leaky-bucket models to guarantee zero boundary error: request timestamps count strictly inside `(t - window_seconds, t]`. Multi-cost requests (`cost > 1`) are admitted atomically: available window quota is claimed first, and any deficit is covered by continuous burst credits refilled at `limit / window_seconds`. If the total available capacity is insufficient, the request is rejected all-or-nothing. Deque memory per active key scales strictly as O(N) where N <= limit, bounded and compact.

Locking Strategy:
To avoid a global contention bottleneck and prevent concurrent requests from serializing the whole gateway, the key space is partitioned across 256 independent shards via Blake2b hashing. Each shard maintains its own `threading.Lock`, dictionary, and access order. Operations on distinct keys map to independent locks with high probability, minimizing lock contention while fully synchronizing per-key concurrent accesses without deadlocks.

Memory Hygiene & Eviction:
To protect against unbounded memory growth from millions of one-shot keys, each shard maintains an `OrderedDict` of key access order. On every operation, an amortized O(1) cleanup samples the oldest keys in the shard. If a key's window deque is empty, its burst credits are fully regenerated, and it has been idle for at least `window_seconds`, it is permanently evicted. Inactive keys naturally expire without background threads.
"""


class RateLimiterTests(unittest.TestCase):
    def setUp(self) -> None:
        self.current_time = 1000.0

    def fake_clock(self) -> float:
        return self.current_time

    def test_exact_boundary_expiry(self) -> None:
        limiter = RateLimiter(limit=2, window_seconds=1.0, burst_credits=0, clock=self.fake_clock)
        d1 = limiter.allow("client-1", cost=1)
        self.assertTrue(d1.allowed)
        self.assertEqual(d1.remaining, 1)

        self.current_time += 0.5
        d2 = limiter.allow("client-1", cost=1)
        self.assertTrue(d2.allowed)
        self.assertEqual(d2.remaining, 0)

        # At exact boundary t = 1001.0, window (1000.0, 1001.0] still includes t=1000.0
        self.current_time = 1001.0
        d3 = limiter.allow("client-1", cost=1)
        self.assertFalse(d3.allowed)

        # Just past boundary t = 1001.000001, t=1000.0 is evicted
        self.current_time = 1001.000001
        d4 = limiter.allow("client-1", cost=1)
        self.assertTrue(d4.allowed)

    def test_burst_credit_refill(self) -> None:
        # limit=10, window=10.0 -> rate=1.0 credit/sec. burst_credits=5
        limiter = RateLimiter(limit=10, window_seconds=10.0, burst_credits=5, clock=self.fake_clock)
        # Consume all 10 window + 5 burst credits
        d = limiter.allow("burst-client", cost=15)
        self.assertTrue(d.allowed)
        self.assertEqual(d.remaining, 0)

        # Immediate next request rejected
        d_fail = limiter.allow("burst-client", cost=1)
        self.assertFalse(d_fail.allowed)

        # Advance 2.5s -> 2.5 credits refilled (integer remaining is 2)
        self.current_time += 2.5
        snap = limiter.snapshot("burst-client", now=self.current_time)
        self.assertEqual(snap.remaining, 2)

        # Advance another 2.5s -> total 5.0 credits refilled (capped at burst_credits=5)
        self.current_time += 2.5
        d_burst = limiter.allow("burst-client", cost=5)
        self.assertTrue(d_burst.allowed)

    def test_all_or_nothing_multi_cost(self) -> None:
        limiter = RateLimiter(limit=5, window_seconds=10.0, burst_credits=2, clock=self.fake_clock)
        # Allowed total = 7. Request cost=8 exceeds limit + burst_credits
        d_too_large = limiter.allow("client-multi", cost=8)
        self.assertFalse(d_too_large.allowed)
        self.assertEqual(d_too_large.remaining, 7)
        self.assertEqual(d_too_large.retry_after, float("inf"))

        # Key has remaining 7. Request cost=6 consumes 5 window and 1 burst
        d_consume = limiter.allow("client-multi", cost=6)
        self.assertTrue(d_consume.allowed)
        self.assertEqual(d_consume.remaining, 1)

        # Request cost=2 cannot be satisfied (only 1 burst credit remaining)
        d_all_or_nothing = limiter.allow("client-multi", cost=2)
        self.assertFalse(d_all_or_nothing.allowed)
        self.assertEqual(d_all_or_nothing.remaining, 1)  # unchanged state

    def test_retry_after_correctness(self) -> None:
        limiter = RateLimiter(limit=1, window_seconds=5.0, burst_credits=0, clock=self.fake_clock)
        d1 = limiter.allow("k", cost=1)
        self.assertTrue(d1.allowed)

        self.current_time += 2.0
        d2 = limiter.allow("k", cost=1)
        self.assertFalse(d2.allowed)
        # Request was at 1000.0, expires at 1005.0. Current time is 1002.0 -> retry_after = 3.0
        self.assertAlmostEqual(d2.retry_after, 3.0, places=5)

    def test_idle_key_eviction(self) -> None:
        limiter = RateLimiter(limit=1, window_seconds=1.0, burst_credits=0, clock=self.fake_clock)
        shard = limiter._get_shard("ephemeral-1")

        limiter.allow("ephemeral-1", cost=1)
        self.assertIn("ephemeral-1", shard.buckets)

        # Advance past window_seconds
        self.current_time += 1.5
        # Eviction is amortized during traffic on that shard
        limiter.allow("ephemeral-2", cost=1)
        self.assertNotIn("ephemeral-1", shard.buckets)

    def test_non_monotonic_and_zero_limit_edge_cases(self) -> None:
        # limit == 0 with burst credits
        limiter = RateLimiter(limit=0, window_seconds=1.0, burst_credits=2, clock=self.fake_clock)
        d = limiter.allow("zero-limit", cost=1)
        self.assertTrue(d.allowed)
        # Non-monotonic clock call
        d_past = limiter.allow("zero-limit", cost=1, now=self.current_time - 10.0)
        self.assertTrue(d_past.allowed)
        d_exhausted = limiter.allow("zero-limit", cost=1)
        self.assertFalse(d_exhausted.allowed)

    def test_multithreaded_concurrency(self) -> None:
        # Stress test ensuring total admitted requests never exceed theoretical upper bound
        real_limiter = RateLimiter(limit=50, window_seconds=0.2, burst_credits=10)
        threads: List[threading.Thread] = []
        admitted_counts: List[int] = [0] * 10

        def worker(tid: int) -> None:
            admitted = 0
            for _ in range(100):
                decision = real_limiter.allow("shared-key", cost=1)
                if decision.allowed:
                    admitted += 1
                time.sleep(0.001)
            admitted_counts[tid] = admitted

        for i in range(10):
            t = threading.Thread(target=worker, args=(i,))
            threads.append(t)
            t.start()
        for t in threads:
            t.join()

        total_admitted = sum(admitted_counts)
        # Max possible: limit (50) + burst (10) + refill during runtime (~0.1-0.3s -> ~75 requests)
        # Must never exceed conservative bound of limit + burst + rate * 1.5s
        self.assertLessEqual(total_admitted, 50 + 10 + int((50 / 0.2) * 1.5))
        self.assertGreater(total_admitted, 0)


if __name__ == "__main__":
    unittest.main()

継続的にリフィリングされるバーストクレジットを備えた、正確なスライディングウィンドウレートリミッター。

設計ノート（300語以内）：
各キーは、(タイムスタンプ、コスト)レコードのデックを所有します。許可された各リクエストは、(t - ウィンドウ, t] で正確にカウントされ、t - ウィンドウ のレコードは期限切れになります。同じ時刻のレコードはマージされます。これはバケット化ではなく正確なものです。時間誤差はゼロですが、メモリはO(アクティブウィンドウで許可されたリクエスト数)であり、クリーンアップは期限切れレコードあたり償却O(1)です。

通常の容量は max(0, limit - window_usage) です。残りのリクエストコストは、継続的にリフィリングされ、キャップされたバーストプールから支払われます。許可はアトミックであり、オールオアナッシングです。リトライ時間は、トークンリフィルと将来のレコード有効期限の境界から正確に計算されます。

キーは64個のシャードに分割されます。シャードロックがマッピングとLRU順序を保護します。各キーには、デックとクレジットを保護する個別のロックがあります。シャードロックは、キーの検索/作成およびロック中にのみ保持され、初回タッチまたはエビクションでの競合を、無関係なシャードをシリアル化することなく防ぎます。その後、キーに対する処理はそのキーのロックのみで実行されます。

アイドルエントリは、window_seconds に空のバーストプールをリフィルするのに必要な時間を加えた後に期限切れになります。機会的なLRU回収は、ワンショットキーワークロード下でバウンドされた保持を提供します。reap() は、決定的/完全なメンテナンスを許可します。limit == 0 の場合、リフィルはゼロで、アイドル状態は1つのウィンドウ後に忘れられるため、そのキャッシュ保持間隔後に戻るIDはリセットされて開始します。