beam: sliding-window rate limiter with injectable clock + tests
This commit is contained in:
1 parent
92745b2a66
commit
a46547c73e
3 files changed
+177
No files matched your search
+6
@@ -0,0 +1,6 @@
|
||||
#!/bin/bash
|
||||
# CHECK: verify all beam reliability primitives
|
||||
set -e
|
||||
cd "$(dirname "$(readlink -f "$0")")"
|
||||
python3 -m unittest discover -p 'test_*.py' -v
|
||||
python3 toposort/test_toposort.py
|
||||
@@ -0,0 +1,64 @@
|
||||
"""Sliding-window rate limiter with injectable clock and thread safety.
|
||||
|
||||
A request is admitted only if fewer than `limit` requests have been recorded
|
||||
in the trailing `window` seconds. Timestamps are kept in a deque; expired
|
||||
entries are pruned on every call. Concurrency is guarded by a lock.
|
||||
|
||||
Subtleties handled:
|
||||
- Clock injection (monotonic-ish) for deterministic tests.
|
||||
- Boundary: a request exactly `window` seconds old is expired (half-open interval).
|
||||
- limit <= 0 rejects everything; window <= 0 rejects everything.
|
||||
- Thread safety.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from typing import Callable, Optional
|
||||
|
||||
|
||||
class RateLimiter:
|
||||
def __init__(
|
||||
self,
|
||||
limit: int,
|
||||
window: float,
|
||||
clock: Optional[Callable[[], float]] = None,
|
||||
) -> None:
|
||||
if limit < 0:
|
||||
raise ValueError("limit must be >= 0")
|
||||
if window < 0:
|
||||
raise ValueError("window must be >= 0")
|
||||
self.limit = limit
|
||||
self.window = window
|
||||
self._clock = clock if clock is not None else time.monotonic
|
||||
self._timestamps: deque[float] = deque()
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def _now(self) -> float:
|
||||
return self._clock()
|
||||
|
||||
def allow(self) -> bool:
|
||||
"""Return True and record the request if admitted."""
|
||||
with self._lock:
|
||||
now = self._now()
|
||||
# prune expired: keep ts where now - ts < window (half-open)
|
||||
while self._timestamps and (now - self._timestamps[0]) >= self.window:
|
||||
self._timestamps.popleft()
|
||||
if self.limit <= 0 or self.window <= 0:
|
||||
return False
|
||||
if len(self._timestamps) < self.limit:
|
||||
self._timestamps.append(now)
|
||||
return True
|
||||
return False
|
||||
|
||||
def reset(self) -> None:
|
||||
with self._lock:
|
||||
self._timestamps.clear()
|
||||
|
||||
def count(self) -> int:
|
||||
with self._lock:
|
||||
now = self._now()
|
||||
while self._timestamps and (now - self._timestamps[0]) >= self.window:
|
||||
self._timestamps.popleft()
|
||||
return len(self._timestamps)
|
||||
@@ -0,0 +1,107 @@
|
||||
import threading
|
||||
import unittest
|
||||
|
||||
from rate_limiter import RateLimiter
|
||||
|
||||
|
||||
class TestRateLimiter(unittest.TestCase):
|
||||
def test_basic_admit_and_reject(self):
|
||||
t = [0.0]
|
||||
rl = RateLimiter(limit=2, window=10.0, clock=lambda: t[0])
|
||||
self.assertTrue(rl.allow())
|
||||
self.assertTrue(rl.allow())
|
||||
self.assertFalse(rl.allow())
|
||||
|
||||
def test_window_expiry(self):
|
||||
t = [0.0]
|
||||
rl = RateLimiter(limit=1, window=5.0, clock=lambda: t[0])
|
||||
self.assertTrue(rl.allow())
|
||||
t[0] = 4.999
|
||||
self.assertFalse(rl.allow())
|
||||
t[0] = 5.0
|
||||
# exactly window old -> expired
|
||||
self.assertTrue(rl.allow())
|
||||
|
||||
def test_boundary_half_open(self):
|
||||
t = [0.0]
|
||||
rl = RateLimiter(limit=2, window=1.0, clock=lambda: t[0])
|
||||
self.assertTrue(rl.allow()) # ts 0
|
||||
t[0] = 0.5
|
||||
self.assertTrue(rl.allow()) # ts 0.5
|
||||
t[0] = 1.0
|
||||
# ts 0 expired (1.0-0>=1), ts 0.5 still live (1.0-0.5=0.5<1)
|
||||
self.assertTrue(rl.allow()) # admits, now 2 live? ts0.5 + new
|
||||
# window now has 0.5 and 1.0 -> full
|
||||
self.assertFalse(rl.allow())
|
||||
|
||||
def test_limit_zero_rejects(self):
|
||||
rl = RateLimiter(limit=0, window=10.0, clock=lambda: 0.0)
|
||||
self.assertFalse(rl.allow())
|
||||
|
||||
def test_window_zero_rejects(self):
|
||||
rl = RateLimiter(limit=5, window=0.0, clock=lambda: 0.0)
|
||||
self.assertFalse(rl.allow())
|
||||
|
||||
def test_negative_args(self):
|
||||
with self.assertRaises(ValueError):
|
||||
RateLimiter(limit=-1, window=1.0)
|
||||
with self.assertRaises(ValueError):
|
||||
RateLimiter(limit=1, window=-1.0)
|
||||
|
||||
def test_reset(self):
|
||||
t = [0.0]
|
||||
rl = RateLimiter(limit=1, window=10.0, clock=lambda: t[0])
|
||||
self.assertTrue(rl.allow())
|
||||
self.assertFalse(rl.allow())
|
||||
rl.reset()
|
||||
self.assertTrue(rl.allow())
|
||||
|
||||
def test_count(self):
|
||||
t = [0.0]
|
||||
rl = RateLimiter(limit=5, window=10.0, clock=lambda: t[0])
|
||||
rl.allow(); rl.allow()
|
||||
self.assertEqual(rl.count(), 2)
|
||||
t[0] = 20.0
|
||||
self.assertEqual(rl.count(), 0)
|
||||
|
||||
def test_thread_safety(self):
|
||||
t = [0.0]
|
||||
rl = RateLimiter(limit=1000, window=10.0, clock=lambda: t[0])
|
||||
admitted = [0]
|
||||
lock = threading.Lock()
|
||||
|
||||
def worker():
|
||||
for _ in range(500):
|
||||
if rl.allow():
|
||||
with lock:
|
||||
admitted[0] += 1
|
||||
|
||||
threads = [threading.Thread(target=worker) for _ in range(4)]
|
||||
for th in threads:
|
||||
th.start()
|
||||
for th in threads:
|
||||
th.join()
|
||||
self.assertEqual(admitted[0], 1000)
|
||||
|
||||
def test_concurrent_never_exceeds_limit(self):
|
||||
t = [0.0]
|
||||
rl = RateLimiter(limit=50, window=10.0, clock=lambda: t[0])
|
||||
admitted = [0]
|
||||
lock = threading.Lock()
|
||||
|
||||
def worker():
|
||||
for _ in range(200):
|
||||
if rl.allow():
|
||||
with lock:
|
||||
admitted[0] += 1
|
||||
|
||||
threads = [threading.Thread(target=worker) for _ in range(8)]
|
||||
for th in threads:
|
||||
th.start()
|
||||
for th in threads:
|
||||
th.join()
|
||||
self.assertLessEqual(admitted[0], 50)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user