Skip to content

Add TokenBucketRateLimiter strategy#364

Open
dimartiro wants to merge 3 commits into
alisaifee:masterfrom
dimartiro:revive/token-bucket
Open

Add TokenBucketRateLimiter strategy#364
dimartiro wants to merge 3 commits into
alisaifee:masterfrom
dimartiro:revive/token-bucket

Conversation

@dimartiro

Copy link
Copy Markdown

This revives #234, which was closed because its implementation no longer worked against the current storage API (it relied on a removed storage.get_current_time() and on storing a (tokens, timestamp) tuple through the int-only get()). This reintroduces the strategy on top of the current API and resolves both problems.

What this adds

  • TokenBucketRateLimiter (sync limits.strategies + async limits.aio.strategies), registered as the token-bucket strategy. Each resource is a bucket holding up to limit tokens that refills continuously at limit / expiry tokens/second. A hit consumes tokens if enough are available; a rejected hit consumes nothing. The bucket starts full, so it allows bursts up to capacity while holding the long-run average to the configured rate.
  • TokenBucketSupport capability mixin (sync + async) with acquire_token_bucket (atomic refill-then-consume) and get_token_bucket (non-mutating read for test/get_window_stats), mirroring the existing MovingWindowSupport / SlidingWindowCounterSupport pattern.

Backends

State is kept in a dedicated per-backend structure so the int counter path is never touched:

  • in-memory — a (tokens, last_refill) map refilled under the existing per-key lock, with its own safety-expiry sweep.
  • redis (incl. cluster, sentinel, valkey) — an atomic Lua script over a hash. Tokens are stored as a string so the fraction is never lost to Lua integer coercion, and the clock is passed in as an argument (like acquire_moving_window.lua).
  • mongodb — a single find_one_and_update aggregation pipeline using $$NOW on a dedicated, TTL-indexed collection (same technique as acquire_sliding_window_entry).
  • memcached — unsupported; instantiating the limiter with it raises NotImplementedError (same as moving window).

Tests & docs

  • sync + async suites covering starts-full, cost-exceeds-capacity, non-consuming test, multi-cost, clear, and time-based refill / cap (parametrised over every supporting backend).
  • Documented in the README, strategies guide and API reference.

Notes

  • get_window_stats reports remaining = floor(tokens) and reset = now + (capacity - tokens) / rate (time until the bucket is full again — analogous to the fixed-window reset). Happy to adjust to time-until-next-token if preferred.
  • Capacity and refill rate are derived from the RateLimitItem (capacity = amount, rate = amount / expiry) so a token bucket is configured like every other strategy. A decoupled burst-vs-rate form would need a new parameter — can add if wanted.

Supersedes #234.

dimartiro added 3 commits July 2, 2026 11:01
Reintroduces the token bucket strategy on top of the current storage API,
resolving the two problems that stalled the original alisaifee#234: it no longer needs
a removed `storage.get_current_time()`, and it no longer tries to store a
`(tokens, timestamp)` tuple through the int-only `get()`.

Each resource is a bucket holding up to `limit` tokens that refills
continuously at `limit / expiry` tokens per second. A `hit` consumes tokens if
enough are available; a rejected hit consumes nothing. The bucket starts full,
so it permits bursts up to capacity while holding the long-run average to the
configured rate.

- Add a `TokenBucketSupport` capability mixin (sync + async) with
  `acquire_token_bucket` (atomic refill-then-consume) and `get_token_bucket`
  (non-mutating read for `test`/`get_window_stats`), mirroring the existing
  moving/sliding window support classes.
- Implement `TokenBucketRateLimiter` for `limits.strategies` and
  `limits.aio.strategies`, registered as the `token-bucket` strategy.
- Back it with dedicated per-backend state so the int counter path is never
  touched:
  - in-memory: a `(tokens, last_refill)` map refilled under the existing key
    lock, with its own safety-expiry sweep.
  - redis (incl. cluster, sentinel, valkey): an atomic Lua script storing a
    hash; tokens are written as a string so the fraction is never truncated,
    and the clock is passed in as an argument.
  - mongodb: a single `find_one_and_update` aggregation pipeline using `$$NOW`
    on a dedicated, TTL-indexed collection.
  - memcached is unsupported and raises `NotImplementedError` (like moving
    window).
- Add sync + async test suites covering starts-full, cost-exceeds-capacity,
  non-consuming `test`, multi-cost, clear, and time-based refill/cap.
- Document the strategy in the README, strategies guide and API reference.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant