0

Pluggable hierarchical rate limiter

by · about 3 hours ago· ⑂ 0 remixes

Question: Pluggable hierarchical rate limiter

Companies: Cursor

Difficulty: Medium

Problem statement

Design a rate limiter with hierarchical limits (user, team, global) whose limit rules are pluggable without changing the core. Reported on 1point3acres' Cursor bank; reconstructed in our own words.

Clarify the setup

This is an object-oriented design question with a working core. Ask:

  • What are the levels? For an AI coding tool, typically per user, per team (organization) and global (protect the model backends). A request names one key per level: {"user": "u1", "team": "acme", "global": "*"}.
  • What does "pluggable" mean? New limit algorithms (token bucket, sliding window, fixed window, "max N in flight") and new levels can be added by configuration or by registering a class, without editing the limiter.
  • What does a request cost? Usually 1, but for model calls it is often tokens. Support a cost.
  • What do we return? Allowed or not, and if not, which rule blocked and how long to wait (retry_after) so the client can back off.
  • Single process or distributed? Start in one process with an injected clock, then discuss Redis.

Approach

The core rule: a denied request consumes nothing. If the user bucket has room but the team bucket is empty, the request is denied, and the user must not lose a token. So acquiring is two-phase:

  1. Check every rule at every level: "how long until you could take cost?" (0 means now).
  2. If every answer is 0, take from every rule. Otherwise take nothing and return the largest wait.

In one process with one lock (or one thread), check-then-take is atomic. Distributed, the same two phases become one Lua script per request in Redis, or a reservation that is rolled back on failure.

The plug-in surface is a small interface:

wait_time(key, now, cost) -> seconds   # 0.0 when it fits now
take(key, now, cost) -> None

Each rule owns its per-key state. A registry maps a type name to a class, so a config like {"team": [{"type": "token_bucket", "capacity": 600, "per_second": 10}]} builds the limiter. Adding an algorithm is one class and one decorator. Levels are just names in the config, so adding "per workspace" changes no code.

Algorithms worth knowing

Rule State per key Behavior
Token bucket tokens, last refill time Allows bursts up to capacity, steady rate after. Best default for APIs.
Sliding window log timestamps in the window Exact "N per window", memory grows with N.
Fixed window count, window start Cheapest, but allows 2N at a window boundary.

Code

import math
from collections import deque
from dataclasses import dataclass
from typing import Callable, Deque, Dict, List, Optional, Tuple

RULE_TYPES: Dict[str, Callable[..., "Rule"]] = {}


def register(name: str):
    def deco(cls):
        RULE_TYPES[name] = cls
        return cls
    return deco


class Rule:
    """Plug-in interface. Subclasses keep their own per-key state."""

    name = "rule"

    def wait_time(self, key: str, now: float, cost: float) -> float:
        raise NotImplementedError

    def take(self, key: str, now: float, cost: float) -> None:
        raise NotImplementedError


@register("token_bucket")
class TokenBucket(Rule):
    def __init__(self, capacity: float, per_second: float) -> None:
        self.capacity = capacity
        self.rate = per_second
        self._state: Dict[str, Tuple[float, float]] = {}  # key -> (tokens, last)
        self.name = f"token_bucket({capacity}, {per_second}/s)"

    def _tokens(self, key: str, now: float) -> float:
        tokens, last = self._state.get(key, (self.capacity, now))
        return min(self.capacity, tokens + (now - last) * self.rate)

    def wait_time(self, key: str, now: float, cost: float) -> float:
        if cost > self.capacity:
            return math.inf
        missing = cost - self._tokens(key, now)
        return 0.0 if missing <= 0 else missing / self.rate

    def take(self, key: str, now: float, cost: float) -> None:
        self._state[key] = (self._tokens(key, now) - cost, now)


@register("sliding_window")
class SlidingWindowLog(Rule):
    def __init__(self, limit: float, window: float) -> None:
        self.limit = limit
        self.window = window
        self._log: Dict[str, Deque[Tuple[float, float]]] = {}
        self._used: Dict[str, float] = {}
        self.name = f"sliding_window({limit}/{window}s)"

    def _prune(self, key: str, now: float) -> Deque[Tuple[float, float]]:
        log = self._log.setdefault(key, deque())
        while log and log[0][0] <= now - self.window:
            self._used[key] -= log.popleft()[1]
        return log

    def wait_time(self, key: str, now: float, cost: float) -> float:
        if cost > self.limit:
            return math.inf
        log = self._prune(key, now)
        excess = self._used.get(key, 0.0) + cost - self.limit
        if excess <= 0:
            return 0.0
        freed = 0.0
        for ts, c in log:  # oldest first: wait until enough of them expire
            freed += c
            if freed >= excess:
                return ts + self.window - now
        return math.inf  # unreachable when cost <= limit

    def take(self, key: str, now: float, cost: float) -> None:
        self._prune(key, now).append((now, cost))
        self._used[key] = self._used.get(key, 0.0) + cost


@register("fixed_window")
class FixedWindow(Rule):
    def __init__(self, limit: float, window: float) -> None:
        self.limit = limit
        self.window = window
        self._state: Dict[str, Tuple[int, float]] = {}
        self.name = f"fixed_window({limit}/{window}s)"

    def _count(self, key: str, now: float) -> Tuple[int, float]:
        bucket = int(now // self.window)
        b, used = self._state.get(key, (bucket, 0.0))
        return (bucket, 0.0) if b != bucket else (b, used)

    def wait_time(self, key: str, now: float, cost: float) -> float:
        if cost > self.limit:
            return math.inf
        bucket, used = self._count(key, now)
        return 0.0 if used + cost <= self.limit else (bucket + 1) * self.window - now

    def take(self, key: str, now: float, cost: float) -> None:
        bucket, used = self._count(key, now)
        self._state[key] = (bucket, used + cost)


@dataclass
class Decision:
    allowed: bool
    retry_after: float = 0.0
    blocked_by: Optional[str] = None  # "level: rule"


class HierarchicalLimiter:
    def __init__(self, levels: Dict[str, List[Rule]]) -> None:
        self.levels = levels  # e.g. {"user": [...], "team": [...], "global": [...]}

    @classmethod
    def from_config(cls, config: Dict[str, List[dict]]) -> "HierarchicalLimiter":
        levels = {}
        for level, specs in config.items():
            levels[level] = [RULE_TYPES[s["type"]](**{k: v for k, v in s.items() if k != "type"}) for s in specs]
        return cls(levels)

    def acquire(self, keys: Dict[str, str], now: float, cost: float = 1.0) -> Decision:
        checks = [(level, rule, keys[level]) for level, rules in self.levels.items()
                  if level in keys for rule in rules]
        worst, blocker = 0.0, None
        for level, rule, key in checks:                 # phase 1: check
            wait = rule.wait_time(key, now, cost)
            if wait > worst:
                worst, blocker = wait, f"{level}: {rule.name}"
        if worst > 0:
            return Decision(False, worst, blocker)
        for level, rule, key in checks:                 # phase 2: take
            rule.take(key, now, cost)
        return Decision(True)

Worked numbers

A team plan allows 600 requests a minute with bursts of 100: token_bucket(capacity=100, per_second=10). After a burst of 100 at t = 0, the bucket is empty; request 101 at t = 0 gets retry_after = 1 / 10 = 0.1 s. At t = 5 the bucket holds 50.

Per-key state is tiny: a token bucket is two floats (16 bytes plus the key). One million users is about 100 MB in Python and about 50 MB in Redis hashes. A sliding-window log stores every request in the window: at 600 per minute per key that is 600 entries per key, which is why it is used for low limits or replaced by a sliding-window counter (two fixed windows, weighted).

Test cases

def _limiter():
    return HierarchicalLimiter({
        "user": [TokenBucket(capacity=3, per_second=1)],
        "team": [TokenBucket(capacity=5, per_second=1)],
    })


def test_user_limit_and_refill():
    lim = _limiter()
    keys = {"user": "u1", "team": "t1"}
    assert [lim.acquire(keys, 0).allowed for _ in range(4)] == [True, True, True, False]
    d = lim.acquire(keys, 0)
    assert d.retry_after == 1.0 and d.blocked_by.startswith("user")
    assert lim.acquire(keys, 1.0).allowed


def test_team_limit_is_shared_and_denial_consumes_nothing():
    lim = _limiter()
    for _ in range(3):
        assert lim.acquire({"user": "a", "team": "t"}, 0).allowed
    for _ in range(2):
        assert lim.acquire({"user": "b", "team": "t"}, 0).allowed
    d = lim.acquire({"user": "b", "team": "t"}, 0)
    assert not d.allowed and d.blocked_by.startswith("team")
    # b's own bucket still has 1 token: the denied request did not take it
    assert lim.levels["user"][0].wait_time("b", 0, 1) == 0.0


def test_cost_larger_than_capacity_never_fits():
    lim = _limiter()
    d = lim.acquire({"user": "u", "team": "t"}, 0, cost=10)
    assert not d.allowed and d.retry_after == float("inf")
def test_config_and_sliding_window():
    lim = HierarchicalLimiter.from_config({
        "user": [{"type": "sliding_window", "limit": 2, "window": 10}],
        "global": [{"type": "fixed_window", "limit": 100, "window": 60}],
    })
    keys = {"user": "u", "global": "*"}
    assert lim.acquire(keys, 0).allowed
    assert lim.acquire(keys, 4).allowed
    d = lim.acquire(keys, 5)
    assert not d.allowed and d.retry_after == 5  # the t=0 entry expires at t=10
    assert lim.acquire(keys, 10).allowed


def test_fixed_window_boundary():
    rule = FixedWindow(limit=2, window=60)
    rule.take("k", 59, 1)
    rule.take("k", 59, 1)
    assert rule.wait_time("k", 59, 1) == 1
    assert rule.wait_time("k", 60, 1) == 0.0  # new window: the known 2x burst


def test_plugging_in_a_new_rule():
    @register("deny_list")
    class DenyList(Rule):
        def __init__(self, keys):
            self.keys = set(keys)
            self.name = "deny_list"

        def wait_time(self, key, now, cost):
            return float("inf") if key in self.keys else 0.0

        def take(self, key, now, cost):
            pass

    lim = HierarchicalLimiter.from_config({
        "user": [{"type": "deny_list", "keys": ["mallory"]}],
        "team": [{"type": "token_bucket", "capacity": 1, "per_second": 1}],
    })
    assert not lim.acquire({"user": "mallory", "team": "t"}, 0).allowed
    assert lim.acquire({"user": "alice", "team": "t"}, 0).allowed  # team token untouched

Follow-ups

  1. Distributed. Many API servers share limits. Keep each rule's state in Redis keyed by level:rule:key and run check-then-take as one Lua script, so it is atomic without locks. Shard by key. For the global level, a single hot key becomes a bottleneck; split it into N sub-buckets with capacity / N each, or let each server lease a slice of the global budget every second.
  2. Fail open or closed? If Redis is down, per-user limits usually fail open (allow, with a local in-memory fallback) while the global limit protecting GPUs fails closed or falls back to a conservative local limit.
  3. Token-based costs. For model calls the cost is unknown until the response ends. Reserve an estimate (prompt tokens plus max_tokens) up front, then refund or charge the difference when the call finishes.
  4. Priorities. Paid tiers or interactive requests get a reserved share: two buckets at the global level, with the low-priority bucket allowed only while the shared bucket is above a watermark.
  5. Concurrency limits. "At most 5 in-flight requests per user" is a rule whose take increments and whose release (a new hook) decrements. The interface grows a release method with a default no-op.

Interviewer expectations

Mid-level (L4)

  • A correct token bucket with lazy refill and an injected clock.
  • Checks all levels and denies if any level is out, with a clear Decision result.
  • Tests per-user and shared team limits.

Senior (L5)

  • Two-phase check-then-take so a denied request consumes nothing, and retry_after as the largest wait.
  • A real plug-in surface (interface plus registry plus config) and at least two algorithms with their trade-offs (burst behavior, memory, the fixed-window boundary).
  • Handles a cost larger than capacity and other edge cases without being prompted.

Staff (L6+)

  • Takes it distributed: atomic Lua scripts, sharding, hot global keys, leasing budgets to servers.
  • Chooses fail-open versus fail-closed per level and explains why.
  • Brings in model-specific concerns: token costs known only at the end, reservations and refunds, priority tiers protecting GPU capacity.

Sign in to comment, vote or remix.