diff --git a/README.md b/README.md index a4293a5..d4421c9 100644 --- a/README.md +++ b/README.md @@ -294,6 +294,47 @@ That is the whole seam for "shared budget when the datastore is there, in-process when it is not" — and the `Storage` stays the ordinary in-memory one in both branches, because it only caches handles. +## Two strategies: token bucket and leaky bucket + +**Given the same parameters they behave identically** — they are duals, and any +comparison implying otherwise has quietly changed the burst between the two +sides. Measured with this package: + +```text +token bucket, 10/s burst 1 → 1 admitted now, 1 more after 250ms +leaky bucket, 100ms cap 1 → 1 admitted now, 1 more after 250ms +``` + +What differs is **how you say it**, and that is worth something: + +| | Token bucket | Leaky bucket | +| --- | --- | --- | +| Configured by | rate + burst | interval + capacity | +| Natural default | absorbs a burst | paces | +| Arithmetic | float tokens | exact integer durations (GCRA) | + +*"No more than 1000 an hour"* is a budget → token bucket. *"No more than one +call every 100ms"* is a pace → leaky bucket, where strict pacing is the obvious +configuration rather than a non-obvious `burst: 1`. + +Pick one by value, so it can come from a config file or a database column: + +```go +strategy, err := ratelimiter.ParseStrategy("leaky_bucket") // or "token_bucket" +limit := ratelimiter.Limit{Requests: 60, Period: time.Minute, Burst: 1} + +newLimiter, err := ratelimiter.NewLimiterFunc(strategy, limit) +bl := ratelimiter.NewBucketLimiter(newLimiter, time.Minute, storage) +``` + +One `Limit` describes both; the strategy decides how it is enforced. +`ParseStrategy` rejects an unrecognised value rather than defaulting — a typo +that silently became `token_bucket` would admit bursts you specifically asked it +not to. + +**[docs/TOKEN_BUCKET.md](docs/TOKEN_BUCKET.md)** and +**[docs/LEAKY_BUCKET.md](docs/LEAKY_BUCKET.md)** are the full guides. + ## Backends — one limit shared across processes `Storage` holds limiters **in this process**. `Backend` holds the **count**, diff --git a/doc.go b/doc.go index 97a923a..3dc8ebd 100644 --- a/doc.go +++ b/doc.go @@ -22,6 +22,16 @@ // // allowed // } // +// # Two strategies +// +// [StrategyTokenBucket] admits a burst and refills continuously; +// [StrategyLeakyBucket] enforces a minimum spacing between admissions. Both can +// be configured for "60 a minute" and they behave completely differently — 60 +// at once versus one per second — so the choice matters more than the numbers. +// [NewLimiterFunc] selects one from a [Strategy] value, which is meant to +// survive a round trip through configuration; [ParseStrategy] validates one +// coming back. See docs/TOKEN_BUCKET.md and docs/LEAKY_BUCKET.md. +// // Limiters are consumed through the [Limiter] interface (Allow, Wait, Burst). // A limiter may optionally also implement [Reserver] to reserve a token and // report the exact delay until it is valid; the default limiter from diff --git a/docs/CUSTOM_STORAGE.md b/docs/CUSTOM_STORAGE.md index 9ee915a..94db9be 100644 --- a/docs/CUSTOM_STORAGE.md +++ b/docs/CUSTOM_STORAGE.md @@ -196,7 +196,7 @@ The rule of thumb: | You want… | Extension point | |------------------------------------------------------|----------------------------| | A different **in-process** store (LRU, metrics, …) | implement **`Storage`** | -| A different **algorithm** (leaky bucket, GCRA, …) | implement **`Limiter`** | +| A different **algorithm** | implement **`Limiter`** — though token bucket and leaky bucket are both [bundled](LEAKY_BUCKET.md) | | A **global** limit shared across instances | implement **[`Backend`](BACKENDS.md)** — `BackendLimiter` then supplies the `Limiter`, plus a local fallback, a circuit breaker and a degraded signal | Distributed limiting is a **`Limiter`** concern, not a `Storage` concern. diff --git a/docs/LEAKY_BUCKET.md b/docs/LEAKY_BUCKET.md new file mode 100644 index 0000000..be80377 --- /dev/null +++ b/docs/LEAKY_BUCKET.md @@ -0,0 +1,221 @@ +# Leaky bucket + +A leaky bucket configured as a *meter* and a token bucket are **the same +policy**. That is not a caveat buried at the bottom of this page — it is the +first thing to understand, because the internet is full of comparisons that +imply otherwise by quietly changing the burst setting between the two examples. + +Measured, with this package: + +```text +token bucket, 10/s burst 1 → 1 admitted immediately, 1 more after 250ms +leaky bucket, 100ms cap 1 → 1 admitted immediately, 1 more after 250ms +``` + +Identical. They are duals: a token bucket that refills at rate `r` with capacity +`b` admits exactly what a leaky bucket draining at `r` with capacity `b` admits. +If you see a comparison claiming one bursts and the other does not, check +whether the two sides were given the same capacity. + +## So why is it here? + +Two reasons, both real and neither of them "different behaviour". + +**1. It is parameterised by spacing, not by rate-and-burst.** You configure it +with an *interval* — "one every 100ms" — which makes strict pacing the obvious +configuration. With a token bucket, the same thing is spelled `burst: 1`, which +is easy to leave at a default and get a burst you did not intend. The +parameterisation is the feature: it makes the safe configuration the natural one +to write. + +**2. The implementation is exact.** GCRA keeps a single `time.Time` — the +theoretical arrival time — and does integer-duration arithmetic. There are no +floating-point tokens to accumulate rounding error over a long uptime, no +allocation, and every operation is O(1). + +If you already have a token bucket with the burst you want, **you do not need to +switch.** Reach for the leaky bucket when you are expressing a *pacing* +requirement — "no more than one call every N" — and want the code to say that. + +## Table of contents + +- [When to choose which](#when-to-choose-which) +- [How it works: virtual scheduling](#how-it-works-virtual-scheduling) +- [Capacity](#capacity) +- [Shaping versus dropping](#shaping-versus-dropping) +- [Cancel actually works here](#cancel-actually-works-here) +- [Usage](#usage) +- [Choosing between them at run time](#choosing-between-them-at-run-time) +- [Further reading](#further-reading) + +## When to choose which + +| | Token bucket | Leaky bucket | +| --- | --- | --- | +| Configured by | rate + burst | interval + capacity | +| Natural default | absorbs a burst | paces | +| Same parameters | **identical behaviour** | **identical behaviour** | +| Arithmetic | float tokens | exact integer durations | +| Shaping (`Wait`) | yes | yes | +| `Cancel()` returns the slot | yes | yes | + +The rule of thumb is about **what you are expressing**, not what you get: + +- *"No more than 1000 an hour"* is a budget → token bucket. +- *"No more than one call every 100ms"* is a pace → leaky bucket. + +Both can be made to do the other's job. Saying it the way you mean it is the +point. + +## How it works: virtual scheduling + +There is no queue and no timer — the "bucket" is one instant. + +The limiter stores a **theoretical arrival time** (TAT): when the next +conforming request may be admitted. Admitting a request pushes the TAT forward +by one emission interval. A request arriving more than the burst allowance ahead +of the TAT is refused, and **the amount it is early by *is* the retry delay**. + +```text +interval = period / requests e.g. 60/min → 1s +tolerance = capacity × interval + +on arrival at t: + tat = max(stored_tat, t) + next = tat + interval + allowAt = next − tolerance + + if t < allowAt: refuse, retry after (allowAt − t) + else: stored_tat = next, admit +``` + +This is GCRA — the Generic Cell Rate Algorithm, from ATM traffic policing. +Every operation is O(1), allocation-free and exact: nothing accumulates, so +there is no floating-point drift to worry about over long uptimes. + +## Capacity + +`capacity` is how many requests may be taken back-to-back from an **idle** +bucket. + +- **`1` is a strict leaky bucket**: perfectly even spacing, no burst at all. + This is what you want in front of a rate-limited third party. +- **`n`** admits `n` immediately and then paces at the interval — a token + bucket's shape with a leaky bucket's recovery. Note the difference from a + token bucket: after the burst, capacity comes back **one interval at a time**, + not as a refilling pool. + +A capacity below one is raised to one. A bucket that admits nothing is a +configuration error, and refusing every request for ever is the least helpful +way to report it. + +## Shaping versus dropping + +This is the choice that matters at the call site: + +```go +lb.Wait(ctx) // SHAPES: sleeps until the request conforms +lb.Allow() // DROPS: returns false immediately +``` + +`Wait` is what makes a leaky bucket a *shaper* — traffic is smoothed rather than +rejected, which is usually what you want for a background worker draining a +queue against a rate-limited API. + +`Allow` is for the cases where waiting is worse than failing: an HTTP handler +would generally rather answer `429` than hold a connection open. + +`Wait` rolls its reservation back if the context ends first, so a caller who +gives up does not make the next caller wait for a request that never happened. + +## Cancel actually works here + +Worth knowing if you are choosing a backend: + +- A **window counter** cannot hand a slot back in a way that is distinguishable + from spending one less, so its `Cancel()` is a documented no-op and a rejected + request still counts. +- A **leaky bucket** *can*: the TAT is a single instant, so cancelling is + rewinding it by one interval. A rejected request need not count against the + caller. + +```go +res := lb.Reserve() +if !res.OK() { + w.Header().Set("Retry-After", strconv.Itoa(int(res.Delay().Seconds()))) + http.Error(w, "slow down", http.StatusTooManyRequests) + res.Cancel() // the slot goes back + return +} +``` + +`Cancel` is idempotent — a `defer` plus an explicit call cannot double-refund. + +## Usage + +```go +// One request per second, strictly evenly spaced. +lb := ratelimiter.NewLeakyBucket(time.Second, 1) + +if lb.Allow() { + // conforms +} +``` + +Per key, through the manager: + +```go +// Every client IP gets one request per second, no bursting. +bl := ratelimiter.NewBucketLimiter( + ratelimiter.NewLeakyBucketFunc(time.Second, 1), + time.Minute, + ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), +) +defer bl.Close() + +if bl.GetOrAdd(clientIP).Allow() { + // ... +} +``` + +## Choosing between them at run time + +The strategy is a value, so it can come from a database column, a YAML field or +an environment variable: + +```go +strategy, err := ratelimiter.ParseStrategy(row.Strategy) // "token_bucket" | "leaky_bucket" +if err != nil { + return err +} + +limit := ratelimiter.Limit{Requests: row.Requests, Period: row.Period, Burst: row.Burst} + +newLimiter, err := ratelimiter.NewLimiterFunc(strategy, limit) +if err != nil { + return err +} + +bl := ratelimiter.NewBucketLimiter(newLimiter, time.Minute, storage) +``` + +**One `Limit` describes both.** The same "N requests per period" an operator +writes down; the strategy decides how it is enforced. The burst means slightly +different things — bucket size for one, back-to-back allowance for the other — +which is stated here rather than left to be discovered. + +`ParseStrategy` **rejects** an unrecognised value rather than defaulting. A typo +that silently became `token_bucket` would admit bursts the operator specifically +asked it not to, and nothing about the running service would say so. + +`ratelimiter.Strategies()` returns every valid value, for populating a form or +validating a schema. + +## Further reading + +- [TOKEN_BUCKET.md](./TOKEN_BUCKET.md) — the default algorithm, and the + comparison table this page is the other half of +- [BACKENDS.md](./BACKENDS.md) — sharing one limit across processes +- [CUSTOM_STORAGE.md](./CUSTOM_STORAGE.md) — in-process containers +- GCRA / leaky bucket — +- Generic Cell Rate Algorithm — diff --git a/docs/TOKEN_BUCKET.md b/docs/TOKEN_BUCKET.md index 802020e..b61a1cb 100644 --- a/docs/TOKEN_BUCKET.md +++ b/docs/TOKEN_BUCKET.md @@ -297,7 +297,7 @@ sequenceDiagram | Algorithm | Bursts | Memory/key | Notes | |----------------------|--------|------------|-------------------------------------------------| | **Token bucket** | Yes | O(1) | This library. Smooth average + bounded burst. | -| Leaky bucket | No | O(1) | Enforces a strictly constant output rate. | +| Leaky bucket | No | O(1) | Enforces a strictly constant output rate. **Bundled** — see [LEAKY_BUCKET.md](LEAKY_BUCKET.md). | | Fixed window counter | Spiky | O(1) | Simple, but allows 2× burst at window edges. | | Sliding window log | Exact | O(requests)| Precise but stores every timestamp. | | Sliding window count | Good | O(1) | Approximates the log cheaply; common in Redis. | @@ -350,6 +350,7 @@ use `RateLimiter` (the token bucket) rather than `MemoryBackend`. - `golang.org/x/time/rate` package docs — - Token bucket — +- [LEAKY_BUCKET.md](LEAKY_BUCKET.md) — the bundled leaky bucket, and when to prefer it - Leaky bucket — - IETF RateLimit header fields — - RFC 9110 (HTTP semantics, `Retry-After`, `429`) — diff --git a/examples/leakybucket/main.go b/examples/leakybucket/main.go new file mode 100644 index 0000000..7feb454 --- /dev/null +++ b/examples/leakybucket/main.go @@ -0,0 +1,109 @@ +// Command leakybucket shows what a leaky bucket is for: pacing calls to +// something that meters you, and shaping rather than dropping. +// +// It also demonstrates the thing most comparisons get wrong — with the same +// parameters, a leaky bucket and a token bucket behave identically. +// +// go run ./examples/leakybucket +package main + +import ( + "context" + "fmt" + "time" + + "golang.org/x/time/rate" + + "github.com/slashdevops/ratelimiter" +) + +func drain(l ratelimiter.Limiter, n int) int { + admitted := 0 + + for range n { + if l.Allow() { + admitted++ + } + } + + return admitted +} + +func main() { + fmt.Println("1. Same parameters, same behaviour") + fmt.Println(" The comparison you usually see is rigged by changing the burst.") + fmt.Println() + + tb := ratelimiter.RateLimiter{Limiter: rate.NewLimiter(rate.Limit(10), 1)} + lb := ratelimiter.NewLeakyBucket(100*time.Millisecond, 1) + + fmt.Printf(" token bucket 10/s burst 1 : %d of 10 admitted now\n", drain(tb, 10)) + fmt.Printf(" leaky bucket 100ms cap 1 : %d of 10 admitted now\n", drain(lb, 10)) + + time.Sleep(250 * time.Millisecond) + + fmt.Printf(" token bucket, 250ms later : %d of 10\n", drain(tb, 10)) + fmt.Printf(" leaky bucket, 250ms later : %d of 10\n", drain(lb, 10)) + fmt.Println("\n Identical. They are duals.") + + // --------------------------------------------------------------------- + fmt.Println("\n2. What it is actually for: shaping") + fmt.Println(" Wait sleeps until the request conforms instead of dropping it,") + fmt.Println(" which is what you want draining a queue against a metered API.") + fmt.Println() + + paced := ratelimiter.NewLeakyBucket(80*time.Millisecond, 1) + start := time.Now() + + for i := range 5 { + if err := paced.Wait(context.Background()); err != nil { + panic(err) + } + + fmt.Printf(" call %d at %v\n", i+1, time.Since(start).Round(10*time.Millisecond)) + } + + // --------------------------------------------------------------------- + fmt.Println("\n3. Retry-After is exact, and Cancel gives the slot back") + fmt.Println() + + bucket := ratelimiter.NewLeakyBucket(time.Second, 1) + + // Reserve TAKES the slot when it conforms. That is what makes Cancel + // meaningful — there is something to give back. + res := bucket.Reserve() + fmt.Printf(" reserved: ok=%v\n", res.OK()) + fmt.Printf(" bucket now exhausted: allow=%v\n", bucket.Allow()) + + res.Cancel() + fmt.Printf(" after Cancel, admitted again: %v\n", bucket.Allow()) + + // A reservation that was REFUSED never took a slot, so cancelling it is a + // no-op rather than free budget. + refused := bucket.Reserve() + fmt.Printf("\n refused reservation: ok=%v, retry after %v\n", + refused.OK(), refused.Delay().Round(10*time.Millisecond)) + + refused.Cancel() + fmt.Printf(" cancelling a refused reservation grants nothing: allow=%v\n", bucket.Allow()) + + fmt.Println("\n A window-counter backend cannot do any of this — it has no way") + fmt.Println(" to hand a slot back that differs from spending one less.") + + // --------------------------------------------------------------------- + fmt.Println("\n4. Per key, through the manager") + fmt.Println() + + bl := ratelimiter.NewBucketLimiter( + ratelimiter.NewLeakyBucketFunc(time.Second, 1), + time.Minute, + ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), + ) + defer bl.Close() + + for _, key := range []string{"alice", "alice", "bob"} { + fmt.Printf(" %-6s allow=%v\n", key, bl.GetOrAdd(key).Allow()) + } + + fmt.Println("\n alice is paced; bob is unaffected. Independent buckets per key.") +} diff --git a/examples/strategy/main.go b/examples/strategy/main.go new file mode 100644 index 0000000..eabb386 --- /dev/null +++ b/examples/strategy/main.go @@ -0,0 +1,100 @@ +// Command strategy shows what the strategy actually changes. +// +// The difference is NOT that one bursts and the other never does — give a leaky +// bucket a capacity and it bursts too. The difference is how the capacity comes +// BACK: a token bucket refills continuously, a leaky bucket releases one slot +// per interval. And a leaky bucket at capacity 1 paces strictly, which a token +// bucket cannot do at all. +// +// go run ./examples/strategy +package main + +import ( + "context" + "fmt" + "time" + + "github.com/slashdevops/ratelimiter" +) + +func drain(l ratelimiter.Limiter, n int) int { + admitted := 0 + + for range n { + if l.Allow() { + admitted++ + } + } + + return admitted +} + +func main() { + // Identical Limit for both: 10 per second, burst defaulting to 10. + limit := ratelimiter.Limit{Requests: 10, Period: time.Second} + + fmt.Println("Same Limit — 10 per second — under each strategy.") + fmt.Println() + + limiters := map[string]ratelimiter.Limiter{} + + for _, name := range []string{"token_bucket", "leaky_bucket"} { + strategy, err := ratelimiter.ParseStrategy(name) + if err != nil { + panic(err) + } + + newLimiter, err := ratelimiter.NewLimiterFunc(strategy, limit) + if err != nil { + panic(err) + } + + l := newLimiter() + limiters[name] = l + + fmt.Printf(" %-14s first burst: %2d of 10 admitted\n", name, drain(l, 10)) + } + + // Both absorbed the burst. The difference is what happens next. + fmt.Println("\nBoth burst. Now wait 300ms and try ten more —") + fmt.Println("this is where they diverge:") + fmt.Println() + + time.Sleep(300 * time.Millisecond) + + for _, name := range []string{"token_bucket", "leaky_bucket"} { + fmt.Printf(" %-14s after 300ms: %2d of 10 admitted\n", name, drain(limiters[name], 10)) + } + + fmt.Println("\n token bucket refilled ~3 slots continuously (300ms at 10/s).") + fmt.Println(" leaky bucket released 3 slots, one per 100ms interval.") + + // Capacity 1 is the configuration a token bucket cannot express: no burst + // at all, ever. + fmt.Println("\nA leaky bucket at capacity 1 paces strictly — no burst, ever:") + + strict, err := ratelimiter.NewLimiterFunc( + ratelimiter.StrategyLeakyBucket, + ratelimiter.Limit{Requests: 10, Period: time.Second, Burst: 1}, + ) + if err != nil { + panic(err) + } + + fmt.Printf(" leaky (burst 1) first burst: %2d of 10 admitted\n", drain(strict(), 10)) + + // Wait shapes rather than drops: it sleeps until each request conforms. + lb := ratelimiter.NewLeakyBucket(50*time.Millisecond, 1) + start := time.Now() + + for range 5 { + if err := lb.Wait(context.Background()); err != nil { + panic(err) + } + } + + fmt.Printf("\nShaping: 5 requests through Wait at one per 50ms took %v\n", + time.Since(start).Round(10*time.Millisecond)) + fmt.Println("Wait sleeps until each request conforms instead of dropping it,") + fmt.Println("which is what you want draining a queue against a metered API.") +} diff --git a/leaky_bucket.go b/leaky_bucket.go new file mode 100644 index 0000000..c0ff0a2 --- /dev/null +++ b/leaky_bucket.go @@ -0,0 +1,275 @@ +package ratelimiter + +import ( + "context" + "sync" + "time" +) + +// LeakyBucket is a [Limiter] that enforces a minimum spacing between +// admissions, implemented by virtual scheduling (GCRA — the Generic Cell Rate +// Algorithm). +// +// # How it relates to the token bucket +// +// Given the same parameters, they behave IDENTICALLY. A token bucket refilling +// at rate r with capacity b admits exactly what a leaky bucket draining at r +// with capacity b admits — they are duals, and measured side by side in this +// package they agree request for request. Comparisons claiming otherwise have +// usually changed the burst between the two examples. +// +// What differs is how the limit is EXPRESSED, and the arithmetic underneath: +// +// - Configured by an interval rather than a rate and a burst, so strict +// pacing — "one every 100ms" — is the obvious configuration rather than a +// non-obvious burst of 1 that is easy to leave at a default. +// - Exact: one time.Time and integer-duration arithmetic, so there are no +// floating-point tokens accumulating rounding error over a long uptime. +// +// Reach for it when you are expressing a PACE. "No more than 1000 an hour" is a +// budget and reads better as a token bucket; "no more than one call every +// 100ms" is a pace and reads better as this. +// +// # Virtual scheduling, not a queue +// +// There is no queue and no timer. The bucket keeps one instant: the theoretical +// arrival time (TAT) at which the next conforming request may be admitted. +// Admitting a request pushes the TAT forward by one emission interval; a +// request that arrives more than the burst allowance ahead of the TAT is +// refused, and the amount it is early by IS the retry delay. +// +// That makes every operation O(1), allocation-free and exact — no accumulated +// floating-point drift, because nothing accumulates. +// +// # Capacity +// +// capacity is how many requests may be taken back-to-back from an idle bucket. +// +// - capacity 1 is a strict leaky bucket: perfectly even spacing, no burst at +// all. This is what you want in front of a rate-limited third party. +// - capacity n allows n immediately and then paces at the interval, which is +// a token bucket's shape with a leaky bucket's recovery. +// +// # Shaping versus dropping +// +// [LeakyBucket.Wait] is the shaping call: it sleeps until the request conforms, +// so traffic is smoothed rather than rejected. [LeakyBucket.Allow] drops +// instead, for the cases where waiting is worse than failing — an HTTP handler +// that would rather answer 429 than hold a connection. +// +// The zero value is not usable; construct one with [NewLeakyBucket]. +type LeakyBucket struct { + mu sync.Mutex + + // tat is the theoretical arrival time of the next conforming request. + tat time.Time + + interval time.Duration // one emission: 1/rate + tolerance time.Duration // how far ahead of the TAT a request may arrive + capacity int + now func() time.Time +} + +// LeakyBucketOption configures a [LeakyBucket]. +type LeakyBucketOption func(*leakyBucketConfig) + +type leakyBucketConfig struct { + now func() time.Time +} + +// WithLeakyClock overrides the time source, which makes pacing deterministic in +// tests. +func WithLeakyClock(now func() time.Time) LeakyBucketOption { + return func(c *leakyBucketConfig) { + if now != nil { + c.now = now + } + } +} + +// NewLeakyBucket returns a [LeakyBucket] admitting one request every interval, +// allowing up to capacity back-to-back from idle. +// +// A capacity below one is raised to one: a bucket that admits nothing is a +// configuration error, and silently refusing everything is the least helpful +// way to report it. +// +// // one request per second, strictly evenly spaced +// lb := ratelimiter.NewLeakyBucket(time.Second, 1) +// +// // 60 per minute, allowing 5 back-to-back after a quiet period +// lb := ratelimiter.NewLeakyBucket(time.Minute/60, 5) +func NewLeakyBucket(interval time.Duration, capacity int, opts ...LeakyBucketOption) *LeakyBucket { + cfg := leakyBucketConfig{now: time.Now} + for _, opt := range opts { + opt(&cfg) + } + + if capacity < 1 { + capacity = 1 + } + + if interval < 0 { + interval = 0 + } + + return &LeakyBucket{ + interval: interval, + tolerance: time.Duration(capacity) * interval, + capacity: capacity, + now: cfg.now, + } +} + +// NewLeakyBucketFunc returns a factory suitable for [NewBucketLimiter], giving +// every key its own independently paced bucket. +// +// // each client IP gets one request per second, no bursting +// newLimiter := ratelimiter.NewLeakyBucketFunc(time.Second, 1) +// bl := ratelimiter.NewBucketLimiter(newLimiter, time.Minute, storage) +func NewLeakyBucketFunc(interval time.Duration, capacity int, opts ...LeakyBucketOption) func() Limiter { + return func() Limiter { + return NewLeakyBucket(interval, capacity, opts...) + } +} + +// Burst implements [Limiter]. It reports the capacity: how many requests may be +// taken back-to-back from an idle bucket. +func (b *LeakyBucket) Burst() int { return b.capacity } + +// Allow implements [Limiter]. It admits the request if it conforms, and never +// blocks. +func (b *LeakyBucket) Allow() bool { + ok, _ := b.reserve(true) + + return ok +} + +// Wait implements [Limiter]. It sleeps until the request conforms, which is how +// a leaky bucket SHAPES traffic rather than dropping it, and returns ctx's +// error if the context ends first. +// +// A request that cannot conform before ctx's deadline does not consume the +// slot: the reservation is rolled back, so a caller who gives up does not make +// the next caller wait for a request that never happened. +func (b *LeakyBucket) Wait(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + + ok, delay := b.reserve(true) + if ok { + return nil + } + + // reserve reported the wait without taking the slot; take it by waiting. + timer := time.NewTimer(delay) + defer timer.Stop() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + // The delay has elapsed, so this request now conforms. Claim it. + if ok, _ := b.reserve(true); ok { + return nil + } + + // Another goroutine took the slot while this one slept. Recurse rather + // than spin: the next delay is bounded by one interval. + return b.Wait(ctx) + } +} + +// Reserve implements [Reserver]. +// +// Unlike a window counter, a leaky bucket CAN give a slot back: the TAT is a +// single instant, so cancelling is rewinding it by one interval. That makes +// Retry-After exact and makes Cancel meaningful — a rejected request does not +// have to count against the caller. +func (b *LeakyBucket) Reserve() Reservation { + ok, delay := b.reserve(true) + + return &leakyReservation{bucket: b, ok: ok, delay: delay, taken: ok} +} + +// reserve is the GCRA step. When commit is true and the request conforms, the +// TAT is advanced. It returns whether the request conforms and, when it does +// not, how long until it would. +func (b *LeakyBucket) reserve(commit bool) (ok bool, delay time.Duration) { + b.mu.Lock() + defer b.mu.Unlock() + + now := b.now() + + // An interval of zero means no pacing at all: admit everything. + if b.interval == 0 { + return true, 0 + } + + tat := b.tat + if tat.Before(now) { + tat = now + } + + next := tat.Add(b.interval) + + // The earliest instant at which this request conforms. Arriving before it + // means the caller is going faster than the bucket drains. + allowAt := next.Add(-b.tolerance) + + if now.Before(allowAt) { + return false, allowAt.Sub(now) + } + + if commit { + b.tat = next + } + + return true, 0 +} + +// rewind returns one interval to the bucket. It never moves the TAT into the +// past relative to now, so cancelling a reservation cannot hand out extra +// budget that was never scheduled. +func (b *LeakyBucket) rewind() { + b.mu.Lock() + defer b.mu.Unlock() + + rewound := b.tat.Add(-b.interval) + + now := b.now() + if rewound.Before(now) { + rewound = now + } + + b.tat = rewound +} + +// leakyReservation is a [Reservation] over a leaky bucket slot. +type leakyReservation struct { + bucket *LeakyBucket + delay time.Duration + ok bool + taken bool + once sync.Once +} + +func (r *leakyReservation) OK() bool { return r.ok } +func (r *leakyReservation) Delay() time.Duration { return r.delay } + +// Cancel returns the slot if this reservation took one. It is safe to call more +// than once; only the first call has an effect, so a defer plus an explicit +// call cannot double-refund. +func (r *leakyReservation) Cancel() { + if !r.taken { + return + } + + r.once.Do(r.bucket.rewind) +} + +var ( + _ Limiter = (*LeakyBucket)(nil) + _ Reserver = (*LeakyBucket)(nil) +) diff --git a/leaky_bucket_test.go b/leaky_bucket_test.go new file mode 100644 index 0000000..9240334 --- /dev/null +++ b/leaky_bucket_test.go @@ -0,0 +1,450 @@ +package ratelimiter + +import ( + "context" + "sync" + "testing" + "time" + + "golang.org/x/time/rate" +) + +// TestLeakyBucketPacesStrictly is the property that separates a leaky bucket +// from a token bucket, and the reason to have one at all. +// +// A token bucket configured for 60/min admits all 60 in the first second and +// nothing for the rest of the minute. In front of a dependency with its own +// per-second quota, that burst is the thing that breaks you. A leaky bucket at +// capacity 1 admits one, then one per interval, for ever. +func TestLeakyBucketPacesStrictly(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := NewLeakyBucket(time.Second, 1, WithLeakyClock(clock)) + + if !b.Allow() { + t.Fatal("the first request from an idle bucket must be admitted") + } + + if b.Allow() { + t.Error("a second request in the same instant was admitted; capacity 1 means no burst at all") + } + + // Not quite an interval later: still too early. + advance(999 * time.Millisecond) + + if b.Allow() { + t.Error("a request 1ms before the interval elapsed was admitted") + } + + advance(time.Millisecond) + + if !b.Allow() { + t.Error("a request exactly one interval later was refused") + } +} + +func TestLeakyBucketCapacityAllowsABurstThenPaces(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := NewLeakyBucket(time.Second, 3, WithLeakyClock(clock)) + + for i := range 3 { + if !b.Allow() { + t.Fatalf("request %d refused; capacity 3 must admit three back-to-back from idle", i+1) + } + } + + if b.Allow() { + t.Error("a fourth back-to-back request was admitted against a capacity of three") + } + + // From here it paces at one per interval rather than refilling the burst. + advance(time.Second) + + if !b.Allow() { + t.Error("one interval later, one request should be admitted") + } + + if b.Allow() { + t.Error("two requests were admitted after a single interval; the burst must not refill in one step") + } +} + +// TestLeakyBucketRetryAfterIsExact: the amount a request is early by IS the +// delay, which is what makes Retry-After exact rather than a guess. A window +// counter can only report the distance to the window edge. +func TestLeakyBucketRetryAfterIsExact(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := NewLeakyBucket(time.Second, 1, WithLeakyClock(clock)) + + if !b.Allow() { + t.Fatal("the first request should be admitted") + } + + advance(400 * time.Millisecond) + + res := b.Reserve() + if res.OK() { + t.Fatal("a request 400ms into a 1s interval should not conform") + } + + if got, want := res.Delay(), 600*time.Millisecond; got != want { + t.Errorf("Delay() = %v, want %v — exactly the remainder of the interval", got, want) + } +} + +// TestLeakyBucketCancelReturnsTheSlot: unlike a window counter, a leaky bucket +// CAN give a slot back, because the TAT is a single instant and cancelling is +// rewinding it. A rejected request therefore need not count against the caller. +func TestLeakyBucketCancelReturnsTheSlot(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + b := NewLeakyBucket(time.Second, 1, WithLeakyClock(clock)) + + res := b.Reserve() + if !res.OK() { + t.Fatal("the first reservation should conform") + } + + // Without the rollback this is refused: the slot is already spent. + res.Cancel() + + if !b.Allow() { + t.Error("the slot was not returned by Cancel") + } +} + +func TestLeakyBucketCancelIsIdempotent(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + + // Capacity 3, not 1. At capacity 1 the clamp inside rewind — which refuses + // to move the TAT into the past — masks a non-idempotent Cancel entirely, + // so a test written that way passes whether or not the guard exists. It is + // only when the TAT is several intervals ahead that repeated cancels can + // refund slots that were never taken. Found by mutating the guard away and + // watching the test still pass. + b := NewLeakyBucket(time.Second, 3, WithLeakyClock(clock)) + + // Spend the whole capacity, so the TAT is three intervals ahead. + res := b.Reserve() + if !res.OK() { + t.Fatal("the first reservation should conform") + } + + for range 2 { + if !b.Allow() { + t.Fatal("the capacity should be spendable") + } + } + + if b.Allow() { + t.Fatal("the bucket should be exhausted") + } + + // Cancel the ONE reservation three times. + res.Cancel() + res.Cancel() + res.Cancel() + + if !b.Allow() { + t.Fatal("the cancelled slot was not returned") + } + + // Exactly one slot came back, not three. + if b.Allow() { + t.Error("cancelling one reservation three times refunded more than one slot; that hands out budget the bucket never scheduled") + } +} + +func TestLeakyBucketCancelOnARefusedReservationDoesNothing(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + b := NewLeakyBucket(time.Second, 1, WithLeakyClock(clock)) + + if !b.Allow() { + t.Fatal("setup") + } + + res := b.Reserve() + if res.OK() { + t.Fatal("this reservation should have been refused") + } + + // It never took a slot, so returning one would be inventing budget. + res.Cancel() + + if b.Allow() { + t.Error("cancelling a refused reservation handed out a slot") + } +} + +// TestLeakyBucketWaitShapes: Wait is the shaping call. It sleeps until the +// request conforms rather than dropping it, which is the whole point of using a +// leaky bucket in front of something with its own quota. +func TestLeakyBucketWaitShapes(t *testing.T) { + t.Parallel() + + // A real clock here: the point is that Wait actually sleeps. + b := NewLeakyBucket(30*time.Millisecond, 1) + + if !b.Allow() { + t.Fatal("the first request should be admitted") + } + + start := time.Now() + + if err := b.Wait(context.Background()); err != nil { + t.Fatalf("Wait: %v", err) + } + + if elapsed := time.Since(start); elapsed < 20*time.Millisecond { + t.Errorf("Wait returned after %v; it should have paced to the interval rather than admitting immediately", elapsed) + } +} + +func TestLeakyBucketWaitHonoursContext(t *testing.T) { + t.Parallel() + + b := NewLeakyBucket(time.Hour, 1) + + if !b.Allow() { + t.Fatal("setup") + } + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + start := time.Now() + + if err := b.Wait(ctx); err == nil { + t.Fatal("Wait returned nil despite an expired context") + } + + if elapsed := time.Since(start); elapsed > time.Second { + t.Errorf("Wait blocked %v past its context; it must not wait out the interval", elapsed) + } +} + +func TestLeakyBucketWaitReturnsImmediatelyOnADeadContext(t *testing.T) { + t.Parallel() + + b := NewLeakyBucket(time.Hour, 1) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if err := b.Wait(ctx); err == nil { + t.Error("Wait admitted a request on an already-cancelled context") + } +} + +func TestLeakyBucketNormalisesConstruction(t *testing.T) { + t.Parallel() + + // A capacity below one would admit nothing at all. Refusing every request + // for ever is the least helpful way to report a configuration error. + b := NewLeakyBucket(time.Second, 0) + if got := b.Burst(); got != 1 { + t.Errorf("Burst() = %d, want 1; a capacity below one is raised to one", got) + } + + if !b.Allow() { + t.Error("a bucket built with capacity 0 admits nothing") + } + + // A zero interval means no pacing. + unpaced := NewLeakyBucket(0, 1) + for range 100 { + if !unpaced.Allow() { + t.Fatal("a zero interval should admit everything") + } + } +} + +func TestLeakyBucketIsSafeUnderConcurrency(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + b := NewLeakyBucket(time.Second, 50, WithLeakyClock(clock)) + + var ( + mu sync.Mutex + allowed int + ) + + var wg sync.WaitGroup + + for range 200 { + wg.Go(func() { + if b.Allow() { + mu.Lock() + allowed++ + mu.Unlock() + } + }) + } + + wg.Wait() + + // The clock never moves, so exactly the capacity should get through — no + // more (a lost update) and no fewer (a spurious refusal). + if allowed != 50 { + t.Errorf("allowed %d of 200 concurrent requests, want exactly 50 (the capacity)", allowed) + } +} + +func TestLeakyBucketFuncBuildsIndependentBuckets(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + + bl := NewBucketLimiter( + NewLeakyBucketFunc(time.Second, 1, WithLeakyClock(clock)), + time.Minute, + NewInMemoryStorage[string, Limiter](), + ) + defer bl.Close() + + if !bl.GetOrAdd("alice").Allow() { + t.Fatal("alice's first request should be admitted") + } + + if bl.GetOrAdd("alice").Allow() { + t.Error("alice got two in one instant against capacity 1") + } + + if !bl.GetOrAdd("bob").Allow() { + t.Error("bob was paced by alice's request; buckets must be independent") + } +} + +// TestLeakyBucketAndTokenBucketAgreeOnTheSameParameters pins the claim the +// documentation makes, because it is the claim most easily lost. +// +// A leaky bucket and a token bucket are DUALS: same rate, same capacity, same +// admissions. The comparison usually seen in the wild — "one bursts, the other +// paces" — is rigged by giving the two sides different capacities. +// +// This matters as a test rather than a footnote: an earlier draft of +// LEAKY_BUCKET.md and the README both made the stronger claim, and it was only +// measuring the two side by side that showed it was false. If someone later +// "fixes" the leaky bucket to behave differently at the same settings, this +// fails and the docs stay true. +func TestLeakyBucketAndTokenBucketAgreeOnTheSameParameters(t *testing.T) { + t.Parallel() + + const ( + perSecond = 10 + capacity = 1 + ) + + tb := RateLimiter{Limiter: rate.NewLimiter(rate.Limit(perSecond), capacity)} + lb := NewLeakyBucket(time.Second/perSecond, capacity) + + drain := func(l Limiter) int { + admitted := 0 + + for range 10 { + if l.Allow() { + admitted++ + } + } + + return admitted + } + + if got, want := drain(tb), drain(lb); got != want { + t.Errorf("first burst: token bucket admitted %d, leaky bucket %d — with equal parameters they must agree", got, want) + } + + time.Sleep(250 * time.Millisecond) + + if got, want := drain(tb), drain(lb); got != want { + t.Errorf("after 250ms: token bucket admitted %d, leaky bucket %d", got, want) + } +} + +func TestLeakyBucketNegativeIntervalIsClamped(t *testing.T) { + t.Parallel() + + // A negative interval would make the TAT run backwards, handing out + // unlimited budget. Clamping to zero degrades to "no pacing", which is at + // least a state the caller can observe. + b := NewLeakyBucket(-time.Second, 1) + for range 10 { + if !b.Allow() { + t.Fatal("a negative interval should degrade to no pacing, not to refusing everything") + } + } +} + +// TestLeakyBucketWaitRetriesWhenAnotherCallerTakesTheSlot covers the race Wait +// has to survive: it sleeps for the reported delay, but by the time it wakes, +// another goroutine may have taken the slot it was waiting for. It must wait +// again rather than admit a request that no longer conforms. +func TestLeakyBucketWaitRetriesWhenAnotherCallerTakesTheSlot(t *testing.T) { + t.Parallel() + + b := NewLeakyBucket(25*time.Millisecond, 1) + + // Exhaust it, then have several goroutines all wait. Each must be paced; + // none may jump the queue. + if !b.Allow() { + t.Fatal("setup") + } + + const waiters = 4 + + start := time.Now() + + var wg sync.WaitGroup + + for range waiters { + wg.Go(func() { + if err := b.Wait(context.Background()); err != nil { + t.Error(err) + } + }) + } + + wg.Wait() + + // Four waiters at one per 25ms cannot finish faster than three intervals. + if elapsed := time.Since(start); elapsed < 60*time.Millisecond { + t.Errorf("%d waiters completed in %v; they must be paced, not admitted together", waiters, elapsed) + } +} + +func TestLeakyBucketRewindNeverMovesTheTATIntoThePast(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := NewLeakyBucket(time.Second, 1, WithLeakyClock(clock)) + + res := b.Reserve() + if !res.OK() { + t.Fatal("setup") + } + + // Let the reservation age past its own interval, then cancel. Rewinding + // blindly would put the TAT a second in the past and hand out a slot that + // was never scheduled. + advance(5 * time.Second) + res.Cancel() + + if !b.Allow() { + t.Fatal("one slot should be available") + } + + if b.Allow() { + t.Error("cancelling an aged reservation granted more than one slot") + } +} diff --git a/rate.go b/rate.go index ea347ad..a1070bf 100644 --- a/rate.go +++ b/rate.go @@ -46,3 +46,9 @@ var ( _ Reserver = RateLimiter{} _ Reservation = (*rate.Reservation)(nil) ) + +// rateLimitOf converts a [Limit] into the tokens-per-second rate that +// golang.org/x/time/rate expects. +func rateLimitOf(l Limit) rate.Limit { + return rate.Limit(l.Rate()) +} diff --git a/strategy.go b/strategy.go new file mode 100644 index 0000000..a90c3c6 --- /dev/null +++ b/strategy.go @@ -0,0 +1,147 @@ +package ratelimiter + +import ( + "fmt" + "slices" + "time" +) + +// Strategy names how a limit is enforced. Two limits with identical numbers can +// behave completely differently depending on which one is chosen. +// +// The values are lowercase snake_case because they are meant to survive a round +// trip through configuration: a database column, a YAML field, an environment +// variable. [ParseStrategy] validates one coming back. +type Strategy string + +const ( + // StrategyTokenBucket admits a burst up to the capacity and refills + // continuously. It answers "have you gone over budget?". + // + // The right default for protecting your own service, where a short burst + // is harmless and the thing you care about is sustained volume. + StrategyTokenBucket Strategy = "token_bucket" + + // StrategyLeakyBucket enforces a minimum spacing between admissions. It + // answers "are you going too fast right now?". + // + // The right choice in front of a dependency with its own limit, where a + // burst is precisely what breaks you. See [LeakyBucket]. + StrategyLeakyBucket Strategy = "leaky_bucket" +) + +// Strategies returns every valid strategy, in a stable order. Useful for +// populating a form or validating a configuration schema. +func Strategies() []Strategy { + return []Strategy{StrategyTokenBucket, StrategyLeakyBucket} +} + +// String implements [fmt.Stringer]. +func (s Strategy) String() string { return string(s) } + +// Valid reports whether s names a strategy this package implements. +func (s Strategy) Valid() bool { return slices.Contains(Strategies(), s) } + +// ParseStrategy converts a configuration value into a [Strategy], rejecting +// anything unrecognised. +// +// It rejects rather than defaulting on purpose. A typo in a config file that +// silently becomes "token bucket" gives you a limiter that admits bursts you +// specifically asked it not to — and nothing about the running service says so. +func ParseStrategy(s string) (Strategy, error) { + if candidate := Strategy(s); candidate.Valid() { + return candidate, nil + } + + return "", fmt.Errorf("ratelimiter: unknown strategy %q, want one of %v", s, Strategies()) +} + +// StrategyOption configures the limiters built by [NewLimiterFunc]. +type StrategyOption func(*strategyConfig) + +type strategyConfig struct { + now func() time.Time +} + +// WithStrategyClock overrides the time source. +// +// It affects [StrategyLeakyBucket] only. The token bucket is +// golang.org/x/time/rate, which has no clock seam — so a test that needs a +// controllable clock across both strategies cannot have one, and this is said +// here rather than discovered. +func WithStrategyClock(now func() time.Time) StrategyOption { + return func(c *strategyConfig) { + if now != nil { + c.now = now + } + } +} + +// NewLimiterFunc returns a [Limiter] factory for the named strategy, suitable +// for [NewBucketLimiter]. +// +// One [Limit] describes both strategies — the same "N requests per period" an +// operator writes down — and the strategy decides how it is enforced: +// +// limit := ratelimiter.Limit{Requests: 60, Period: time.Minute} +// +// token_bucket → 60 available at once, refilling continuously +// leaky_bucket → one every second, evenly spaced +// +// The burst differs too: the token bucket uses [Limit.Capacity] as its bucket +// size, and the leaky bucket uses it as how many may be taken back-to-back from +// idle. Leave Burst unset for a strict leaky bucket by setting it to 1. +// +// It returns an error for an unknown strategy or a limit that describes +// nothing, because both usually arrive from configuration and both are worth +// failing loudly at startup rather than quietly at request time. +func NewLimiterFunc(strategy Strategy, limit Limit, opts ...StrategyOption) (func() Limiter, error) { + if !strategy.Valid() { + return nil, fmt.Errorf("ratelimiter: unknown strategy %q, want one of %v", strategy, Strategies()) + } + + if limit.Requests <= 0 || limit.Period <= 0 { + return nil, fmt.Errorf( + "ratelimiter: invalid limit {Requests:%d Period:%v}; both must be positive", + limit.Requests, limit.Period, + ) + } + + cfg := strategyConfig{} + for _, opt := range opts { + opt(&cfg) + } + + switch strategy { + case StrategyLeakyBucket: + // One emission per request across the period: 60/minute is one per + // second. Derived rather than configured, so the two strategies are + // described by the same numbers. + interval := limit.Period / time.Duration(limit.Requests) + + var leakyOpts []LeakyBucketOption + if cfg.now != nil { + leakyOpts = append(leakyOpts, WithLeakyClock(cfg.now)) + } + + return NewLeakyBucketFunc(interval, limit.Capacity(), leakyOpts...), nil + + case StrategyTokenBucket: + fallthrough + + default: + return NewRateLimiterFunc(rateLimitOf(limit), limit.Capacity()), nil + } +} + +// MustNewLimiterFunc is [NewLimiterFunc] for a strategy and limit that are +// known good at compile time. It panics on an error, so keep it out of any path +// that handles configuration supplied at run time. +func MustNewLimiterFunc(strategy Strategy, limit Limit, opts ...StrategyOption) func() Limiter { + f, err := NewLimiterFunc(strategy, limit, opts...) + if err != nil { + panic(err) + } + + return f +} diff --git a/strategy_test.go b/strategy_test.go new file mode 100644 index 0000000..3e32bb7 --- /dev/null +++ b/strategy_test.go @@ -0,0 +1,177 @@ +package ratelimiter + +import ( + "testing" + "time" +) + +func TestParseStrategy(t *testing.T) { + t.Parallel() + + for _, want := range Strategies() { + got, err := ParseStrategy(string(want)) + if err != nil { + t.Errorf("ParseStrategy(%q): %v", want, err) + } + + if got != want { + t.Errorf("ParseStrategy(%q) = %q", want, got) + } + } + + // Rejected rather than defaulted. A typo that silently became a token + // bucket would admit bursts the operator specifically asked it not to, and + // nothing about the running service would say so. + for _, bad := range []string{"", "tokenbucket", "Token_Bucket", "leaky", "gcra"} { + if _, err := ParseStrategy(bad); err == nil { + t.Errorf("ParseStrategy(%q) returned no error; an unrecognised strategy must not default", bad) + } + } +} + +func TestStrategyValidAndString(t *testing.T) { + t.Parallel() + + if !StrategyTokenBucket.Valid() || !StrategyLeakyBucket.Valid() { + t.Error("a bundled strategy reported itself invalid") + } + + if Strategy("nonsense").Valid() { + t.Error("an unknown strategy reported itself valid") + } + + if got := StrategyLeakyBucket.String(); got != "leaky_bucket" { + t.Errorf("String() = %q", got) + } +} + +// TestNewLimiterFuncSelectsTheAlgorithm is the point of the whole type: the +// same Limit, enforced two different ways, and the difference is visible in one +// instant. +func TestNewLimiterFuncSelectsTheAlgorithm(t *testing.T) { + t.Parallel() + + // 60 per minute. A token bucket makes all 60 available at once; a leaky + // bucket paces them one per second. + limit := Limit{Requests: 60, Period: time.Minute} + + t.Run("token bucket bursts", func(t *testing.T) { + t.Parallel() + + f, err := NewLimiterFunc(StrategyTokenBucket, limit) + if err != nil { + t.Fatal(err) + } + + l := f() + + allowed := 0 + + for range 60 { + if l.Allow() { + allowed++ + } + } + + if allowed != 60 { + t.Errorf("token bucket admitted %d of 60 immediately, want 60 — absorbing a burst is what it is for", allowed) + } + }) + + t.Run("leaky bucket paces", func(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + + f, err := NewLimiterFunc(StrategyLeakyBucket, Limit{Requests: 60, Period: time.Minute, Burst: 1}, + WithStrategyClock(clock)) + if err != nil { + t.Fatal(err) + } + + l := f() + + allowed := 0 + + for range 60 { + if l.Allow() { + allowed++ + } + } + + if allowed != 1 { + t.Errorf("leaky bucket admitted %d of 60 in one instant, want 1 — the burst is exactly what it exists to prevent", allowed) + } + + // One second later — 60 per minute — the next one conforms. + advance(time.Second) + + if !l.Allow() { + t.Error("the leaky bucket refused a request one interval later") + } + }) +} + +func TestNewLimiterFuncRejectsBadInput(t *testing.T) { + t.Parallel() + + good := Limit{Requests: 1, Period: time.Second} + + if _, err := NewLimiterFunc(Strategy("nope"), good); err == nil { + t.Error("an unknown strategy was accepted") + } + + for name, bad := range map[string]Limit{ + "zero requests": {Requests: 0, Period: time.Second}, + "zero period": {Requests: 1, Period: 0}, + "negative": {Requests: -1, Period: time.Second}, + } { + t.Run(name, func(t *testing.T) { + // These usually arrive from configuration, so failing at startup + // beats failing at request time. + if _, err := NewLimiterFunc(StrategyTokenBucket, bad); err == nil { + t.Error("an impossible limit was accepted") + } + }) + } +} + +func TestMustNewLimiterFuncPanicsOnBadInput(t *testing.T) { + t.Parallel() + + defer func() { + if recover() == nil { + t.Error("MustNewLimiterFunc did not panic on an unknown strategy") + } + }() + + _ = MustNewLimiterFunc(Strategy("nope"), Limit{Requests: 1, Period: time.Second}) +} + +func TestStrategyRoundTripsThroughConfiguration(t *testing.T) { + t.Parallel() + + // What a database column or a YAML field actually does with it. + for _, s := range Strategies() { + parsed, err := ParseStrategy(s.String()) + if err != nil || parsed != s { + t.Errorf("%q did not survive a round trip: %v, %v", s, parsed, err) + } + } +} + +func TestNewLimiterFuncBuildsIndependentLimiters(t *testing.T) { + t.Parallel() + + f := MustNewLimiterFunc(StrategyTokenBucket, Limit{Requests: 1, Period: time.Hour}) + + a, b := f(), f() + + if !a.Allow() { + t.Fatal("the first limiter should admit one") + } + + if !b.Allow() { + t.Error("the second limiter shared the first one's budget; the factory must build independent limiters") + } +}