Skip to content

← Testing with pytest and Hypothesis step 7 of 19

Medium Primitives

Build a Rate Limiter That Reads a Fake Clock

A rate limiter stops a client from sending too many requests too fast. A token bucket is a common design. The bucket holds up to capacity tokens and starts full. Each request spends tokens, and the bucket gains refill_per_second tokens every second until it is full again.

In Pass In a Fake Clock, the clock was a plain function. Here it is an object with a monotonic() method, described by the Protocol named Clock. The driver's FakeClock fits Clock without inheriting from it.

clock = FakeClock()
allow = make_rate_limiter(capacity=2, refill_per_second=0.5, clock=clock)
allow(1)  # True
allow(1)  # True, and the bucket is empty
clock.now = 1.0
allow(1)  # False, only half a token has come back
clock.now = 2.0
allow(1)  # True, the half token was kept, so now there is one

What to write

def make_rate_limiter(
    *, capacity: int, refill_per_second: float, clock: Clock
) -> Callable[[int], bool]

Return a function allow(n). If the bucket holds at least n tokens, it spends n and returns True. Otherwise it spends nothing and returns False.

  • When you create the limiter, start with float(capacity) tokens. Read clock.monotonic() once and store the reading as last. The clock can start at any number.
  • Each call to allow(n) reads clock.monotonic() exactly once.
  • If that reading, now, is greater than last, add (now - last) * refill_per_second tokens, capped at capacity. Then set last to now.
  • Keep fractional tokens as they are, without rounding.
  • Refill first, then decide. A denied request keeps the tokens the refill added.
  • Read the time only from clock, and never sleep.

When the clock goes backward

The real time.monotonic() only goes forward, but the fake clock in the tests can go back. If a reading is at or below last, add no tokens and leave last alone. So last is always the highest reading so far. After 10, a reading of 3 adds nothing, and a later 11 adds one second of tokens. While the clock is behind, requests can still spend the tokens already in the bucket.

Edge cases

  • allow(0) always returns True.
  • A request for more than capacity tokens always returns False.
  • With capacity=0, only allow(0) succeeds.
  • With refill_per_second=0, spent tokens never come back.
  • Each limiter has its own tokens and its own last, even when several share one clock. A limiter refills only during its own calls.

capacity, n and the rate are zero or more, and clock readings are finite. The driver calls one limiter at a time, so you need no lock.

What the driver reports

Keep solve as it is. It creates one FakeClock at 0.0 and buckets limiters that share it. It handles the requests in the order given, not sorted by time. For each one, it sets the clock to at and calls limiter number bucket, counting from 0, with tokens. It returns the list of results, which is [] when there are no requests.