Rate Limiter

Input: lines of the form epoch_seconds user endpoint, in time order. Lines that do not have three fields, or whose time is not an integer, are skipped and counted.

Stages

  1. Counting. Count requests per user and endpoint, and print the three endpoints with the most requests overall.
  2. Rate limit. A user may make at most N requests to an endpoint in any rolling 60-second window. Print each request with ALLOW or DENY. A request at time t leaves the window at t + 60. Only allowed requests use up the quota.
  3. Per-endpoint limits. Limits come from a config string endpoint:limit,endpoint:limit. Endpoints not listed use N.
  4. Blocking. A user denied 3 times within 10 seconds is blocked for 5 minutes, on every endpoint. Requests during the block are denied and do not extend it.

Files

  • rate_limiter.py: the solution for all four stages.
  • test_rate_limiter.py: tests grouped by stage. Run with python3 test_rate_limiter.py or pytest.

Solutions

rate_limiter.py

from collections import Counter, defaultdict, deque
from dataclasses import dataclass

WINDOW = 60  # seconds, rolling rate-limit window
BURST_WINDOW = 10  # seconds
BURST_DENIES = 3  # denies within BURST_WINDOW that trigger a block
BLOCK = 5 * 60  # seconds


@dataclass
class Request:
    ts: int
    user: str
    endpoint: str


def parse(line):
    """Input line: 'epoch_seconds user endpoint'. Returns None if malformed."""
    parts = line.split()
    if len(parts) != 3 or not parts[0].isdigit():
        return None
    ts, user, endpoint = parts
    return Request(int(ts), user, endpoint)


def parse_limits(limits, n):
    """'endpoint:limit,endpoint:limit' -> endpoint -> limit; unknown endpoints get n."""
    config = defaultdict(lambda: n)
    for part in limits.split(","):
        if not part:  # empty config string
            continue
        k, _, v = part.partition(":")
        config[k] = int(v)
    return config


class RateLimiter:
    def __init__(self, limits):
        self.limits = limits  # endpoint -> limit
        self.hits = defaultdict(deque)  # (user, endpoint) -> timestamps
        self.denies = defaultdict(deque)  # user -> timestamps of recent DENIED requests
        self.blocked_until = {}  # user -> ts when the block ends
        self.user_endpoints = defaultdict(Counter)  # user -> endpoint -> count
        self.endpoint_totals = Counter()  # endpoint -> count, all users

    def record(self, req):
        """count every request, allowed or not."""
        self.user_endpoints[req.user][req.endpoint] += 1
        self.endpoint_totals[req.endpoint] += 1

    def allow(self, req):
        """True = ALLOW, False = DENY."""
        # blocked users are denied first; these denies don't extend the block
        if self.blocked_until.get(req.user, -1) > req.ts:
            return False

        # rolling window per (user, endpoint), limit per endpoint
        q = self.hits[(req.user, req.endpoint)]
        while q and q[0] <= req.ts - WINDOW:
            q.popleft()
        if len(q) < self.limits[req.endpoint]:
            q.append(req.ts)  # only allowed requests use up the quota
            return True

        # a user denied BURST_DENIES times in BURST_WINDOW s is blocked
        r = self.denies[req.user]
        r.append(req.ts)
        while r and r[0] <= req.ts - BURST_WINDOW:
            r.popleft()
        if len(r) >= BURST_DENIES:
            self.blocked_until[req.user] = req.ts + BLOCK
            r.clear()
        return False

    def top_endpoints(self, k=3):
        return self.endpoint_totals.most_common(k)


def run(lines, n, limits=""):
    limiter, decisions, bad = RateLimiter(parse_limits(limits, n)), [], 0
    for line in lines:
        req = parse(line)
        if req is None:
            bad += 1
            continue
        limiter.record(req)
        d = "ALLOW" if limiter.allow(req) else "DENY"
        decisions.append(d)
        print(f"{req.ts} {req.user} {req.endpoint} {d}")
    for ep, c in limiter.top_endpoints():
        print(f"{ep}: {c}")
    print(f"malformed lines: {bad}")
    return limiter, decisions, bad

test_rate_limiter.py

from rate_limiter import parse, parse_limits, run

LINES = [
    "0 alice login",  # ALLOW
    "5 alice login",  # DENY: login limit is 1
    "6 alice search",  # ALLOW: search has its own window (default 2)
    "7 alice search",  # ALLOW
    "8 alice search",  # DENY: 3rd search in 60 s
    "9 bob login",  # ALLOW: bob has his own quota
    "60 alice login",  # ALLOW: ts 0 left the window (0 <= 60-60)
    "bad line",
    "x alice login",
]


def test_stage1_counting():
    l, _, bad = run(LINES, n=2, limits="login:1")
    assert bad == 2
    assert parse("x alice login") is None
    assert l.top_endpoints() == [("login", 4), ("search", 3)]
    assert l.user_endpoints["alice"] == {"login": 3, "search": 3}


def test_stage2_rolling_window():
    _, d, _ = run(["0 u a", "1 u a", "2 u a", "59 u a", "60 u a"], n=2)
    # third request denied; at 60 the request from 0 has left the window
    assert d == ["ALLOW", "ALLOW", "DENY", "DENY", "ALLOW"]


def test_stage3_per_endpoint_limits():
    _, d, _ = run(LINES, n=2, limits="login:1")
    assert d == ["ALLOW", "DENY", "ALLOW", "ALLOW", "DENY", "ALLOW", "ALLOW"]
    cfg = parse_limits("login:5,search:100", 7)
    assert cfg["login"] == 5 and cfg["search"] == 100 and cfg["other"] == 7
    assert parse_limits("", 3)["any"] == 3


def test_stage4_blocking():
    lines = [
        "0 u a",  # ALLOW
        "1 u a",  # DENY 1
        "2 u a",  # DENY 2
        "3 u a",  # DENY 3 -> blocked until 303
        "100 u b",  # DENY: blocked, even on another endpoint
        "302 u b",  # DENY: still blocked
        "303 u b",  # ALLOW: block over
    ]
    l, d, _ = run(lines, n=1)
    assert d == ["ALLOW", "DENY", "DENY", "DENY", "DENY", "DENY", "ALLOW"]
    assert l.blocked_until["u"] == 303
    l, _, _ = run(LINES, n=2, limits="login:1")
    assert "alice" not in l.blocked_until  # only 2 denies, no block


if __name__ == "__main__":
    for name, f in list(globals().items()):
        if name.startswith("test_"):
            f()
    print("all tests passed")