From bf12ace89821236257f043c5b141d74f1561ab7a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Gonz=C3=A1lez=20Di=20Antonio?= Date: Tue, 25 Aug 2026 19:43:52 +0200 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20Backend=20=E2=80=94=20one=20limit?= =?UTF-8?q?=20shared=20across=20processes,=20any=20KV=20store?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Storage holds limiters IN THIS PROCESS. There was no seam for holding the COUNT somewhere else, so a limit shared across instances was out of reach without abandoning this package. Backend is that seam: type Backend interface { Take(ctx context.Context, key string, limit Limit, cost int) (Decision, error) } Implement it over Valkey, Redis, DynamoDB, Postgres — or use the bundled MemoryBackend — and every process shares one budget. WHY IT IS TAKE-SHAPED AND NOT GET/SET-SHAPED. The obvious design is a small key-value interface — Get, Set, Incr — with the token arithmetic in this package. That works in one process and is a race everywhere else: the update is read (tokens, ts) -> refill by elapsed time -> compare -> write back a read-modify-write. Two processes read the same state, both decide they may proceed, both write, and the limit silently becomes 2x. Making it safe needs a transaction with a retry loop on a key that is contended by definition, or a server-side script; neither is expressible through Get/Set. So the decision has to run where the state lives, and the interface has to be the DECISION. That is also precisely why Storage cannot do this: Storage is a container of Limiter values in this process, and its Load/Store shape IS the Get/Set shape. WHAT BackendLimiter ADDS, and why it is here rather than copied into every consumer: - A LOCAL FALLBACK. Neither answer to "the datastore is down" is acceptable alone: refusing turns a blip into a total outage, allowing deletes the limiter exactly when it is needed. With a fallback there is nothing to choose between -- an outage degrades to per-process limiting, which is what you had before adding a backend at all. - A CIRCUIT BREAKER. Falling back is only half a fallback. Without one, every request during an outage pays a failed round trip before reaching the local answer it was always going to get: the limiter keeps working and makes the whole service slower by the timeout. - A DEGRADED SIGNAL. A limiter silently enforcing N x the intended limit is invisible from a request. Alert on the state, not the error rate. MemoryBackend is deliberately a sliding-window counter rather than a token bucket, even though this package has a better in-process limiter. It is the same algorithm a datastore backend runs, so moving between them does not move the behaviour underneath you. RateLimiter remains the right choice when you want the best in-process limiter. TWO BUGS THE TESTS FOUND, both in code I had already written and believed: - The single-prober guard used a mutex with `defer Unlock()`, which releases when shouldAsk RETURNS -- nanoseconds later, long before the backend answers. Every caller acquired it in turn. Measured: 20 of 20 racers reached the backend, which is the thundering herd the guard exists to prevent. It is a CAS on an atomic held across the call now. - A FAILED probe did not release the slot, so the breaker wedged shut for ever: the backend was never re-tested and the limiter stayed degraded permanently, including long after recovery. Worse than the outage, because it does not end when the outage does. Benchmarks drove one API addition. BackendLimiter.Allow measured 391 ns/op against a backend that answers in 100 ns, all of it context.WithTimeout arming a timer per request. WithoutBackendTimeout removes the deadline for in-process backends: 110 ns/op. It is documented as never for a network backend, where the deadline is what bounds a hung datastore. Every test verified to fail by mutating the code it guards. Coverage 99.7%. Docs: new docs/BACKENDS.md; doc.go rewritten (it claimed distributed limiting was out of scope for the bundled types, which this makes false); TOKEN_BUCKET.md's "deliberately out of scope" section replaced with the real trade-off table; CUSTOM_STORAGE.md now routes readers to BACKENDS.md instead of only explaining what not to do; MIGRATION.md notes this release is additive. All four cross-link. Runnable example in examples/backend. Purely additive: every existing call keeps working. --- README.md | 52 ++- backend.go | 108 +++++++ backend_limiter.go | 408 ++++++++++++++++++++++++ backend_limiter_test.go | 664 +++++++++++++++++++++++++++++++++++++++ benchmark_test.go | 92 ++++++ doc.go | 62 +++- docs/BACKENDS.md | 205 ++++++++++++ docs/CUSTOM_STORAGE.md | 6 +- docs/MIGRATION.md | 9 + docs/TOKEN_BUCKET.md | 53 +++- examples/backend/main.go | 81 +++++ memory_backend.go | 254 +++++++++++++++ memory_backend_test.go | 290 +++++++++++++++++ 13 files changed, 2258 insertions(+), 26 deletions(-) create mode 100644 backend.go create mode 100644 backend_limiter.go create mode 100644 backend_limiter_test.go create mode 100644 docs/BACKENDS.md create mode 100644 examples/backend/main.go create mode 100644 memory_backend.go create mode 100644 memory_backend_test.go diff --git a/README.md b/README.md index 8167078..a4293a5 100644 --- a/README.md +++ b/README.md @@ -294,10 +294,58 @@ 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. +## Backends — one limit shared across processes + +`Storage` holds limiters **in this process**. `Backend` holds the **count**, +wherever you want it: an in-process map, Valkey, Redis, DynamoDB, Postgres. + +```go +type Backend interface { + Take(ctx context.Context, key string, limit Limit, cost int) (Decision, error) +} +``` + +That is the whole interface. Implement it and every process shares one budget. + +```go +backend := ratelimiter.NewMemoryBackend() // or your own +limit := ratelimiter.Limit{Requests: 100, Period: time.Minute} + +bl := ratelimiter.NewBucketLimiter(nil, time.Minute, + ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), + ratelimiter.WithLimiterFactoryForKey( + ratelimiter.NewBackendLimiterFunc(backend, limit, + ratelimiter.WithFallback(local)), + ), +) +``` + +It is `Take`-shaped rather than `Get`/`Set`-shaped for a reason: a token bucket's +update is a read-modify-write, and split across a network that is a race in +which the limit silently becomes 2×. The decision has to run **where the state +lives**, so the interface is the *decision*, not the storage. + +`BackendLimiter` adds the three things nobody should have to reimplement: + +- a **local fallback**, because neither refusing nor allowing is an acceptable + answer to "the datastore is down" on its own; +- a **circuit breaker**, because falling back without one makes every request + during an outage pay a failed round trip first; +- a **degraded signal**, because a limiter silently enforcing N× the intended + limit is invisible from a request. + +**[docs/BACKENDS.md](docs/BACKENDS.md)** is the full guide, including a ~25-line +Valkey implementation. Runnable example: + +```bash +go run ./examples/backend +``` + ## Scope: single-process only -Token state lives in memory inside each `*rate.Limiter`, so this library -enforces limits **within one process**. Running N instances behind a load +Token state lives in memory inside each `*rate.Limiter`, so the **default** +limiter enforces limits **within one process**. Use a [`Backend`](#backends--one-limit-shared-across-processes) +to share one budget across instances. Running N instances behind a load balancer yields an effective global limit of up to N × `limit`. Global, cross-instance limiting requires a distributed algorithm (e.g. a Redis script) and is out of scope. The `Storage` interface is for custom *in-process* stores, diff --git a/backend.go b/backend.go new file mode 100644 index 0000000..f0a59c5 --- /dev/null +++ b/backend.go @@ -0,0 +1,108 @@ +package ratelimiter + +import ( + "context" + "time" +) + +// Limit is how much a rule allows, expressed the way an operator states it. +// +// Requests over Period, not a float rate: "300 per minute" is what a person +// says, what a form collects, and what a database row should hold. The rate is +// derived. Storing the float instead makes every caller guess at the window it +// came from. +type Limit struct { + // Requests is the budget over one Period. + Requests int + + // Period is the window the budget applies to. One second, one minute, one + // hour — whatever the rule says. + Period time.Duration + + // Burst is the capacity available in a single instant. Zero means + // Requests, which is the sensible default: a window-based backend has no + // separate notion of burst. + Burst int +} + +// Rate returns the sustained refill rate in tokens per second. +func (l Limit) Rate() float64 { + if l.Period <= 0 { + return 0 + } + + return float64(l.Requests) / l.Period.Seconds() +} + +// Capacity returns Burst, defaulting to Requests when Burst is unset. +func (l Limit) Capacity() int { + if l.Burst > 0 { + return l.Burst + } + + return l.Requests +} + +// Decision is what a [Backend] answers. +type Decision struct { + // Allowed reports whether the tokens were granted. + Allowed bool + + // Remaining is the budget left in the current window after this call. It + // is best-effort: a backend that cannot compute it cheaply may report -1, + // and callers must treat a negative value as "unknown" rather than zero. + Remaining int + + // RetryAfter is how long until the request would be granted. It is only + // meaningful when Allowed is false, and it may be an estimate — a + // window-based backend derives it from the window edge rather than from a + // continuous refill. + RetryAfter time.Duration +} + +// Backend is the pluggable state store for a rate limiter: an in-process map, +// Valkey, Redis, DynamoDB, or anything else that can count atomically per key. +// +// # Why this is Take-shaped and not Get/Set-shaped +// +// It is tempting to define a backend as a small key-value interface — Get, +// Set, Incr — and let this package do the token arithmetic. That does not work +// for anything but an in-process store, and the reason is the whole design: +// +// read (tokens, timestamp) → refill by elapsed time → compare → write back +// +// is a read-modify-write. Split across a network it is a race: two replicas +// read the same state, both decide they may proceed, and both write. The limit +// silently becomes 2x. Making it safe needs either a transaction with a retry +// loop on a key that is contended by definition, or a server-side script — and +// neither is expressible through a Get/Set interface. +// +// So the decision has to run WHERE THE STATE LIVES, and the interface has to be +// the decision, not the storage. [Backend.Take] is that: one call, one answer, +// atomicity owned by the implementation because only the implementation can +// provide it. +// +// This is also exactly why [Storage] cannot be used for distributed limiting. +// Storage holds [Limiter] VALUES in this process; it is a container, and its +// Load/Store shape is the Get/Set shape described above. See +// docs/CUSTOM_STORAGE.md. +// +// # The contract +// +// Implementations MUST be safe for concurrent use, and Take MUST be atomic per +// key: two concurrent Takes for the same key must not both succeed on the +// strength of the same remaining token. Everything else — the algorithm, the +// key encoding, the expiry — is the implementation's business. +// +// An implementation SHOULD expire its own state. A rate limiter's key space is +// usually unbounded (one entry per client address), so a backend that never +// forgets is a leak. +type Backend interface { + // Take attempts to consume cost tokens for key under limit, and reports + // what happened. cost is normally 1. + // + // An error means the decision could not be made — not that it was denied. + // Callers decide what an unavailable backend implies; [BackendLimiter] + // falls back to a local limiter rather than guessing. + Take(ctx context.Context, key string, limit Limit, cost int) (Decision, error) +} diff --git a/backend_limiter.go b/backend_limiter.go new file mode 100644 index 0000000..8359aef --- /dev/null +++ b/backend_limiter.go @@ -0,0 +1,408 @@ +package ratelimiter + +import ( + "context" + "log/slog" + "sync/atomic" + "time" +) + +// Default circuit-breaker settings for [BackendLimiter]. Small numbers on +// purpose: being wrong in the pessimistic direction costs a few seconds of +// local-only limiting, and being wrong in the optimistic direction costs a +// timeout on every request. +const ( + DefaultBackendFailureThreshold = 3 + DefaultBackendCooldown = 5 * time.Second + DefaultBackendTimeout = 100 * time.Millisecond +) + +// BackendLimiter is a [Limiter] whose state lives in a [Backend] — an +// in-process map, Valkey, Redis, anything that can count atomically per key. +// +// It is bound to one key, because a Limiter is: Allow and Wait take no +// arguments, so the instance IS the bucket. Build them with +// [NewBackendLimiterFunc] and [WithLimiterFactoryForKey], which is what gives +// each one its key. +// +// # What it adds over calling a Backend directly +// +// Three things, and they are the reason this is in the library rather than +// copied into every consumer: +// +// 1. A LOCAL FALLBACK. A remote backend can be unreachable, and neither +// answer to that is acceptable on its own: refusing turns a cache blip into +// a total outage, and allowing deletes the limiter exactly when it is +// needed. With a fallback limiter there is nothing to choose between — +// there is always a local answer, so an outage degrades to per-process +// limiting instead of to no limiting or to no service. +// +// 2. A CIRCUIT BREAKER. Falling back is only half a fallback. Without this, +// every request during an outage pays a failed round trip — connect, wait, +// time out — before reaching the local answer it was always going to get. +// The limiter would keep limiting correctly and make the whole service +// slower by the timeout, for the duration of the outage. After +// FailureThreshold consecutive failures the backend is skipped entirely +// until Cooldown elapses, then probed once. +// +// 3. A TIMEOUT. This call sits in front of everything else the process does, +// so it needs a budget far tighter than a normal query. +// +// # Being in the degraded state must be visible +// +// It is invisible from a request: the service keeps working, keeps limiting, +// and is silently enforcing N times the intended limit across N processes. +// [BackendLimiter.Degraded] reports it, and OnDegraded fires on each +// transition. Alert on the state, not on the error rate — a handful of errors +// is noise, a sustained fallback is the thing somebody has to know about. +type BackendLimiter struct { + backend Backend + key string + limit Limit + fallback Limiter + + timeout time.Duration + failureThreshold int + cooldown time.Duration + + logger *slog.Logger + onDegraded func(key string, degraded bool) + clock func() time.Time + failures atomic.Int64 + skipUntil atomic.Int64 // unix nanos; zero means "ask the backend" + degraded atomic.Bool + // probing is held by the single caller re-testing a failed backend after + // the cooldown. See shouldAsk for why this is a CAS and not a mutex. + probing atomic.Bool +} + +// BackendLimiterOption configures a [BackendLimiter]. +type BackendLimiterOption func(*backendLimiterConfig) + +type backendLimiterConfig struct { + fallback Limiter + timeout time.Duration + failureThreshold int + cooldown time.Duration + logger *slog.Logger + onDegraded func(key string, degraded bool) + clock func() time.Time +} + +// WithFallback supplies the limiter consulted when the backend cannot answer. +// +// Strongly recommended, and the reason is in the type: [Backend.Take] returns +// an error, [Limiter.Allow] returns a bool. Without a fallback there is nowhere +// for "I do not know" to go, and the limiter has to invent an answer. With one, +// it does not have to. +// +// When it is absent an unreachable backend ALLOWS the request, because refusing +// would make the backend a hard dependency of every request — the failure mode +// that is worse than the one being prevented. The choice is logged. +func WithFallback(l Limiter) BackendLimiterOption { + return func(c *backendLimiterConfig) { c.fallback = l } +} + +// WithBackendTimeout bounds a single Take. Defaults to +// [DefaultBackendTimeout]. A value <= 0 is ignored; use +// [WithoutBackendTimeout] to remove the bound. +func WithBackendTimeout(d time.Duration) BackendLimiterOption { + return func(c *backendLimiterConfig) { + if d > 0 { + c.timeout = d + } + } +} + +// WithoutBackendTimeout removes the per-call deadline. +// +// Use it for an IN-PROCESS backend, where the timeout can only cost and never +// help: context.WithTimeout allocates and arms a timer on every request, which +// measurably dominates a backend that answers in ~100ns. Measured with +// BenchmarkBackendLimiterAllow at the time this was added: +// +// with a timeout: 391 ns/op +// without: 110 ns/op +// +// Never use it for a network backend. Without a deadline a hung datastore +// blocks the request that touched it for however long the client's own +// timeouts allow, which is the failure the circuit breaker exists to bound — +// and the breaker cannot trip on a call that has not returned. +func WithoutBackendTimeout() BackendLimiterOption { + return func(c *backendLimiterConfig) { c.timeout = 0 } +} + +// WithCircuitBreaker sets how many consecutive failures skip the backend, and +// for how long. A threshold of zero or less disables the breaker, which means +// accepting a failed round trip on every request during an outage. +func WithCircuitBreaker(threshold int, cooldown time.Duration) BackendLimiterOption { + return func(c *backendLimiterConfig) { + c.failureThreshold = threshold + if cooldown > 0 { + c.cooldown = cooldown + } + } +} + +// WithBackendLogger sets the logger used for backend failures and for +// transitions in and out of the degraded state. Defaults to [slog.Default]. +func WithBackendLogger(l *slog.Logger) BackendLimiterOption { + return func(c *backendLimiterConfig) { + if l != nil { + c.logger = l + } + } +} + +// WithOnDegraded registers a callback fired on each transition into or out of +// the degraded state. Use it to drive a gauge. +func WithOnDegraded(f func(key string, degraded bool)) BackendLimiterOption { + return func(c *backendLimiterConfig) { c.onDegraded = f } +} + +// WithBackendClock overrides the time source, for tests. +func WithBackendClock(now func() time.Time) BackendLimiterOption { + return func(c *backendLimiterConfig) { + if now != nil { + c.clock = now + } + } +} + +// NewBackendLimiter returns a [BackendLimiter] for one key. +// +// Prefer [NewBackendLimiterFunc] with [WithLimiterFactoryForKey]: a +// BucketLimiter needs a factory, and building limiters by hand means keeping +// the key and the limiter in step yourself. +func NewBackendLimiter(b Backend, key string, limit Limit, opts ...BackendLimiterOption) *BackendLimiter { + cfg := backendLimiterConfig{ + timeout: DefaultBackendTimeout, + failureThreshold: DefaultBackendFailureThreshold, + cooldown: DefaultBackendCooldown, + logger: slog.Default(), + clock: time.Now, + } + for _, opt := range opts { + opt(&cfg) + } + + return &BackendLimiter{ + backend: b, + key: key, + limit: limit, + fallback: cfg.fallback, + timeout: cfg.timeout, + failureThreshold: cfg.failureThreshold, + cooldown: cfg.cooldown, + logger: cfg.logger, + onDegraded: cfg.onDegraded, + clock: cfg.clock, + } +} + +// NewBackendLimiterFunc returns a factory suitable for +// [WithLimiterFactoryForKey], building one [BackendLimiter] per key. +// +// backend := ratelimiter.NewMemoryBackend() // or a Valkey one +// limit := ratelimiter.Limit{Requests: 100, Period: time.Minute} +// +// bl := ratelimiter.NewBucketLimiter(nil, time.Minute, +// ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), +// ratelimiter.WithLimiterFactoryForKey( +// ratelimiter.NewBackendLimiterFunc(backend, limit, +// ratelimiter.WithFallback(localLimiter)), +// ), +// ) +func NewBackendLimiterFunc(b Backend, limit Limit, opts ...BackendLimiterOption) func(string) Limiter { + return func(key string) Limiter { + return NewBackendLimiter(b, key, limit, opts...) + } +} + +// Burst implements [Limiter]. +func (l *BackendLimiter) Burst() int { return l.limit.Capacity() } + +// Degraded reports whether the backend is currently being skipped. Alert on +// this rather than on an error count: it is the state that changes what the +// limiter enforces. +func (l *BackendLimiter) Degraded() bool { return l.degraded.Load() } + +// Allow implements [Limiter]. +func (l *BackendLimiter) Allow() bool { + return l.take(context.Background(), 1).Allowed +} + +// Wait implements [Limiter]. It sleeps until the backend would grant a token or +// ctx is done, whichever comes first. +func (l *BackendLimiter) Wait(ctx context.Context) error { + for { + d := l.take(ctx, 1) + if d.Allowed { + return nil + } + + delay := d.RetryAfter + if delay <= 0 { + delay = l.limit.Period + } + + timer := time.NewTimer(delay) + + select { + case <-ctx.Done(): + timer.Stop() + + return ctx.Err() + case <-timer.C: + } + } +} + +// Reserve implements [Reserver], so callers such as HTTP middleware can emit +// Retry-After and RateLimit-Reset for a backend-managed limit. +// +// The reservation is NOT cancellable in the way a token bucket's is: a window +// counter has no way to hand a token back that is distinguishable from +// spending one less. Cancel is therefore a no-op, and that is stated rather +// than hidden — a caller that rejects a request has still consumed from the +// window. Over-counting refused requests is the conservative direction, and +// closing it would cost a second round trip on the rejection path. +func (l *BackendLimiter) Reserve() Reservation { + d := l.take(context.Background(), 1) + + return backendReservation{ok: d.Allowed, delay: d.RetryAfter} +} + +func (l *BackendLimiter) take(ctx context.Context, cost int) Decision { + ask, isProbe := l.shouldAsk() + if !ask { + return l.fallbackDecision() + } + + if isProbe { + // Release the probe slot however this call ends, so a failed probe + // does not wedge the breaker shut until the process restarts. + defer l.probing.Store(false) + } + + if l.timeout > 0 { + var cancel context.CancelFunc + + ctx, cancel = context.WithTimeout(ctx, l.timeout) + defer cancel() + } + + d, err := l.backend.Take(ctx, l.key, l.limit, cost) + if err != nil { + l.recordFailure(err) + + return l.fallbackDecision() + } + + l.recordSuccess() + + return d +} + +// shouldAsk reports whether the backend is worth calling, and whether this +// caller is the one probing a failed backend. +// +// Once the failure threshold is crossed the backend is skipped until the +// cooldown elapses. Exactly one caller then gets through to re-test it, and the +// rest keep using the fallback until it reports — a cooldown expiry must not +// release every in-flight request at a datastore that is, by hypothesis, +// already unwell. +// +// The single-prober guarantee is a CAS on probing, not a mutex held across the +// call. A mutex looks right and is not: with `defer Unlock()` the lock is +// released when this function returns, nanoseconds later and long before the +// backend answers, so every caller acquires it in turn and all of them probe. +// TestBackendLimiterProbesOnceUnderConcurrency measured exactly that — 20 of 20 +// callers reached the backend. +func (l *BackendLimiter) shouldAsk() (ask, isProbe bool) { + if l.failureThreshold <= 0 { + return true, false + } + + if l.failures.Load() < int64(l.failureThreshold) { + return true, false + } + + if l.clock().UnixNano() < l.skipUntil.Load() { + return false, false + } + + // Cooldown elapsed. The winner of this CAS holds the probe slot until its + // call finishes; see take. + if !l.probing.CompareAndSwap(false, true) { + return false, false + } + + return true, true +} + +func (l *BackendLimiter) recordFailure(err error) { + n := l.failures.Add(1) + l.skipUntil.Store(l.clock().Add(l.cooldown).UnixNano()) + + if l.failureThreshold > 0 && n == int64(l.failureThreshold) && l.degraded.CompareAndSwap(false, true) { + l.logger.Warn("rate limiter backend unavailable; falling back to the local limiter", + "key", l.key, + "consecutive_failures", n, + "cooldown", l.cooldown, + "effect", "the limit is now enforced per process, not globally", + "error", err, + ) + + if l.onDegraded != nil { + l.onDegraded(l.key, true) + } + + return + } + + l.logger.Debug("rate limiter backend call failed", "key", l.key, "error", err) +} + +func (l *BackendLimiter) recordSuccess() { + l.failures.Store(0) + l.skipUntil.Store(0) + + if l.degraded.CompareAndSwap(true, false) { + l.logger.Info("rate limiter backend recovered; the limit is global again", "key", l.key) + + if l.onDegraded != nil { + l.onDegraded(l.key, false) + } + } +} + +// fallbackDecision answers from the local limiter, or allows when there is +// none. See [WithFallback] for why allowing is the default. +func (l *BackendLimiter) fallbackDecision() Decision { + if l.fallback == nil { + return Decision{Allowed: true, Remaining: -1} + } + + if !l.fallback.Allow() { + return Decision{Allowed: false, Remaining: 0, RetryAfter: l.limit.Period} + } + + return Decision{Allowed: true, Remaining: -1} +} + +// backendReservation is a Reservation over a decision already made. See +// [BackendLimiter.Reserve] for why Cancel does nothing. +type backendReservation struct { + ok bool + delay time.Duration +} + +func (r backendReservation) OK() bool { return r.ok } +func (r backendReservation) Delay() time.Duration { return r.delay } +func (r backendReservation) Cancel() {} + +var ( + _ Limiter = (*BackendLimiter)(nil) + _ Reserver = (*BackendLimiter)(nil) +) diff --git a/backend_limiter_test.go b/backend_limiter_test.go new file mode 100644 index 0000000..54e92f6 --- /dev/null +++ b/backend_limiter_test.go @@ -0,0 +1,664 @@ +package ratelimiter + +import ( + "context" + "errors" + "io" + "log/slog" + "sync" + "sync/atomic" + "testing" + "time" +) + +var errBackendDown = errors.New("backend unreachable") + +// scriptedBackend answers however the test tells it to, and counts calls — +// which is what the circuit-breaker tests actually assert, because the OUTCOME +// is identical whether or not the breaker works. Only the call count differs. +type scriptedBackend struct { + mu sync.Mutex + fail bool + allowed bool + calls int +} + +func (b *scriptedBackend) Take(context.Context, string, Limit, int) (Decision, error) { + b.mu.Lock() + defer b.mu.Unlock() + + b.calls++ + + if b.fail { + return Decision{}, errBackendDown + } + + return Decision{Allowed: b.allowed, Remaining: 1}, nil +} + +func (b *scriptedBackend) setFail(v bool) { + b.mu.Lock() + defer b.mu.Unlock() + + b.fail = v +} + +func (b *scriptedBackend) callCount() int { + b.mu.Lock() + defer b.mu.Unlock() + + return b.calls +} + +// alwaysLimiter is a fallback whose answer the test controls. +type alwaysLimiter struct { + allow bool + calls atomic.Int64 +} + +func (l *alwaysLimiter) Burst() int { return 1 } +func (l *alwaysLimiter) Allow() bool { + l.calls.Add(1) + + return l.allow +} +func (l *alwaysLimiter) Wait(context.Context) error { return nil } + +func quietLogger() *slog.Logger { + return slog.New(slog.NewTextHandler(io.Discard, nil)) +} + +func TestBackendLimiterAnswersFromTheBackend(t *testing.T) { + t.Parallel() + + for _, allowed := range []bool{true, false} { + b := &scriptedBackend{allowed: allowed} + l := NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithBackendLogger(quietLogger())) + + if got := l.Allow(); got != allowed { + t.Errorf("Allow() = %v, want %v", got, allowed) + } + } +} + +// TestBackendLimiterFallsBackWhenTheBackendFails is the property that makes a +// remote backend safe to depend on: neither answer to "the backend is down" is +// acceptable on its own. Refusing turns a cache blip into a total outage; +// allowing deletes the limiter exactly when it is needed. With a fallback there +// is nothing to choose between. +func TestBackendLimiterFallsBackWhenTheBackendFails(t *testing.T) { + t.Parallel() + + t.Run("the fallback refuses", func(t *testing.T) { + t.Parallel() + + b := &scriptedBackend{fail: true} + fb := &alwaysLimiter{allow: false} + + l := NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithFallback(fb), WithBackendLogger(quietLogger())) + + if l.Allow() { + t.Error("allowed while the backend was down and the fallback refused") + } + + if fb.calls.Load() == 0 { + t.Error("the fallback was never consulted") + } + }) + + t.Run("the fallback allows", func(t *testing.T) { + t.Parallel() + + b := &scriptedBackend{fail: true} + fb := &alwaysLimiter{allow: true} + + l := NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithFallback(fb), WithBackendLogger(quietLogger())) + + if !l.Allow() { + t.Error("refused while the fallback allowed; a backend outage must not reject traffic") + } + }) + + t.Run("no fallback allows, rather than making the backend a hard dependency", func(t *testing.T) { + t.Parallel() + + b := &scriptedBackend{fail: true} + l := NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithBackendLogger(quietLogger())) + + if !l.Allow() { + t.Error("refused with no fallback configured; that would make every request depend on the backend being up") + } + }) +} + +// TestBackendLimiterStopsCallingAFailedBackend asserts the CALL COUNT, not the +// outcome, and that is the whole point: with or without the breaker the request +// is decided the same way. What differs is whether every request first pays a +// failed round trip — connect, wait, time out — which during an outage makes +// the whole service slower by the timeout while the limiter keeps working. +func TestBackendLimiterStopsCallingAFailedBackend(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + b := &scriptedBackend{fail: true} + fb := &alwaysLimiter{allow: true} + + l := NewBackendLimiter(b, "k", Limit{Requests: 100, Period: time.Second}, + WithFallback(fb), + WithCircuitBreaker(3, 5*time.Second), + WithBackendClock(clock), + WithBackendLogger(quietLogger()), + ) + + for range 50 { + l.Allow() + } + + if got := b.callCount(); got != 3 { + t.Errorf("backend called %d times across 50 requests, want 3; after the threshold it must be skipped entirely", got) + } + + if !l.Degraded() { + t.Error("Degraded() is false while the backend is being skipped; the state is invisible from a request and has to be reportable") + } +} + +func TestBackendLimiterProbesAfterTheCooldownAndRecovers(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := &scriptedBackend{fail: true} + fb := &alwaysLimiter{allow: true} + + var transitions []bool + + var mu sync.Mutex + + l := NewBackendLimiter(b, "k", Limit{Requests: 100, Period: time.Second}, + WithFallback(fb), + WithCircuitBreaker(3, 5*time.Second), + WithBackendClock(clock), + WithBackendLogger(quietLogger()), + WithOnDegraded(func(_ string, d bool) { + mu.Lock() + transitions = append(transitions, d) + mu.Unlock() + }), + ) + + for range 10 { + l.Allow() + } + + if got := b.callCount(); got != 3 { + t.Fatalf("backend called %d times, want 3", got) + } + + // Inside the cooldown nothing gets through. + advance(2 * time.Second) + l.Allow() + + if got := b.callCount(); got != 3 { + t.Errorf("backend called %d times inside the cooldown, want 3", got) + } + + // Past it, exactly one probe. + advance(4 * time.Second) + b.setFail(false) + b.allowed = true + + if !l.Allow() { + t.Error("the probe should have been allowed once the backend recovered") + } + + if got := b.callCount(); got != 4 { + t.Errorf("backend called %d times after the cooldown, want 4", got) + } + + if l.Degraded() { + t.Error("still degraded after a successful probe") + } + + // And it is asking again, every time. + l.Allow() + + if got := b.callCount(); got != 5 { + t.Errorf("backend called %d times after recovery, want 5; a success must reset the breaker", got) + } + + mu.Lock() + defer mu.Unlock() + + if len(transitions) != 2 || !transitions[0] || transitions[1] { + t.Errorf("degraded transitions = %v, want [true false]; a gauge needs both edges", transitions) + } +} + +// TestBackendLimiterReportsASecondOutage covers what a recovery has to leave +// behind. The degraded transition fires when the consecutive-failure count +// REACHES the threshold, so if a success does not reset that count the second +// outage counts 4, 5, 6... and never equals it again: the limiter would go +// degraded silently, for ever after, having reported it correctly exactly once. +// +// A monitoring signal that works the first time and never again is worse than +// none, because it is trusted. +func TestBackendLimiterReportsASecondOutage(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := &scriptedBackend{fail: true} + + var transitions []bool + + var mu sync.Mutex + + l := NewBackendLimiter(b, "k", Limit{Requests: 100, Period: time.Second}, + WithFallback(&alwaysLimiter{allow: true}), + WithCircuitBreaker(3, 5*time.Second), + WithBackendClock(clock), + WithBackendLogger(quietLogger()), + WithOnDegraded(func(_ string, d bool) { + mu.Lock() + transitions = append(transitions, d) + mu.Unlock() + }), + ) + + // Outage one. + for range 5 { + l.Allow() + } + + // Recover. + advance(6 * time.Second) + b.setFail(false) + l.Allow() + + // Outage two. + b.setFail(true) + + for range 10 { + advance(6 * time.Second) // past each cooldown, so every call reaches the backend + l.Allow() + } + + if !l.Degraded() { + t.Error("Degraded() is false during a second outage") + } + + mu.Lock() + defer mu.Unlock() + + if len(transitions) != 3 { + t.Fatalf("degraded transitions = %v, want [true false true]; a recovery must reset the failure count or the second outage is never reported", transitions) + } + + if !transitions[0] || transitions[1] || !transitions[2] { + t.Errorf("degraded transitions = %v, want [true false true]", transitions) + } +} + +func TestBackendLimiterWithoutABreakerKeepsCalling(t *testing.T) { + t.Parallel() + + b := &scriptedBackend{fail: true} + l := NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithFallback(&alwaysLimiter{allow: true}), + WithCircuitBreaker(0, 0), // disabled + WithBackendLogger(quietLogger()), + ) + + for range 10 { + l.Allow() + } + + if got := b.callCount(); got != 10 { + t.Errorf("backend called %d times with the breaker disabled, want 10", got) + } +} + +func TestBackendLimiterReserveReportsTheDelay(t *testing.T) { + t.Parallel() + + b := &scriptedBackend{allowed: false} + l := NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithBackendLogger(quietLogger())) + + var r Reserver = l + + res := r.Reserve() + if res.OK() { + t.Error("Reserve reported OK for a refused decision") + } + + // Cancel is a no-op by design; calling it must not panic and must not + // change the answer. + res.Cancel() + + if res.OK() { + t.Error("Cancel changed the reservation") + } +} + +func TestBackendLimiterBurstReportsCapacity(t *testing.T) { + t.Parallel() + + l := NewBackendLimiter(&scriptedBackend{}, "k", + Limit{Requests: 10, Period: time.Second}, WithBackendLogger(quietLogger())) + + if got := l.Burst(); got != 10 { + t.Errorf("Burst() = %d, want 10 (Requests, since Burst is unset)", got) + } + + l = NewBackendLimiter(&scriptedBackend{}, "k", + Limit{Requests: 10, Period: time.Second, Burst: 25}, WithBackendLogger(quietLogger())) + + if got := l.Burst(); got != 25 { + t.Errorf("Burst() = %d, want 25", got) + } +} + +// TestBackendLimiterThroughBucketLimiter is the whole assembly: a BucketLimiter +// handing out one key-bound BackendLimiter per key, which is how a consumer +// actually wires this. +func TestBackendLimiterThroughBucketLimiter(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + backend := NewMemoryBackend(WithMemoryClock(clock), WithMemorySweepInterval(0)) + defer backend.Close() + + limit := Limit{Requests: 2, Period: time.Second} + + bl := NewBucketLimiter(nil, time.Minute, + NewInMemoryStorage[string, Limiter](), + WithLimiterFactoryForKey(NewBackendLimiterFunc(backend, limit, + WithBackendLogger(quietLogger()))), + ) + defer bl.Close() + + for i := range 2 { + if !bl.GetOrAdd("alice").Allow() { + t.Fatalf("alice request %d refused inside her budget", i+1) + } + } + + if bl.GetOrAdd("alice").Allow() { + t.Error("alice's third request was allowed against a budget of two") + } + + if !bl.GetOrAdd("bob").Allow() { + t.Error("bob was refused because alice spent her budget") + } +} + +func TestLimitHelpers(t *testing.T) { + t.Parallel() + + l := Limit{Requests: 300, Period: time.Minute} + if got := l.Rate(); got != 5 { + t.Errorf("Rate() = %v, want 5 (300 per minute)", got) + } + + if got := l.Capacity(); got != 300 { + t.Errorf("Capacity() = %d, want 300 when Burst is unset", got) + } + + if got := (Limit{Requests: 1, Period: 0}).Rate(); got != 0 { + t.Errorf("Rate() with a zero period = %v, want 0 rather than a division by zero", got) + } +} + +func TestBackendLimiterWaitReturnsWhenAllowed(t *testing.T) { + t.Parallel() + + l := NewBackendLimiter(&scriptedBackend{allowed: true}, "k", + Limit{Requests: 1, Period: time.Second}, + WithBackendLogger(quietLogger()), WithBackendTimeout(time.Second)) + + if err := l.Wait(context.Background()); err != nil { + t.Errorf("Wait returned %v for an allowed request", err) + } +} + +// TestBackendLimiterWaitHonoursContext: Wait must not outlive its context. A +// limiter that blocks past cancellation holds a request handler open long after +// the client has gone. +func TestBackendLimiterWaitHonoursContext(t *testing.T) { + t.Parallel() + + l := NewBackendLimiter(&scriptedBackend{allowed: false}, "k", + Limit{Requests: 1, Period: time.Hour}, + WithBackendLogger(quietLogger()), WithoutBackendTimeout()) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + start := time.Now() + + err := l.Wait(ctx) + if err == nil { + t.Fatal("Wait returned nil for a permanently refused limiter") + } + + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("Wait returned %v, want the context's error", err) + } + + if elapsed := time.Since(start); elapsed > time.Second { + t.Errorf("Wait blocked for %v past its context; it must not wait out the limit period", elapsed) + } +} + +func TestBackendReservationReportsDelay(t *testing.T) { + t.Parallel() + + b := &scriptedBackend{} + l := NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithBackendLogger(quietLogger())) + + // A refusal from the memory backend carries a real RetryAfter. + mem := NewMemoryBackend(WithMemorySweepInterval(0)) + defer mem.Close() + + limit := Limit{Requests: 1, Period: time.Minute} + ml := NewBackendLimiter(mem, "k", limit, WithBackendLogger(quietLogger())) + + if !ml.Allow() { + t.Fatal("the first request should be allowed") + } + + res := ml.Reserve() + if res.OK() { + t.Error("Reserve reported OK past the budget") + } + + if res.Delay() <= 0 { + t.Error("a refused reservation must report how long until it would be granted") + } + + res.Cancel() // no-op by design; must not panic + + _ = l +} + +func TestWithoutBackendTimeoutRemovesTheDeadline(t *testing.T) { + t.Parallel() + + // A backend that reports whether it saw a deadline. + var sawDeadline atomic.Bool + + b := backendFunc(func(ctx context.Context, _ string, _ Limit, _ int) (Decision, error) { + _, ok := ctx.Deadline() + sawDeadline.Store(ok) + + return Decision{Allowed: true}, nil + }) + + NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithBackendLogger(quietLogger()), WithoutBackendTimeout()).Allow() + + if sawDeadline.Load() { + t.Error("a deadline was applied despite WithoutBackendTimeout") + } + + NewBackendLimiter(b, "k", Limit{Requests: 1, Period: time.Second}, + WithBackendLogger(quietLogger()), WithBackendTimeout(time.Second)).Allow() + + if !sawDeadline.Load() { + t.Error("no deadline was applied despite WithBackendTimeout") + } +} + +// backendFunc adapts a function to Backend. +type backendFunc func(context.Context, string, Limit, int) (Decision, error) + +func (f backendFunc) Take(ctx context.Context, key string, limit Limit, cost int) (Decision, error) { + return f(ctx, key, limit, cost) +} + +// TestBackendLimiterProbesOnceUnderConcurrency: when the cooldown elapses, +// exactly ONE caller reaches the backend to find out whether it is back. +// +// Letting every in-flight request through at that instant turns each cooldown +// expiry into a thundering herd against a datastore that is, by hypothesis, +// already unwell — and the requests that pile up are precisely the ones the +// fallback could have answered for free. +func TestBackendLimiterProbesOnceUnderConcurrency(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + + var ( + blocking atomic.Bool + probes atomic.Int64 + inFlight atomic.Int64 + maxInFlight atomic.Int64 + release = make(chan struct{}) + ) + + b := backendFunc(func(context.Context, string, Limit, int) (Decision, error) { + probes.Add(1) + + n := inFlight.Add(1) + defer inFlight.Add(-1) + + for { + m := maxInFlight.Load() + if n <= m || maxInFlight.CompareAndSwap(m, n) { + break + } + } + + if blocking.Load() { + <-release + } + + return Decision{}, errBackendDown + }) + + l := NewBackendLimiter(b, "k", Limit{Requests: 100, Period: time.Second}, + WithFallback(&alwaysLimiter{allow: true}), + WithCircuitBreaker(1, 5*time.Second), + WithBackendClock(clock), + WithBackendLogger(quietLogger()), + WithoutBackendTimeout(), + ) + + // Trip the breaker with one non-blocking failure. + l.Allow() + + if got := probes.Load(); got != 1 { + t.Fatalf("setup: probes = %d, want 1", got) + } + + // Past the cooldown, race twenty callers at the probe window. + blocking.Store(true) + advance(6 * time.Second) + + go func() { + time.Sleep(50 * time.Millisecond) + close(release) + }() + + var wg sync.WaitGroup + + for range 20 { + wg.Go(func() { l.Allow() }) + } + + wg.Wait() + + if got := probes.Load(); got != 2 { + t.Errorf("probes = %d, want 2 (one to trip the breaker, one to re-test it); a cooldown expiry must not become a thundering herd", got) + } + + if got := maxInFlight.Load(); got != 1 { + t.Errorf("max concurrent backend calls = %d, want 1", got) + } +} + +// TestBackendLimiterRetriesAfterAFailedProbe: a probe that fails must not wedge +// the breaker shut. +// +// The single-prober guard is a slot one caller claims. If it is only released +// when the probe SUCCEEDS, a probe that fails keeps it for ever: no later +// caller can claim it, so the backend is never re-tested and the limiter stays +// degraded permanently — including long after the datastore has recovered. +// +// That is worse than the outage it is reacting to, because it does not end when +// the outage does, and nothing about a healthy-looking service says why. +func TestBackendLimiterRetriesAfterAFailedProbe(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := &scriptedBackend{fail: true} + + l := NewBackendLimiter(b, "k", Limit{Requests: 100, Period: time.Second}, + WithFallback(&alwaysLimiter{allow: true}), + WithCircuitBreaker(1, 5*time.Second), + WithBackendClock(clock), + WithBackendLogger(quietLogger()), + ) + + l.Allow() // trips the breaker + + if got := b.callCount(); got != 1 { + t.Fatalf("setup: backend called %d times, want 1", got) + } + + // First probe — still down. + advance(6 * time.Second) + l.Allow() + + if got := b.callCount(); got != 2 { + t.Fatalf("first probe: backend called %d times, want 2", got) + } + + // Second probe, another cooldown later. This is the one that never happens + // if a failed probe keeps the slot. + advance(6 * time.Second) + l.Allow() + + if got := b.callCount(); got != 3 { + t.Errorf("backend called %d times, want 3; a failed probe must release the slot or the backend is never re-tested", got) + } + + // And once it recovers, the limiter comes back. + b.setFail(false) + b.allowed = true + + advance(6 * time.Second) + + if !l.Allow() { + t.Error("the limiter never recovered after the backend came back") + } + + if l.Degraded() { + t.Error("still degraded after a successful probe") + } +} diff --git a/benchmark_test.go b/benchmark_test.go index fc464b7..8df0690 100644 --- a/benchmark_test.go +++ b/benchmark_test.go @@ -5,6 +5,11 @@ import ( "testing" "time" + "context" + "errors" + "io" + "log/slog" + "golang.org/x/time/rate" ) @@ -43,3 +48,90 @@ func BenchmarkGetOrAdd_Parallel(b *testing.B) { } }) } + +func BenchmarkMemoryBackendTake(b *testing.B) { + backend := NewMemoryBackend(WithMemorySweepInterval(0)) + defer backend.Close() + + limit := Limit{Requests: 1 << 30, Period: time.Hour} + ctx := context.Background() + + b.ResetTimer() + + for i := 0; b.Loop(); i++ { + _, _ = backend.Take(ctx, "key", limit, 1) + } +} + +// BenchmarkMemoryBackendTakeParallel is the one that matters: a rate limiter's +// hot path is contended by construction, which is why the backend is sharded. +func BenchmarkMemoryBackendTakeParallel(b *testing.B) { + backend := NewMemoryBackend(WithMemorySweepInterval(0)) + defer backend.Close() + + limit := Limit{Requests: 1 << 30, Period: time.Hour} + + b.ResetTimer() + + b.RunParallel(func(pb *testing.PB) { + ctx := context.Background() + i := 0 + + for pb.Next() { + _, _ = backend.Take(ctx, keys[i%len(keys)], limit, 1) + i++ + } + }) +} + +func BenchmarkBackendLimiterAllow(b *testing.B) { + backend := NewMemoryBackend(WithMemorySweepInterval(0)) + defer backend.Close() + + l := NewBackendLimiter(backend, "key", + Limit{Requests: 1 << 30, Period: time.Hour}, + WithoutBackendTimeout()) + + b.ResetTimer() + + for b.Loop() { + l.Allow() + } +} + +// BenchmarkBackendLimiterDegraded measures the path an outage puts every +// request on. It must not involve the backend at all — that is the whole point +// of the circuit breaker. +func BenchmarkBackendLimiterDegraded(b *testing.B) { + l := NewBackendLimiter(failingBackend{}, "key", + Limit{Requests: 1 << 30, Period: time.Hour}, + WithFallback(RateLimiter{rate.NewLimiter(rate.Inf, 1)}), + WithCircuitBreaker(1, time.Hour), + WithBackendLogger(slog.New(slog.NewTextHandler(io.Discard, nil))), + ) + + l.Allow() // trip the breaker + + b.ResetTimer() + + for b.Loop() { + l.Allow() + } +} + +type failingBackend struct{} + +func (failingBackend) Take(context.Context, string, Limit, int) (Decision, error) { + return Decision{}, errBenchDown +} + +var errBenchDown = errors.New("down") + +var keys = func() []string { + k := make([]string, 64) + for i := range k { + k[i] = "key-" + strconv.Itoa(i) + } + + return k +}() diff --git a/doc.go b/doc.go index 70a1dce..97a923a 100644 --- a/doc.go +++ b/doc.go @@ -4,14 +4,14 @@ // // The central type is [BucketLimiter], a manager that hands out an independent // [Limiter] per key (for example a user ID or IP address). Each key gets its -// own token bucket, so exhausting one key never affects another. +// own bucket, so exhausting one key never affects another. // // Limiters are created lazily on first use and evicted after they have been // idle (not accessed) for a configurable duration. Eviction is performed by a // single background goroutine that is started when the manager is created and // stopped by [BucketLimiter.Close]. // -// Basic usage: +// # Basic usage // // storage := ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter]() // newLimiter := ratelimiter.NewRateLimiterFunc(rate.Limit(5), 10) // 5 rps, burst 10 @@ -26,13 +26,53 @@ // A limiter may optionally also implement [Reserver] to reserve a token and // report the exact delay until it is valid; the default limiter from // [NewRateLimiterFunc] does, which lets HTTP middleware emit accurate -// Retry-After headers for any backend. See the examples directory for a runnable -// HTTP middleware. -// -// The [Storage] interface can be implemented to plug in a custom in-process -// store; see [Storage] and the customstorage example. Note that the -// token-bucket state lives in memory inside each *rate.Limiter, so this package -// targets single-process rate limiting. Distributed rate limiting across -// multiple instances requires a different algorithm (a datastore-backed -// [Limiter]) and is out of scope for the bundled types. +// Retry-After headers for any backend. See the examples directory for a +// runnable HTTP middleware. +// +// # Three extension points, and which one you want +// +// They are easy to confuse, and picking the wrong one produces a limiter that +// looks like it works. The rule of thumb: +// +// a different in-process CONTAINER (LRU, metrics, sharding) → Storage +// a different ALGORITHM (leaky bucket, GCRA, …) → Limiter +// a limit SHARED across processes → Backend +// +// [Storage] holds the per-key [Limiter] values in this process. It is a +// container, and it is in-process by definition: [BucketLimiter.GetOrAdd] hands +// the caller the limiter and the caller mutates it, so a Storage that +// serialised limiter state into a datastore would return a full bucket on every +// call and never limit anything. See docs/CUSTOM_STORAGE.md. +// +// [Limiter] is the decision. Implement it to change how tokens are accounted — +// including accounting them somewhere other than this process. +// +// [Backend] is where the count lives: an in-process map, Valkey, Redis, or +// anything else that can count atomically per key. It is deliberately +// Take-shaped rather than Get/Set-shaped, because the token update is a +// read-modify-write and splitting that across a network is a race in which the +// limit silently becomes twice what was configured. The decision has to run +// where the state lives. +// +// # Sharing one limit across processes +// +// backend := ratelimiter.NewMemoryBackend() // or your own, over Valkey/Redis +// limit := ratelimiter.Limit{Requests: 100, Period: time.Minute} +// +// bl := ratelimiter.NewBucketLimiter(nil, time.Minute, +// ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), +// ratelimiter.WithLimiterFactoryForKey( +// ratelimiter.NewBackendLimiterFunc(backend, limit, +// ratelimiter.WithFallback(local)), +// ), +// ) +// +// [WithLimiterFactoryForKey] is what makes this possible: a [Limiter] is bound +// to exactly one key — Allow and Wait take no arguments, so the instance IS the +// bucket — and a limiter whose counter lives elsewhere has to know which remote +// key is its own. +// +// [BackendLimiter] adds a local fallback, a circuit breaker and a degraded +// signal, because a remote backend can be unreachable and neither refusing nor +// allowing is an acceptable answer to that on its own. See docs/BACKENDS.md. package ratelimiter diff --git a/docs/BACKENDS.md b/docs/BACKENDS.md new file mode 100644 index 0000000..c688e46 --- /dev/null +++ b/docs/BACKENDS.md @@ -0,0 +1,205 @@ +# Backends — one limit shared across processes + +`Storage` holds limiters **in this process**. `Backend` holds the **count**, +wherever you want it: an in-process map, Valkey, Redis, DynamoDB, Postgres. + +```go +type Backend interface { + Take(ctx context.Context, key string, limit Limit, cost int) (Decision, error) +} +``` + +That is the whole interface. Implement it and every process shares one budget. + +## Why it is `Take`-shaped and not `Get`/`Set`-shaped + +This is the part worth reading before writing an implementation, because the +obvious design does not work. + +It is tempting to define a backend as a small key-value interface — `Get`, +`Set`, `Incr` — and let this package do the token arithmetic. For an in-process +store that is fine. Across a network it is a race: + +```text +read (tokens, timestamp) → refill by elapsed time → compare → write back +``` + +is a read-modify-write. Two processes read the same state, both decide they may +proceed, and both write. **The limit silently becomes 2×.** Making it safe needs +either a transaction with a retry loop — on a key that is contended by +definition, because one caller hammering one endpoint is the case a rate limiter +exists for — or a server-side script. Neither is expressible through `Get`/`Set`. + +So the decision has to run **where the state lives**, and the interface has to be +the *decision*, not the storage. `Take` is that: one call, one answer, atomicity +owned by the implementation because only the implementation can provide it. + +This is also exactly why [`Storage`](./CUSTOM_STORAGE.md) cannot be used for +distributed limiting. `Storage` holds `Limiter` **values** in this process; its +`Load`/`Store` shape *is* the `Get`/`Set` shape above. + +| You want… | Extension point | +| --- | --- | +| a different **in-process container** (LRU, metrics) | `Storage` | +| a different **algorithm** | `Limiter` | +| a **shared** limit across processes | **`Backend`** | + +## Wiring it up + +```go +backend := ratelimiter.NewMemoryBackend() // or your own +limit := ratelimiter.Limit{Requests: 100, Period: time.Minute} + +bl := ratelimiter.NewBucketLimiter(nil, time.Minute, + ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), + ratelimiter.WithLimiterFactoryForKey( + ratelimiter.NewBackendLimiterFunc(backend, limit, + ratelimiter.WithFallback(local), + ), + ), +) +``` + +The `Storage` is the ordinary in-memory one, and stays that way whichever backend +you use: it caches one lightweight handle per key. The count lives in the +backend. + +## Switching on configuration + +The whole "shared when the datastore is there, local when it is not" decision is +one branch: + +```go +newLimiter := func(key string) ratelimiter.Limiter { + local := ratelimiter.RateLimiter{Limiter: rate.NewLimiter(limit.Rate(), limit.Capacity())} + + if backend == nil { // cache disabled + return local + } + + return ratelimiter.NewBackendLimiter(backend, key, limit, + ratelimiter.WithFallback(local)) +} +``` + +## What `BackendLimiter` adds + +Three things, and they are in the library rather than copied into every consumer +because getting any of them wrong is expensive. + +### 1. A local fallback + +A remote backend can be unreachable, and **neither answer is acceptable on its +own**: + +- *Refusing* turns a cache blip into a total outage — self-inflicted, larger than + any attack the limiter prevents, arriving at the worst possible moment. +- *Allowing* deletes the limiter exactly when it is needed. + +With `WithFallback` there is nothing to choose between: there is always a local +answer, so an outage degrades to **per-process limiting** — which is what you +had before you added a backend at all. + +Without a fallback an unreachable backend allows the request, because the +alternative is making the backend a hard dependency of every request. The choice +is logged. + +### 2. A circuit breaker + +Falling back is only **half** a fallback. Without a breaker, every request +during an outage pays a failed round trip — connect, wait, time out — before +reaching the local answer it was always going to get. The limiter keeps limiting +correctly *and* makes the whole service slower by the timeout, for the duration +of the outage. + +After `FailureThreshold` consecutive failures the backend is skipped entirely +until `Cooldown` elapses, then probed once by a single caller. + +Recovery needs no repair: a window-keyed backend either finds or creates the +current window on the next call. + +### 3. A timeout + +This call sits in front of everything else the process does, so its budget +should be far tighter than a normal query. Default 100 ms. + +## Observability + +Being in the degraded state is **invisible from a request**: the service keeps +working, keeps limiting, and is silently enforcing N × the intended limit across +N processes. + +```go +ratelimiter.WithOnDegraded(func(key string, degraded bool) { + gauge.Set(key, degraded) +}) +``` + +**Alert on the state, not on the error rate.** A handful of errors is noise; a +sustained fallback is the thing somebody has to know about. + +## Writing a Valkey backend + +This package has no Valkey dependency and will not grow one, so the client code +is yours. It is short: + +```go +type ValkeyBackend struct{ client valkey.Client } + +func (b ValkeyBackend) Take(ctx context.Context, key string, limit ratelimiter.Limit, cost int) (ratelimiter.Decision, error) { + now := time.Now() + window := now.UnixNano() / int64(limit.Period) + k := fmt.Sprintf("rl:%s:%d", key, window) + + // INCR is atomic on its own — no script, no WATCH, no dedicated connection. + n, err := b.client.Do(ctx, b.client.B().Incr().Key(k).Build()).AsInt64() + if err != nil { + return ratelimiter.Decision{}, err + } + + // Only the first hit of a window needs an expiry. + if n == 1 { + _ = b.client.Do(ctx, b.client.B().Pexpire().Key(k). + Milliseconds(int64(2*limit.Period/time.Millisecond)).Build()).Error() + } + + if n > int64(limit.Requests) { + elapsed := time.Duration(now.UnixNano() % int64(limit.Period)) + + return ratelimiter.Decision{Allowed: false, RetryAfter: limit.Period - elapsed}, nil + } + + return ratelimiter.Decision{Allowed: true, Remaining: limit.Requests - int(n)}, nil +} +``` + +That is a **fixed** window — simple, one round trip, and it admits up to 2× the +budget across a boundary. `MemoryBackend` shows the two-window weighting that +reduces that to a small approximation error while still costing one counter per +window; the same technique works over `INCR` with two keys read together. + +If you want an exact token bucket with continuous refill, that needs a +server-side script, because the refill is a read-modify-write. `Take` does not +care which you choose — that is the point of putting the decision behind an +interface. + +## See also + +- [CUSTOM_STORAGE.md](./CUSTOM_STORAGE.md) — the in-process container interface, + and why it is not this one +- [TOKEN_BUCKET.md](./TOKEN_BUCKET.md) — the default in-process algorithm, and + how a window counter differs from it +- [MIGRATION.md](./MIGRATION.md) — upgrading; `Backend` is purely additive + +## Contract for implementers + +- **Safe for concurrent use.** +- **`Take` must be atomic per key.** Two concurrent calls for the same key must + not both succeed on the strength of the same remaining token. +- **Expire your own state.** A rate limiter's key space is usually unbounded — + one entry per client address — so a backend that never forgets is a leak. +- **Return an error rather than a guess.** An error means "I could not decide", + and `BackendLimiter` knows what to do with that. A backend that invents an + answer takes that choice away from the caller. +- `Remaining` may be `-1` for "unknown"; callers must not read a negative value + as zero. diff --git a/docs/CUSTOM_STORAGE.md b/docs/CUSTOM_STORAGE.md index 06e3961..9ee915a 100644 --- a/docs/CUSTOM_STORAGE.md +++ b/docs/CUSTOM_STORAGE.md @@ -177,6 +177,10 @@ func TestLoadOrStoreAtomic(t *testing.T) { ## Storage is in-process — read this before reaching for Redis +> **Looking for a limit shared across processes? You want +> [`Backend`](BACKENDS.md), not `Storage`.** This section explains why, and the +> rest of this page is about in-process containers. + It is tempting to implement `Storage` on top of Redis or Valkey to get a *distributed* limit shared across instances. **That does not work, and it is a subtle footgun.** Here is why: @@ -193,7 +197,7 @@ The rule of thumb: |------------------------------------------------------|----------------------------| | A different **in-process** store (LRU, metrics, …) | implement **`Storage`** | | A different **algorithm** (leaky bucket, GCRA, …) | implement **`Limiter`** | -| A **global** limit shared across instances | implement **`Limiter`** (backed by Redis/Valkey) | +| 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/MIGRATION.md b/docs/MIGRATION.md index d12be4a..9853b41 100644 --- a/docs/MIGRATION.md +++ b/docs/MIGRATION.md @@ -9,6 +9,15 @@ mechanically. > constructing an independent limiter per key, which is the root of most of the > API changes below. +> **Upgrading to the release that added `Backend`?** Nothing to do. That change +> is purely additive: `WithLimiterFactoryForKey`, `Backend`, `MemoryBackend` and +> `BackendLimiter` are new API, and every existing call — including +> `NewBucketLimiter(newLimiter, d, storage)` and every `Option` — keeps working +> unchanged. `Option` was deliberately **not** made generic for exactly this +> reason: `Option[K]` would have forced every existing `WithClock(now)` call to +> be explicitly instantiated. The guide below covers the earlier breaking +> release. + ## At a glance | Area | Before | After | diff --git a/docs/TOKEN_BUCKET.md b/docs/TOKEN_BUCKET.md index afceff0..802020e 100644 --- a/docs/TOKEN_BUCKET.md +++ b/docs/TOKEN_BUCKET.md @@ -148,9 +148,10 @@ The same bucket can be consumed three ways. `ratelimiter.Limiter` exposes If you decide not to proceed, call `reservation.Cancel()` to return the token. The default `RateLimiter` from `NewRateLimiterFunc` implements `Reserver` (it wraps `*rate.Limiter`, whose richer `ReserveN`/`DelayFrom` API is still - reachable through the embedded field); a custom backend such as Redis/Valkey - can implement `Reserver` too, so middleware stays backend-agnostic. See - [docs/CUSTOM_STORAGE.md](CUSTOM_STORAGE.md#distributed-limiting-with-redis--valkey). + reachable through the embedded field); `BackendLimiter` implements `Reserver` + too, so middleware stays backend-agnostic. Note that a window counter cannot + hand a token back in a way that is distinguishable from spending one less, so + its `Cancel()` is a documented no-op — see [docs/BACKENDS.md](BACKENDS.md). ## Choosing parameters @@ -307,15 +308,43 @@ fixed windows suffer from. ## Single-process vs. distributed -This library limits within a **single process**: the token state lives in memory -inside each `*rate.Limiter`. If you run N instances behind a load balancer, each -enforces the limit independently, so the effective global limit is up to N·`r`. - -For a *global* limit shared across instances you need a distributed algorithm -(commonly a sliding-window or token-bucket script in Redis, or a dedicated rate -limiting service). That is deliberately out of scope here — the `Storage` -interface exists to plug in custom **in-process** stores (e.g. a size-bounded -LRU), not to synchronize token state across machines. +The **default** limiter limits within a single process: the token state lives in +memory inside each `*rate.Limiter`. If you run N instances behind a load +balancer, each enforces the limit independently, so the effective global limit +is up to N·`r`. + +For a *global* limit shared across instances, implement a +[`Backend`](BACKENDS.md) — the count then lives wherever you put it, and every +process shares one budget. + +> **The `Storage` interface is not the way to do that.** It exists to plug in +> custom **in-process** containers (a size-bounded LRU, a metrics wrapper), not +> to synchronise token state across machines. `GetOrAdd` hands the caller the +> limiter and never writes it back, so a `Storage` that serialised state into a +> datastore would hand out a full bucket on every request and never limit +> anything. See [CUSTOM_STORAGE.md](CUSTOM_STORAGE.md). + +### The algorithm changes, and that is the honest part + +A distributed backend will usually be a **window counter**, because that is what +can be done in one atomic command. A token bucket's update is a +read-modify-write — read `(tokens, ts)`, refill, compare, write — which across a +network needs either a transaction with a retry loop on a contended key, or a +server-side script. + +So a shared limit does not behave identically to the in-process one: + +| | Token bucket (in-process) | Window counter (backend) | +| --- | --- | --- | +| Refill | continuous | stepwise | +| Burst | exactly `b` | approximately `b`; over-admits slightly at a boundary | +| `Retry-After` | exact, from the reservation | derived from the window edge | +| Cost | none | one round trip | + +`MemoryBackend` is deliberately a window counter for this reason: it is the +*same* algorithm a datastore backend runs, so switching between them does not +move the behaviour underneath you. When you want the best in-process limiter, +use `RateLimiter` (the token bucket) rather than `MemoryBackend`. ## Further reading diff --git a/examples/backend/main.go b/examples/backend/main.go new file mode 100644 index 0000000..2fb48a4 --- /dev/null +++ b/examples/backend/main.go @@ -0,0 +1,81 @@ +// Command backend demonstrates the Backend seam: the same limiter, backed by +// in-process state or by a shared store, chosen with one branch. +// +// The "remote" backend here is a map with an artificial failure switch rather +// than Valkey, so the example runs with no dependencies — but it exercises the +// three things BackendLimiter adds: the local fallback, the circuit breaker, +// and the degraded signal. +// +// go run ./examples/backend +package main + +import ( + "context" + "fmt" + "time" + + "golang.org/x/time/rate" + + "github.com/slashdevops/ratelimiter" +) + +// flakyBackend wraps a real backend and can be told to start failing, standing +// in for a datastore that becomes unreachable. +type flakyBackend struct { + inner ratelimiter.Backend + down bool + calls int +} + +func (b *flakyBackend) Take(ctx context.Context, key string, limit ratelimiter.Limit, cost int) (ratelimiter.Decision, error) { + b.calls++ + + if b.down { + return ratelimiter.Decision{}, fmt.Errorf("backend unreachable") + } + + return b.inner.Take(ctx, key, limit, cost) +} + +func main() { + limit := ratelimiter.Limit{Requests: 3, Period: time.Minute} + + shared := &flakyBackend{inner: ratelimiter.NewMemoryBackend()} + + // The local limiter is both the cache-disabled path and the fallback. + newLocal := func() ratelimiter.Limiter { + return ratelimiter.RateLimiter{Limiter: rate.NewLimiter(rate.Limit(limit.Rate()), limit.Capacity())} + } + + degraded := false + + bl := ratelimiter.NewBucketLimiter(nil, time.Minute, + ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), + ratelimiter.WithLimiterFactoryForKey(func(key string) ratelimiter.Limiter { + return ratelimiter.NewBackendLimiter(shared, key, limit, + ratelimiter.WithFallback(newLocal()), + ratelimiter.WithCircuitBreaker(2, time.Minute), + ratelimiter.WithOnDegraded(func(_ string, d bool) { degraded = d }), + ) + }), + ) + defer bl.Close() + + fmt.Println("shared backend healthy — the budget is 3 and it is global") + + for i := range 5 { + fmt.Printf(" request %d allow=%-5v backend_calls=%d\n", i+1, bl.GetOrAdd("alice").Allow(), shared.calls) + } + + fmt.Println("\nbackend goes down — the limiter falls back to the local budget") + shared.down = true + + for i := range 6 { + allow := bl.GetOrAdd("alice").Allow() + fmt.Printf(" request %d allow=%-5v backend_calls=%d degraded=%v\n", i+1, allow, shared.calls, degraded) + } + + fmt.Println("\nNote the backend call count stops climbing: after the threshold") + fmt.Println("the limiter answers locally with no network call at all, which is") + fmt.Println("what stops an outage adding a timeout to every request.") +} diff --git a/memory_backend.go b/memory_backend.go new file mode 100644 index 0000000..f65ee43 --- /dev/null +++ b/memory_backend.go @@ -0,0 +1,254 @@ +package ratelimiter + +import ( + "context" + "hash/maphash" + "sync" + "time" +) + +// memoryShards is how many independently-locked maps [MemoryBackend] spreads +// keys across. A rate limiter's hot path is one map operation per request, so a +// single mutex would serialise the whole server on it. 64 is enough that +// contention is negligible at any realistic core count while the memory +// overhead stays trivial. +const memoryShards = 64 + +// MemoryBackend is an in-process [Backend]: a sliding-window counter, sharded +// for concurrency, with its own expiry. +// +// # Why a sliding window and not a token bucket +// +// This package already has a token bucket — [RateLimiter], wrapping +// golang.org/x/time/rate — and it is the better in-process limiter: continuous +// refill, exact reservations, no window edges. MemoryBackend deliberately does +// NOT reproduce it. +// +// It exists so that the local and distributed paths agree. A Valkey or Redis +// backend will be a window counter, because that is what can be done in one +// atomic command; if the in-process backend were a token bucket, turning a +// cache on or off would silently change how traffic is admitted at a window +// boundary. Matching the algorithm makes [Backend] a seam you can move through +// without the behaviour moving under you. +// +// Use [RateLimiter] when you want the best in-process limiter. Use +// MemoryBackend when you want the same limiter you would get from a datastore, +// without the datastore. +// +// # The algorithm +// +// Two adjacent fixed windows, weighted by how far into the current one the +// clock is: +// +// weighted = previous*(1 - elapsed/period) + current +// +// A plain fixed window admits up to 2x the budget across a boundary — spend it +// all at the end of one window and again at the start of the next. The +// weighting reduces that to a small approximation error while still costing one +// counter per window. +// +// The zero value is not usable; construct one with [NewMemoryBackend]. +type MemoryBackend struct { + shards [memoryShards]memoryShard + seed maphash.Seed + now func() time.Time + + stop chan struct{} + done chan struct{} + closeOnce sync.Once +} + +type memoryShard struct { + mu sync.Mutex + entries map[string]*memoryEntry +} + +type memoryEntry struct { + window int64 // the current window index + current int + previous int + seen time.Time +} + +// MemoryBackendOption configures a [MemoryBackend]. +type MemoryBackendOption func(*memoryBackendConfig) + +type memoryBackendConfig struct { + now func() time.Time + sweepInterval time.Duration +} + +// WithMemoryClock overrides the time source, which makes window boundaries +// deterministic in tests. +func WithMemoryClock(now func() time.Time) MemoryBackendOption { + return func(c *memoryBackendConfig) { + if now != nil { + c.now = now + } + } +} + +// WithMemorySweepInterval overrides how often idle keys are evicted. Zero +// disables the background sweeper, leaving only the lazy expiry that happens +// when a key is next touched — which is enough for a bounded key space and a +// leak for an unbounded one. +func WithMemorySweepInterval(d time.Duration) MemoryBackendOption { + return func(c *memoryBackendConfig) { + c.sweepInterval = d + } +} + +// NewMemoryBackend returns a ready-to-use MemoryBackend and starts its sweeper. +// Call [MemoryBackend.Close] to stop it. +func NewMemoryBackend(opts ...MemoryBackendOption) *MemoryBackend { + cfg := memoryBackendConfig{now: time.Now, sweepInterval: time.Minute} + for _, opt := range opts { + opt(&cfg) + } + + b := &MemoryBackend{ + seed: maphash.MakeSeed(), + now: cfg.now, + stop: make(chan struct{}), + done: make(chan struct{}), + } + + for i := range b.shards { + b.shards[i].entries = make(map[string]*memoryEntry) + } + + if cfg.sweepInterval > 0 { + go b.sweepLoop(cfg.sweepInterval) + } else { + close(b.done) + } + + return b +} + +func (b *MemoryBackend) shard(key string) *memoryShard { + return &b.shards[maphash.String(b.seed, key)%memoryShards] +} + +// Take implements [Backend]. +func (b *MemoryBackend) Take(_ context.Context, key string, limit Limit, cost int) (Decision, error) { + if limit.Period <= 0 || limit.Requests <= 0 { + // A limit that allows nothing, or a nonsense window. Refusing is the + // safe reading: a misconfigured rule must not become an open door. + return Decision{Allowed: false, Remaining: 0}, nil + } + + if cost <= 0 { + cost = 1 + } + + now := b.now() + window := now.UnixNano() / int64(limit.Period) + elapsed := float64(now.UnixNano()%int64(limit.Period)) / float64(limit.Period) + + sh := b.shard(key) + + sh.mu.Lock() + defer sh.mu.Unlock() + + e, ok := sh.entries[key] + if !ok { + e = &memoryEntry{window: window} + sh.entries[key] = e + } + + // Roll the windows forward. A gap of two or more windows means everything + // known is stale, so both counters reset rather than one shifting into a + // slot it does not belong in. + switch delta := window - e.window; { + case delta == 1: + e.previous, e.current = e.current, 0 + case delta > 1: + e.previous, e.current = 0, 0 + } + + e.window = window + e.seen = now + + weighted := float64(e.previous)*(1-elapsed) + float64(e.current) + + if int(weighted)+cost > limit.Requests { + return Decision{ + Allowed: false, + Remaining: 0, + RetryAfter: time.Duration((1 - elapsed) * float64(limit.Period)), + }, nil + } + + e.current += cost + + remaining := limit.Requests - (int(weighted) + cost) + if remaining < 0 { + remaining = 0 + } + + return Decision{Allowed: true, Remaining: remaining}, nil +} + +// sweepLoop evicts keys nobody has touched for long enough that both their +// windows are stale. +func (b *MemoryBackend) sweepLoop(interval time.Duration) { + defer close(b.done) + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-b.stop: + return + case <-ticker.C: + b.sweep(interval) + } + } +} + +func (b *MemoryBackend) sweep(idleFor time.Duration) { + cutoff := b.now().Add(-2 * idleFor) + + for i := range b.shards { + sh := &b.shards[i] + + sh.mu.Lock() + + for key, e := range sh.entries { + if e.seen.Before(cutoff) { + delete(sh.entries, key) + } + } + + sh.mu.Unlock() + } +} + +// Len reports how many keys are currently held. It is intended for tests and +// metrics, and is a best-effort snapshot rather than a consistent one. +func (b *MemoryBackend) Len() int { + n := 0 + + for i := range b.shards { + sh := &b.shards[i] + + sh.mu.Lock() + n += len(sh.entries) + sh.mu.Unlock() + } + + return n +} + +// Close stops the background sweeper and waits for it to exit. It is safe to +// call more than once and from several goroutines. Take keeps working after +// Close; only eviction stops. +func (b *MemoryBackend) Close() { + b.closeOnce.Do(func() { close(b.stop) }) + <-b.done +} + +// compile-time check that the backend satisfies the port. +var _ Backend = (*MemoryBackend)(nil) diff --git a/memory_backend_test.go b/memory_backend_test.go new file mode 100644 index 0000000..4f8ff79 --- /dev/null +++ b/memory_backend_test.go @@ -0,0 +1,290 @@ +package ratelimiter + +import ( + "context" + "sync" + "testing" + "time" +) + +func fixedClock(t time.Time) (func() time.Time, func(time.Duration)) { + var mu sync.Mutex + + now := t + + return func() time.Time { + mu.Lock() + defer mu.Unlock() + + return now + }, func(d time.Duration) { + mu.Lock() + defer mu.Unlock() + + now = now.Add(d) + } +} + +func TestMemoryBackendAllowsExactlyTheBudget(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + b := NewMemoryBackend(WithMemoryClock(clock), WithMemorySweepInterval(0)) + defer b.Close() + + limit := Limit{Requests: 5, Period: time.Second} + + for i := range 5 { + d, err := b.Take(context.Background(), "k", limit, 1) + if err != nil { + t.Fatalf("Take %d: %v", i, err) + } + + if !d.Allowed { + t.Fatalf("request %d refused inside the budget", i+1) + } + + if want := 5 - (i + 1); d.Remaining != want { + t.Errorf("request %d: Remaining = %d, want %d", i+1, d.Remaining, want) + } + } + + d, _ := b.Take(context.Background(), "k", limit, 1) + if d.Allowed { + t.Error("the sixth request was allowed against a budget of five") + } + + if d.RetryAfter <= 0 { + t.Error("a refusal must say how long until it would be granted") + } +} + +func TestMemoryBackendKeysAreIndependent(t *testing.T) { + t.Parallel() + + clock, _ := fixedClock(time.Unix(0, 0)) + b := NewMemoryBackend(WithMemoryClock(clock), WithMemorySweepInterval(0)) + defer b.Close() + + limit := Limit{Requests: 1, Period: time.Second} + + if d, _ := b.Take(context.Background(), "alice", limit, 1); !d.Allowed { + t.Fatal("alice's first request should be allowed") + } + + if d, _ := b.Take(context.Background(), "bob", limit, 1); !d.Allowed { + t.Error("bob was refused because alice spent her budget; keys must not share one") + } +} + +// TestMemoryBackendSlidingWindowSmoothsTheBoundary is the reason this is not a +// plain fixed window. +// +// A fixed window admits the whole budget at the end of one window and the whole +// budget again at the start of the next — 2x the limit across the boundary, in +// an instant. The weighting carries the previous window's spend forward in +// proportion to how far into the new one the clock is, so the doubling cannot +// happen. +func TestMemoryBackendSlidingWindowSmoothsTheBoundary(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(100, 0)) + b := NewMemoryBackend(WithMemoryClock(clock), WithMemorySweepInterval(0)) + defer b.Close() + + limit := Limit{Requests: 10, Period: time.Second} + + // Spend the whole budget at the very end of a window. + advance(900 * time.Millisecond) + + for range 10 { + if d, _ := b.Take(context.Background(), "k", limit, 1); !d.Allowed { + t.Fatal("the budget should be spendable inside one window") + } + } + + // Step just over the boundary. A fixed window would hand back all ten. + advance(200 * time.Millisecond) + + allowed := 0 + + for range 10 { + if d, _ := b.Take(context.Background(), "k", limit, 1); d.Allowed { + allowed++ + } + } + + if allowed >= 10 { + t.Errorf("allowed %d immediately after the boundary; a fixed window would give 10, which is the 2x doubling this design exists to avoid", allowed) + } + + if allowed == 0 { + t.Error("allowed 0 after the boundary; the window should have released some budget") + } +} + +func TestMemoryBackendResetsAfterAFullGap(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := NewMemoryBackend(WithMemoryClock(clock), WithMemorySweepInterval(0)) + defer b.Close() + + limit := Limit{Requests: 1, Period: time.Second} + + if d, _ := b.Take(context.Background(), "k", limit, 1); !d.Allowed { + t.Fatal("the first request should be allowed") + } + + // Two whole windows later nothing is carried forward: a counter shifted + // into a slot it does not belong in would refuse this. + advance(3 * time.Second) + + if d, _ := b.Take(context.Background(), "k", limit, 1); !d.Allowed { + t.Error("a request three windows later was refused; stale counters must be dropped, not shifted") + } +} + +func TestMemoryBackendRefusesANonsenseLimit(t *testing.T) { + t.Parallel() + + b := NewMemoryBackend(WithMemorySweepInterval(0)) + defer b.Close() + + for name, limit := range map[string]Limit{ + "zero period": {Requests: 10, Period: 0}, + "zero requests": {Requests: 0, Period: time.Second}, + "negative": {Requests: -1, Period: time.Second}, + } { + t.Run(name, func(t *testing.T) { + // A misconfigured rule must not become an open door. Refusing is + // visible; allowing everything looks exactly like working. + if d, _ := b.Take(context.Background(), "k", limit, 1); d.Allowed { + t.Error("a limit that allows nothing allowed a request") + } + }) + } +} + +func TestMemoryBackendEvictsIdleKeys(t *testing.T) { + t.Parallel() + + clock, advance := fixedClock(time.Unix(0, 0)) + b := NewMemoryBackend(WithMemoryClock(clock), WithMemorySweepInterval(0)) + defer b.Close() + + limit := Limit{Requests: 10, Period: time.Second} + + for _, k := range []string{"a", "b", "c"} { + if _, err := b.Take(context.Background(), k, limit, 1); err != nil { + t.Fatal(err) + } + } + + if got := b.Len(); got != 3 { + t.Fatalf("Len = %d, want 3", got) + } + + advance(time.Hour) + b.sweep(time.Minute) + + if got := b.Len(); got != 0 { + t.Errorf("Len = %d after the sweep, want 0; an unbounded key space that is never evicted is a leak", got) + } +} + +func TestMemoryBackendIsSafeUnderConcurrency(t *testing.T) { + t.Parallel() + + b := NewMemoryBackend(WithMemorySweepInterval(0)) + defer b.Close() + + // A budget far larger than the number of requests, so every one should be + // allowed and any lost update shows up as a refusal. + limit := Limit{Requests: 10_000, Period: time.Hour} + + var ( + wg sync.WaitGroup + mu sync.Mutex + refused int + ) + + for range 200 { + wg.Go(func() { + d, err := b.Take(context.Background(), "hot", limit, 1) + if err != nil { + t.Error(err) + } + + if !d.Allowed { + mu.Lock() + refused++ + mu.Unlock() + } + }) + } + + wg.Wait() + + if refused != 0 { + t.Errorf("%d of 200 concurrent requests were refused against a 10000 budget", refused) + } +} + +// TestMemoryBackendSweeperRuns covers the background goroutine rather than only +// the sweep it calls: a sweeper that is never started is a leak that no unit +// test of sweep() would notice. +func TestMemoryBackendSweeperRuns(t *testing.T) { + t.Parallel() + + b := NewMemoryBackend(WithMemorySweepInterval(5 * time.Millisecond)) + defer b.Close() + + limit := Limit{Requests: 10, Period: time.Millisecond} + + if _, err := b.Take(context.Background(), "k", limit, 1); err != nil { + t.Fatal(err) + } + + deadline := time.Now().Add(2 * time.Second) + for b.Len() > 0 { + if time.Now().After(deadline) { + t.Fatal("the key was never evicted; the background sweeper is not running") + } + + time.Sleep(2 * time.Millisecond) + } +} + +func TestMemoryBackendCloseIsIdempotent(t *testing.T) { + t.Parallel() + + b := NewMemoryBackend(WithMemorySweepInterval(time.Hour)) + b.Close() + b.Close() // must not panic or block + + // Take keeps working after Close; only eviction stops. + if _, err := b.Take(context.Background(), "k", + Limit{Requests: 1, Period: time.Second}, 1); err != nil { + t.Errorf("Take after Close: %v", err) + } +} + +func TestMemoryBackendNormalisesCost(t *testing.T) { + t.Parallel() + + b := NewMemoryBackend(WithMemorySweepInterval(0)) + defer b.Close() + + limit := Limit{Requests: 1, Period: time.Minute} + + // A zero or negative cost is treated as one rather than as free, so a + // caller cannot spend nothing forever. + if d, _ := b.Take(context.Background(), "k", limit, 0); !d.Allowed { + t.Fatal("the first request should be allowed") + } + + if d, _ := b.Take(context.Background(), "k", limit, -5); d.Allowed { + t.Error("a negative cost was treated as free; it must count as one") + } +} From d227d1009d08eea571d6ba9639b91e0ed9f24c82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Gonz=C3=A1lez=20Di=20Antonio?= Date: Tue, 25 Aug 2026 19:45:58 +0200 Subject: [PATCH 2/3] chore: build on Go 1.27 and apply its modernisers go.mod moves from 1.26.0 to 1.27.0. CI resolves its toolchain from `go-version-file: ./go.mod`, so this is what the build actually runs on, and it matches the toolchain the consumers of this library are already using. golang.org/x/time is already at v0.15.0, the latest; `go get -u ./...` and `go mod tidy` produce no change. The dependency list stays one entry, which is worth keeping that way. `go fix ./...` had one thing to say, in code added by the previous commit: a three-line clamp collapsed to `max(...)`. That branch was the last statement in the package no test reached, so coverage goes back to 100.0% -- the level this repo was at before the Backend work, and the reason to check rather than assume it was still true. NOTE FOR REVIEWERS: raising the `go` directive raises the minimum Go version for everyone who imports this library. That is a deliberate choice rather than a side effect of the tooling, and it is the one part of this change that is not purely additive. --- go.mod | 2 +- memory_backend.go | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/go.mod b/go.mod index ebf5f3f..ffdf017 100644 --- a/go.mod +++ b/go.mod @@ -1,5 +1,5 @@ module github.com/slashdevops/ratelimiter -go 1.26.0 +go 1.27.0 require golang.org/x/time v0.15.0 diff --git a/memory_backend.go b/memory_backend.go index f65ee43..4b50b8a 100644 --- a/memory_backend.go +++ b/memory_backend.go @@ -182,10 +182,7 @@ func (b *MemoryBackend) Take(_ context.Context, key string, limit Limit, cost in e.current += cost - remaining := limit.Requests - (int(weighted) + cost) - if remaining < 0 { - remaining = 0 - } + remaining := max(limit.Requests-(int(weighted)+cost), 0) return Decision{Allowed: true, Remaining: remaining}, nil } From d73ad37b99fc572566b6746150057836841c7ab4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Gonz=C3=A1lez=20Di=20Antonio?= Date: Tue, 25 Aug 2026 19:52:46 +0200 Subject: [PATCH 3/3] fix(ci): keep go.mod on 1.26, and pin Go in the CodeQL workflow Two changes, one of them undoing a mistake I made in the previous commit. REVERT THE go DIRECTIVE to 1.26.0. Bumping it to 1.27.0 broke CodeQL: go: go.mod requires go >= 1.27.0 (running go 1.26.6; GOTOOLCHAIN=local) Extraction failed for all discovered Go projects. The autobuilder runs whatever Go the runner image ships with GOTOOLCHAIN=local, so it cannot fetch a newer toolchain and the analysis simply fails. But the failure is the smaller half of the argument. The bump was justified as "this is what CI builds with, and it matches the consumers" -- and CodeQL is a counter-example to the first half, while the second is not a reason that binds anybody else. Nothing in this package needs 1.27: `go vet` with the directive at 1.26 reports no use of a newer standard-library symbol, and the full suite passes at 100% coverage. Raising the minimum Go version of a PUBLIC library excludes every consumer still on the previous release, and Go supports the two most recent majors. Doing that for alignment rather than for a feature is a cost paid by other people for our convenience. The `max(...)` rewrite `go fix` produced stays: it predates 1.27 and it is why coverage reached 100%. PIN GO IN THE CODEQL WORKFLOW anyway, from go.mod, before codeql-action/init. The workflow is fragile independently of this revert: it analyses with whatever Go the runner happens to have that week, so the day the directive legitimately moves, this breaks again. And it breaks illegibly -- the reported error is "We were unable to automatically build your code", which points nowhere near a version mismatch. This is the same class of problem as the pr.yaml coverage step noted in the PR: a workflow that depends on an unstated assumption about the runner. Verified at 1.26.0: go vet clean (including the stdversion check, which is what proves no 1.27-only symbol is used), go fix -diff clean, race tests green, coverage 100.0%. --- .github/workflows/codeql.yml | 23 ++++++++++++++++++----- go.mod | 2 +- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index a67e94e..659fd6b 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -59,11 +59,24 @@ jobs: - name: Checkout repository uses: actions/checkout@v7 - # Add any setup steps before running the `github/codeql-action/init` action. - # This includes steps like installing compilers or runtimes (`actions/setup-node` - # or others). This is typically only required for manual builds. - # - name: Setup runtime (example) - # uses: actions/setup-example@v1 + # Install the Go version this module declares, BEFORE codeql-action/init. + # + # Without this the autobuilder uses whatever Go the runner happens to + # ship, with GOTOOLCHAIN=local — so it cannot download a newer one, and + # the whole analysis fails the moment go.mod's `go` directive moves ahead + # of the image: + # + # go: go.mod requires go >= 1.27.0 (running go 1.26.6; GOTOOLCHAIN=local) + # Extraction failed for all discovered Go projects. + # + # That is a failure of the workflow rather than of the code, and it + # reports as "We were unable to automatically build your code", which + # points nowhere near the cause. Pinning from go.mod means this job builds + # what the project declares, whatever the runner image is that week. + - name: Setup Go + uses: actions/setup-go@v6 + with: + go-version-file: ./go.mod # Initializes the CodeQL tools for scanning. - name: Initialize CodeQL diff --git a/go.mod b/go.mod index ffdf017..ebf5f3f 100644 --- a/go.mod +++ b/go.mod @@ -1,5 +1,5 @@ module github.com/slashdevops/ratelimiter -go 1.27.0 +go 1.26.0 require golang.org/x/time v0.15.0