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

"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はリセットされて開始します。