Add TokenBucketRateLimiter strategy#364
Open
dimartiro wants to merge 3 commits into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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-onlyget()). This reintroduces the strategy on top of the current API and resolves both problems.What this adds
TokenBucketRateLimiter(synclimits.strategies+ asynclimits.aio.strategies), registered as thetoken-bucketstrategy. Each resource is a bucket holding up tolimittokens that refills continuously atlimit / expirytokens/second. Ahitconsumes 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.TokenBucketSupportcapability mixin (sync + async) withacquire_token_bucket(atomic refill-then-consume) andget_token_bucket(non-mutating read fortest/get_window_stats), mirroring the existingMovingWindowSupport/SlidingWindowCounterSupportpattern.Backends
State is kept in a dedicated per-backend structure so the int counter path is never touched:
(tokens, last_refill)map refilled under the existing per-key lock, with its own safety-expiry sweep.acquire_moving_window.lua).find_one_and_updateaggregation pipeline using$$NOWon a dedicated, TTL-indexed collection (same technique asacquire_sliding_window_entry).NotImplementedError(same as moving window).Tests & docs
test, multi-cost, clear, and time-based refill / cap (parametrised over every supporting backend).Notes
get_window_statsreportsremaining = floor(tokens)andreset = 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.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.