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/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..4b50b8a --- /dev/null +++ b/memory_backend.go @@ -0,0 +1,251 @@ +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 := max(limit.Requests-(int(weighted)+cost), 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") + } +}