diff --git a/wizard-work/beam/sliding_window.py b/wizard-work/beam/sliding_window.py new file mode 100644 index 0000000..09ae228 --- /dev/null +++ b/wizard-work/beam/sliding_window.py @@ -0,0 +1,41 @@ +"""Sliding-window rate limiter. + +Public API: + SlidingWindow(limit, window) with .allow(key, now=None) -> bool +Counts events per key within a rolling window of `window` seconds. +Thread-safe? No (single-threaded). Deterministic given `now`. +""" +from __future__ import annotations + +from collections import defaultdict, deque +from typing import Deque, Dict, Optional + + +class SlidingWindow: + def __init__(self, limit: int, window: float): + if limit < 1: + raise ValueError("limit must be >= 1") + if window <= 0: + raise ValueError("window must be > 0") + self.limit = limit + self.window = window + self._events: Dict[object, Deque[float]] = defaultdict(deque) + + def allow(self, key, now: Optional[float] = None) -> bool: + if now is None: + import time + now = time.monotonic() + q = self._events[key] + cutoff = now - self.window + while q and q[0] <= cutoff: + q.popleft() + if len(q) < self.limit: + q.append(now) + return True + return False + + def reset(self, key=None): + if key is None: + self._events.clear() + else: + self._events.pop(key, None) diff --git a/wizard-work/beam/test_sliding_window.py b/wizard-work/beam/test_sliding_window.py new file mode 100644 index 0000000..2b3b936 --- /dev/null +++ b/wizard-work/beam/test_sliding_window.py @@ -0,0 +1,40 @@ +import sys +sys.path.insert(0, "/root/.hermes/wizard-outputs/beam") +from sliding_window import SlidingWindow + +def test_basic(): + w = SlidingWindow(2, 10.0) + assert w.allow("k", 0) is True + assert w.allow("k", 1) is True + assert w.allow("k", 2) is False + assert w.allow("k", 11) is True # first expired + +def test_window_boundary(): + w = SlidingWindow(1, 5.0) + assert w.allow("k", 0) is True + assert w.allow("k", 5) is True # exactly window: old event cutoff 0 <= 0 expired + +def test_per_key(): + w = SlidingWindow(1, 10.0) + assert w.allow("a", 0) is True + assert w.allow("b", 0) is True + assert w.allow("a", 1) is False + +def test_invalid(): + for kw in ({"limit":0,"window":1}, {"limit":1,"window":0}, {"limit":1,"window":-1}): + try: + SlidingWindow(**kw); raise SystemExit("expected ValueError") + except ValueError: + pass + +def test_reset(): + w = SlidingWindow(1, 10.0) + assert w.allow("k",0) is True + assert w.allow("k",1) is False + w.reset("k") + assert w.allow("k",1) is True + w.allow("a",0); w.allow("b",0) + w.reset() + assert w.allow("a",0) is True + +print("ALL TESTS PASSED")