diff --git a/README.md b/README.md index 521f1bd..bf811a6 100644 --- a/README.md +++ b/README.md @@ -37,9 +37,10 @@ idle for a configurable duration. - **Type-safe & generic** — `Storage[K, V]` and `BucketLimiter[K]` use Go generics (Go 1.26+). - **Extensible** — implement the `Storage` interface for a custom in-process - store, or the `Limiter` interface for a custom algorithm. + store, or the `Limiter` interface for a custom algorithm (optionally adding + `Reserver` for accurate `Retry-After`, even with a Redis/Valkey backend). - **HTTP middleware example** — with accurate `Retry-After` and `RateLimit-*` - response headers. + response headers, driven by the backend-agnostic `Reserver` interface. ## Architecture @@ -155,10 +156,12 @@ func main() { | Type / func | Role | |------------------------------------|----------------------------------------------------------------------------| | `Limiter` | Minimal interface (`Allow`, `Wait`, `Burst`). `*rate.Limiter` satisfies it. | +| `Reserver` / `Reservation` | Optional capability: reserve a token and read its delay. Enables accurate `Retry-After` for any backend. | +| `RateLimiter` | Default limiter: wraps `*rate.Limiter`, implements `Limiter` **and** `Reserver`. | | `Storage[K, V]` | Pluggable, concurrency-safe store for per-key limiters. | | `InMemoryStorage[K, V]` | Default `sync.Map`-backed store. | | `BucketLimiter[K]` | Manager: hands out one `Limiter` per key, handles creation and eviction. | -| `NewRateLimiterFunc(limit, burst)` | Convenience factory for the common `*rate.Limiter` case. | +| `NewRateLimiterFunc(limit, burst)` | Convenience factory producing `RateLimiter` values. | ### `limit` and `burst` @@ -186,6 +189,33 @@ if err := lim.Wait(ctx); err != nil { } ``` +### Reserve (for accurate `Retry-After`) + +When you need the exact delay until the next token — to set a `Retry-After` or +`RateLimit-Reset` header — use the optional `Reserver` capability. Feature-detect +it so your code works with any limiter and degrades gracefully: + +```go +lim := manager.GetOrAdd(key) + +if r, ok := lim.(ratelimiter.Reserver); ok { + res := r.Reserve() + if res.OK() && res.Delay() == 0 { + // proceed now + } else { + res.Cancel() // return the token + retryAfter := res.Delay() // tell the client exactly how long to wait + } +} else { + _ = lim.Allow() // limiter without reservation support: no timing info +} +``` + +The default `RateLimiter` from `NewRateLimiterFunc` implements `Reserver`, and a +custom (e.g. Redis/Valkey-backed) `Limiter` can too — so the same middleware +produces accurate headers regardless of backend. See +[docs/CUSTOM_STORAGE.md](docs/CUSTOM_STORAGE.md#distributed-limiting-with-redis--valkey). + ## HTTP middleware The [`examples/middleware`](examples/middleware/main.go) program limits requests @@ -211,10 +241,12 @@ go run ./examples/key -limit 1 -burst 3 ## Custom storage -Implement `Storage[K, V]` to back the manager with your own in-process store — -for example a size-bounded LRU to cap memory instead of (or in addition to) -time-based eviction. Implementations must be safe for concurrent use, and -`LoadOrStore` must be atomic. +`BucketLimiter` talks to the `Storage[K, V]` interface, never a concrete map, +and you inject the implementation at construction time. `InMemoryStorage` is +just the bundled **default** — implement the interface to bring your own +in-process store (for example a size-bounded LRU to cap memory instead of, or in +addition to, time-based eviction). Implementations must be safe for concurrent +use, and `LoadOrStore` must be atomic. ```go type Storage[K comparable, V any] interface { @@ -226,6 +258,19 @@ type Storage[K comparable, V any] interface { } ``` +A complete, runnable size-bounded LRU store lives in +[`examples/customstorage`](examples/customstorage/main.go): + +```bash +go run ./examples/customstorage -cap 2 +``` + +**[docs/CUSTOM_STORAGE.md](docs/CUSTOM_STORAGE.md)** is a full guide: the method +contracts, how to test atomicity, and — importantly — why a custom `Storage` is +**in-process only**, plus the correct pattern for **distributed limiting with +Redis / [Valkey](https://github.com/valkey-io/valkey-go)** (a datastore-backed +`Limiter` wired through a `Storage` resolver, with the token-bucket Lua script). + ## Scope: single-process only Token state lives in memory inside each `*rate.Limiter`, so this library diff --git a/doc.go b/doc.go index f95fd8f..70a1dce 100644 --- a/doc.go +++ b/doc.go @@ -22,9 +22,17 @@ // // allowed // } // +// Limiters are consumed through the [Limiter] interface (Allow, Wait, Burst). +// A limiter may optionally also implement [Reserver] to reserve a token and +// report the exact delay until it is valid; the default limiter from +// [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. 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 and is out of scope. +// 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. package ratelimiter diff --git a/docs/CUSTOM_STORAGE.md b/docs/CUSTOM_STORAGE.md new file mode 100644 index 0000000..4425e44 --- /dev/null +++ b/docs/CUSTOM_STORAGE.md @@ -0,0 +1,472 @@ +# Implementing a custom `Storage` + +`BucketLimiter` never talks to a concrete map type. It talks to the +[`Storage[K, V]`](../storage.go) interface, and you inject the implementation at +construction time: + +```go +manager := ratelimiter.NewBucketLimiter(newLimiter, time.Minute, storage) +// ▲ +// any Storage[K, Limiter] you like +``` + +The bundled [`InMemoryStorage`](../memory_store.go) is just the **default** +implementation, offered for convenience. Anyone can supply their own by +implementing the five-method interface — and if they don't, they use the default. + +- [When to implement your own](#when-to-implement-your-own) +- [The interface and its contract](#the-interface-and-its-contract) +- [Where the store sits](#where-the-store-sits) +- [A complete in-process example: bounded LRU](#a-complete-in-process-example-bounded-lru) +- [Testing your implementation](#testing-your-implementation) +- [Storage is in-process — read this before reaching for Redis](#storage-is-in-process--read-this-before-reaching-for-redis) +- [Distributed limiting with Redis / Valkey](#distributed-limiting-with-redis--valkey) + +## When to implement your own + +Reach for a custom `Storage` when the default `sync.Map` doesn't fit how you want +to **hold the per-key limiters in memory**. Common reasons: + +- **Bound memory for an unbounded key space** — cap the number of live keys with + an LRU or a fixed-size ring, so per-IP limiting exposed to the internet cannot + grow without limit. (This is the [example below](#a-complete-in-process-example-bounded-lru).) +- **Instrument the store** — count hits/misses, export metrics, or log key churn. +- **Change the backing structure** — a sharded map to cut lock contention, a + `weak`-pointer store, an arena, etc. +- **Bind per-key metadata** — use the store as the seam that constructs + key-aware limiters (this is what the [Valkey pattern](#distributed-limiting-with-redis--valkey) + relies on). + +If none of those apply, use `NewInMemoryStorage` — it is correct, atomic, and +fast. + +## The interface and its contract + +```go +type Storage[K comparable, V any] interface { + Store(key K, value V) + Load(key K) (value V, ok bool) + LoadOrStore(key K, value V) (actual V, loaded bool) + Delete(key K) + Range(f func(key K, value V) bool) +} +``` + +Every implementation **MUST be safe for concurrent use by multiple goroutines.** +Beyond that: + +| Method | Contract | Called by the manager? | +|---------------|---------------------------------------------------------------------------------------------------|-------------------------------------------------| +| `Store` | Unconditionally set `value` for `key`. | No — provided for completeness / your own use. | +| `Load` | Return the stored value and whether it was present. | **Yes** — first thing `GetOrAdd` does. | +| `LoadOrStore` | If `key` exists, return the existing value; else store and return the given one. **Must be atomic.** | **Yes** — the create-on-miss path in `GetOrAdd`. | +| `Delete` | Remove `key`; deleting an absent key is a no-op. | **Yes** — from `Remove` and idle eviction. | +| `Range` | Call `f` for each entry until `f` returns `false`. Best-effort snapshot. | No — the manager tracks idle timers separately. | + +Two consequences worth internalizing: + +1. **`LoadOrStore` atomicity is load-bearing.** `GetOrAdd` relies on it so that + two goroutines racing on a brand-new key both receive the *same* limiter + instance — otherwise one client could momentarily get two independent buckets. + Guard it with a mutex or use an atomic primitive (`sync.Map.LoadOrStore`). +2. **The manager only calls `Load`, `LoadOrStore`, and `Delete`.** It keeps its + own last-use map for idle eviction, so it never calls your `Range`, and it + never calls `Store`. You still must implement all five to satisfy the type, + but only those three are on the hot path. + +## Where the store sits + +```mermaid +flowchart LR + App["Your code"] -->|"GetOrAdd(key)"| BL["BucketLimiter[K]"] + BL -->|"Load / LoadOrStore / Delete"| S{{"Storage[K, Limiter]
(interface)"}} + S -.->|"default"| IMS["InMemoryStorage
(sync.Map)"] + S -.->|"or your own"| Custom["LRUStorage · MetricsStorage
ValkeyStore · …"] + BL -->|"newLimiter() on miss"| F["factory func() Limiter"] + F -->|"fresh Limiter"| S +``` + +`BucketLimiter` is storage-agnostic: swap the box on the right and everything +else stays the same. + +## A complete in-process example: bounded LRU + +[`examples/customstorage`](../examples/customstorage/main.go) is a runnable, +size-bounded LRU store. It keeps at most `cap` distinct keys and evicts the +least-recently-used one on overflow — capping memory without relying on the +manager's time-based eviction. + +```bash +go run ./examples/customstorage -cap 2 +``` + +``` +store capacity = 2 + +request alice -> allowed=true +request bob -> allowed=true + [LRUStorage] capacity 2 exceeded, evicted alice +request carol -> allowed=true + [LRUStorage] capacity 2 exceeded, evicted bob +request alice -> allowed=true +``` + +The shape of any in-process implementation is the same — a guarded map plus +whatever policy you want: + +```go +type LRUStorage[K comparable, V any] struct { + mu sync.Mutex + capacity int + ll *list.List // front = most-recently-used + items map[K]*list.Element +} + +func (s *LRUStorage[K, V]) LoadOrStore(key K, value V) (V, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if el, ok := s.items[key]; ok { // already present → return it (atomic) + s.ll.MoveToFront(el) + return el.Value.(*lruEntry[K, V]).value, true + } + s.set(key, value) // insert; evict LRU if over capacity + return value, false +} +// ... Store, Load, Delete, Range — see the example file for the full type. +``` + +> **Heads-up on eviction interaction.** If your store evicts a key that the +> manager still considers active, the next `GetOrAdd` simply rebuilds a fresh, +> full bucket via the factory — identical to what idle eviction does. That is +> usually fine, but if you want *only* size-based bounding, construct the manager +> with `deleteAfter <= 0` to turn off the time-based sweeper (as the example +> does). + +## Testing your implementation + +Two properties are worth an explicit test: + +```go +// 1. LoadOrStore is atomic: N goroutines racing on one new key all get the +// same instance. +func TestLoadOrStoreAtomic(t *testing.T) { + s := NewMyStore[string, int]() + const g = 100 + var wg sync.WaitGroup + got := make([]int, g) + for i := range g { + wg.Add(1) + go func(i int) { + defer wg.Done() + v, _ := s.LoadOrStore("k", i) // each offers a different value + got[i] = v + }(i) + } + wg.Wait() + for _, v := range got { + if v != got[0] { + t.Fatalf("LoadOrStore not atomic: saw %d and %d", v, got[0]) + } + } +} + +// 2. Run the race detector against a Load/Store/Delete mix: +// go test -race ./... +``` + +## Storage is in-process — read this before reaching for Redis + +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: + +- What a `Storage` holds is the **`Limiter` object itself** — a live Go value + (`*rate.Limiter`) whose token count lives in its own memory. +- Serializing that object into Redis and back gives every process (and every + reload) its **own** counter. Nothing is shared. You would pay network latency + for zero coordination. + +The rule of thumb: + +| You want… | Extension point | +|------------------------------------------------------|----------------------------| +| A different **in-process** store (LRU, metrics, …) | implement **`Storage`** | +| A different **algorithm** (leaky bucket, GCRA, …) | implement **`Limiter`** | +| A **global** limit shared across instances | implement **`Limiter`** (backed by Redis/Valkey) | + +Distributed limiting is a **`Limiter`** concern, not a `Storage` concern. + +## Distributed limiting with Redis / Valkey + +To share one budget across every instance, the *token math must run where the +state lives* — in the datastore, via an atomic server-side script. So you +implement the [`Limiter`](../limiter.go) interface (`Allow`, `Wait`, `Burst`) +backed by [valkey-go](https://github.com/valkey-io/valkey-go), and you use a +custom `Storage` as the **key→limiter resolver** so each limiter knows *its* +Valkey key. + +> **Design note.** The `newLimiter func() Limiter` factory does not receive the +> key, and `Allow()`/`Wait()` take no key argument — so one `Limiter` instance +> represents exactly one key's bucket. The only injection point that sees both +> the **key** and a place to hold a **shared client** is `Storage.LoadOrStore`. +> That makes a custom `Storage` the natural seam: it caches one lightweight, +> key-bound Valkey limiter per key. This uses the store as a factory/cache +> rather than as a value container — a deliberate, supported reinterpretation +> for this use case. (If you would rather not reuse `BucketLimiter` at all, a +> ~30-line standalone manager over the same `valkeyLimiter` works too and keeps +> `Storage` out of it entirely.) + +### The flow + +```mermaid +sequenceDiagram + autonumber + participant App as Your code + participant BL as BucketLimiter + participant St as ValkeyStore (resolver) + participant Lim as valkeyLimiter (per key) + participant VK as Valkey / Redis + + App->>BL: GetOrAdd(key).Allow() + BL->>St: Load(key) / LoadOrStore(key, _) + St-->>BL: valkeyLimiter bound to key
(cached in-process) + BL-->>App: limiter + App->>Lim: Allow() + Lim->>VK: EVAL token-bucket.lua
KEYS=[rl:key] ARGV=[rate,burst,now,1] + Note over VK: refill + try-consume
run atomically server-side + VK-->>Lim: {allowed, retry_after_ms} + Lim-->>App: true / false +``` + +Every instance shares the same Valkey key, so the limit is now **global**. + +### The token-bucket script + +This Lua runs atomically inside Valkey. It refills the bucket from elapsed time, +then tries to take one token. (Valkey keeps the `redis.*` scripting API for +compatibility.) + +```lua +-- KEYS[1] = bucket key +-- ARGV[1] = rate (tokens per second) +-- ARGV[2] = burst (bucket capacity) +-- ARGV[3] = now_ms (caller's clock, milliseconds) +-- ARGV[4] = cost (tokens requested, usually 1) +-- returns { allowed (1/0), retry_after_ms } +local rate = tonumber(ARGV[1]) +local burst = tonumber(ARGV[2]) +local now = tonumber(ARGV[3]) +local cost = tonumber(ARGV[4]) + +local state = redis.call('HMGET', KEYS[1], 'tokens', 'ts') +local tokens = tonumber(state[1]) +local ts = tonumber(state[2]) +if tokens == nil then + tokens = burst + ts = now +end + +-- continuous refill since the last update +tokens = math.min(burst, tokens + math.max(0, now - ts) / 1000.0 * rate) + +local allowed, retry = 0, 0 +if tokens >= cost then + allowed = 1 + tokens = tokens - cost +else + retry = math.ceil((cost - tokens) / rate * 1000.0) +end + +redis.call('HSET', KEYS[1], 'tokens', tokens, 'ts', now) +-- let idle buckets expire once they would be full again +redis.call('PEXPIRE', KEYS[1], math.ceil(burst / rate * 1000.0) + 1000) +return { allowed, retry } +``` + +### The `Limiter` implementation + +> The snippet below targets `github.com/valkey-io/valkey-go`. It is **illustrative** — +> it is intentionally *not* compiled into this module (the library has no Valkey +> dependency). Drop it into your own package, add `valkey-go` to your `go.mod`, +> and test it against a real server before relying on it. + +```go +package valkeylimiter + +import ( + "context" + "strconv" + "time" + + "github.com/slashdevops/ratelimiter" + "github.com/valkey-io/valkey-go" +) + +var script = valkey.NewLuaScript(tokenBucketLua) // the Lua from above + +// valkeyLimiter implements ratelimiter.Limiter against a shared Valkey bucket. +type valkeyLimiter struct { + client valkey.Client + key string // the Valkey key, e.g. "rl:user-123" + rate float64 // tokens per second + burst int +} + +func (l *valkeyLimiter) Burst() int { return l.burst } + +func (l *valkeyLimiter) Allow() bool { + allowed, _, err := l.eval(context.Background(), 1) + // Fail-open or fail-closed on transport errors is your policy choice. + return err == nil && allowed +} + +func (l *valkeyLimiter) Wait(ctx context.Context) error { + for { + allowed, retry, err := l.eval(ctx, 1) + if err != nil { + return err + } + if allowed { + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(retry): + } + } +} + +func (l *valkeyLimiter) eval(ctx context.Context, cost int) (allowed bool, retry time.Duration, err error) { + now := strconv.FormatInt(time.Now().UnixMilli(), 10) + res := script.Exec(ctx, l.client, []string{l.key}, []string{ + strconv.FormatFloat(l.rate, 'f', -1, 64), + strconv.Itoa(l.burst), + now, + strconv.Itoa(cost), + }) + arr, err := res.ToArray() + if err != nil { + return false, 0, err + } + ok, _ := arr[0].ToInt64() + ms, _ := arr[1].ToInt64() + return ok == 1, time.Duration(ms) * time.Millisecond, nil +} + +// Compile-time check. +var _ ratelimiter.Limiter = (*valkeyLimiter)(nil) +``` + +### Accurate `Retry-After`: implement `Reserver` too + +The [`Reserver`](../reservation.go) interface lets HTTP middleware read the exact +delay until the next token and emit accurate `Retry-After` / `RateLimit-Reset` +headers **for any backend** — the [middleware example](../examples/middleware/main.go) +feature-detects it. Your Lua script already returns `retry_after_ms`, so exposing +it is a few lines and makes the Valkey limiter a first-class citizen of that +middleware (instead of degrading to `Allow`-only): + +```go +type valkeyReservation struct { + limiter *valkeyLimiter + ok bool + delay time.Duration +} + +func (r valkeyReservation) OK() bool { return r.ok } +func (r valkeyReservation) Delay() time.Duration { return r.delay } +func (r valkeyReservation) Cancel() { /* optional: return the token via a compensating EVAL */ } + +func (l *valkeyLimiter) Reserve() ratelimiter.Reservation { + allowed, retry, err := l.eval(context.Background(), 1) + if err != nil { + return valkeyReservation{limiter: l, ok: false} + } + // allowed -> delay 0; otherwise the caller waits retry before the token is valid. + return valkeyReservation{limiter: l, ok: true, delay: retry} +} + +var _ ratelimiter.Reserver = (*valkeyLimiter)(nil) +``` + +> `Cancel` is best-effort for a distributed limiter: returning a token means an +> extra round-trip (a compensating script that credits one token back). If you +> reserve-then-reject on every over-limit request, a cheap option is to not +> pre-consume in `Reserve` at all — compute the delay with a read-only variant of +> the script and only consume in `Allow`. Choose the trade-off that fits your +> accuracy vs. round-trip budget. + +### Wiring it through a `Storage` resolver + +The store injects the shared client and binds the key. It ignores the +in-memory value the manager passes (that throwaway `*rate.Limiter`) and returns +the Valkey-backed limiter instead — caching one per key: + +```go +type ValkeyStore struct { + client valkey.Client + rate float64 + burst int + + mu sync.Mutex + cache map[string]ratelimiter.Limiter +} + +func NewValkeyStore(c valkey.Client, r float64, burst int) *ValkeyStore { + return &ValkeyStore{client: c, rate: r, burst: burst, cache: map[string]ratelimiter.Limiter{}} +} + +func (s *ValkeyStore) LoadOrStore(key string, _ ratelimiter.Limiter) (ratelimiter.Limiter, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if l, ok := s.cache[key]; ok { + return l, true + } + l := &valkeyLimiter{client: s.client, key: "rl:" + key, rate: s.rate, burst: s.burst} + s.cache[key] = l + return l, false +} + +func (s *ValkeyStore) Load(key string) (ratelimiter.Limiter, bool) { + s.mu.Lock() + defer s.mu.Unlock() + l, ok := s.cache[key] + return l, ok +} + +func (s *ValkeyStore) Store(key string, v ratelimiter.Limiter) { /* mu + s.cache[key] = v */ } +func (s *ValkeyStore) Delete(key string) { /* mu + delete(s.cache, key) */ } +func (s *ValkeyStore) Range(f func(string, ratelimiter.Limiter) bool) { /* mu + iterate */ } +``` + +Then assemble as usual. Because the store supplies the real limiter, the +`newLimiter` factory is never meaningfully used — pass a trivial one: + +```go +client, _ := valkey.NewClient(valkey.ClientOption{InitAddress: []string{"127.0.0.1:6379"}}) +defer client.Close() + +store := NewValkeyStore(client, 5, 10) // 5 rps, burst 10 — shared across instances +noop := func() ratelimiter.Limiter { return nil } // ignored by ValkeyStore.LoadOrStore + +manager := ratelimiter.NewBucketLimiter(noop, time.Minute, store) +defer manager.Close() + +if manager.GetOrAdd("user-123").Allow() { // consults Valkey atomically + // handle request +} +``` + +Now `user-123` draws from **one** bucket no matter which instance serves the +request. The manager's idle eviction still runs — it just evicts the cheap +in-process `valkeyLimiter` wrapper from `ValkeyStore`, while the authoritative +state (and its own TTL) lives in Valkey. + +### Redis instead of Valkey + +The same design works with any Redis client (`go-redis/redis`, +`redis/rueidis`, …): the Lua script is unchanged, only the `Exec`/`EVAL` call +differs. If you would rather not hand-roll the script at all, a purpose-built +library such as [`go-redis/redis_rate`](https://github.com/go-redis/redis_rate) +implements a GCRA limiter over Redis and may be the faster path to production. diff --git a/docs/MIGRATION.md b/docs/MIGRATION.md index 808abc4..d12be4a 100644 --- a/docs/MIGRATION.md +++ b/docs/MIGRATION.md @@ -21,6 +21,7 @@ mechanically. | Calling on the manager | `manager.Allow()` / `.Wait()` / `.Burst()` | removed — use `manager.GetOrAdd(key).Allow()` etc. | | Manual removal | `err := manager.Remove(key)` | `manager.Remove(key)` *(no return value)* | | `Storage` interface | non-generic, `Store/Load/Delete` | generic `Storage[K, V]`, adds `LoadOrStore` and `Range` | +| Default limiter type | `NewRateLimiterFunc` builds a bare `*rate.Limiter` | builds a `RateLimiter` (wraps `*rate.Limiter`, adds `Reserver`) | ## 1. `NewInMemoryStorage` now requires type parameters @@ -122,6 +123,38 @@ type Storage[K comparable, V any] interface { - `Range` is used by the manager only indirectly; implement it to iterate the current entries (return `false` from `f` to stop early). +## 7. `NewRateLimiterFunc` returns `RateLimiter`, not `*rate.Limiter` + +The factory now hands out a `ratelimiter.RateLimiter` — a thin wrapper that +embeds `*rate.Limiter` and additionally implements the new +[`Reserver`](../reservation.go) interface. This is what lets middleware compute +an accurate `Retry-After` for any backend without importing `golang.org/x/time/rate`. + +Almost all code is unaffected: `Allow`, `Wait`, and `Burst` are promoted from the +embedded `*rate.Limiter`, so `manager.GetOrAdd(key).Allow()` is unchanged. The +only thing that breaks is a **type assertion to the concrete `*rate.Limiter`**: + +```go +// Before — downcast to reach Reserve/ReserveN: +if rl, ok := manager.GetOrAdd(key).(*rate.Limiter); ok { + res := rl.ReserveN(time.Now(), 1) + // ... +} + +// After — feature-detect the backend-agnostic Reserver instead: +if r, ok := manager.GetOrAdd(key).(ratelimiter.Reserver); ok { + res := r.Reserve() + _ = res.Delay() +} + +// Or, if you specifically need *rate.Limiter's richer API, assert the wrapper +// and reach the embedded field: +if rl, ok := manager.GetOrAdd(key).(ratelimiter.RateLimiter); ok { + res := rl.Limiter.ReserveN(time.Now(), 1) + // ... +} +``` + ## Behavioral changes (no code change required) - **Per-key isolation now works.** Distinct keys have independent buckets; diff --git a/docs/TOKEN_BUCKET.md b/docs/TOKEN_BUCKET.md index 02bc4e8..afceff0 100644 --- a/docs/TOKEN_BUCKET.md +++ b/docs/TOKEN_BUCKET.md @@ -132,8 +132,8 @@ With `burst = 5` instead, the first **five** requests at `t = 0` are allowed ## Allow vs. Wait vs. Reserve -The same bucket can be consumed three ways. `ratelimiter.Limiter` exposes the -first two; the underlying `*rate.Limiter` also offers `Reserve`. +The same bucket can be consumed three ways. `ratelimiter.Limiter` exposes +`Allow` and `Wait`; the optional `ratelimiter.Reserver` interface adds `Reserve`. - **`Allow() bool`** — non-blocking. Returns `true` and consumes a token if one is available, else `false`. Use it to **drop** work on overload (HTTP 429, @@ -142,10 +142,15 @@ first two; the underlying `*rate.Limiter` also offers `Reserve`. is cancelled/expired. Use it to **shape** work — a background job or client that should slow down rather than fail. Beware unbounded goroutine buildup if arrivals persistently exceed `r`; always pass a `ctx` with a deadline. -- **`Reserve()` / `ReserveN()`** (on `*rate.Limiter`) — reserves a token and - tells you the `Delay()` until it becomes valid, without blocking. This is what - the middleware example uses to compute an accurate `Retry-After`. If you decide - not to proceed, call `reservation.Cancel()` to return the token. +- **`Reserve() Reservation`** (via the optional `Reserver` interface) — reserves + a token and tells you the `Delay()` until it becomes valid, without blocking. + This is what the middleware example uses to compute an accurate `Retry-After`. + 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). ## Choosing parameters diff --git a/examples/customstorage/main.go b/examples/customstorage/main.go new file mode 100644 index 0000000..8681a7e --- /dev/null +++ b/examples/customstorage/main.go @@ -0,0 +1,173 @@ +// Command customstorage shows how to plug your own [ratelimiter.Storage] +// implementation into a [ratelimiter.BucketLimiter] instead of the bundled +// [ratelimiter.InMemoryStorage]. +// +// The store here is a size-bounded LRU: it keeps at most -cap distinct keys and +// evicts the least-recently-used key when it would overflow. This is the classic +// reason to bring your own store — capping memory for an unbounded key space +// (for example per-IP limiting exposed to the internet) instead of, or in +// addition to, the manager's time-based idle eviction. +// +// Run it: +// +// go run ./examples/customstorage -cap 2 +// +// Watch the log lines: asking for a third distinct key evicts the +// least-recently-used one, and a limiter rebuilt after eviction starts from a +// fresh, full bucket. +package main + +import ( + "container/list" + "flag" + "fmt" + "sync" + "time" + + "github.com/slashdevops/ratelimiter" + "golang.org/x/time/rate" +) + +// LRUStorage is a goroutine-safe, size-bounded implementation of +// [ratelimiter.Storage]. It holds at most capacity entries; inserting beyond +// that evicts the least-recently-used key. Load and LoadOrStore count as uses +// and move a key to the most-recently-used position. +// +// It stores the Limiter instances in process memory, exactly like the bundled +// InMemoryStorage — the only added behavior is the LRU size cap. +type LRUStorage[K comparable, V any] struct { + mu sync.Mutex + capacity int + ll *list.List // front = most-recently-used + items map[K]*list.Element +} + +type lruEntry[K comparable, V any] struct { + key K + value V +} + +// NewLRUStorage returns an empty store that keeps at most capacity keys. +func NewLRUStorage[K comparable, V any](capacity int) *LRUStorage[K, V] { + if capacity < 1 { + capacity = 1 + } + return &LRUStorage[K, V]{ + capacity: capacity, + ll: list.New(), + items: make(map[K]*list.Element, capacity), + } +} + +// Store implements [ratelimiter.Storage]. +func (s *LRUStorage[K, V]) Store(key K, value V) { + s.mu.Lock() + defer s.mu.Unlock() + s.set(key, value) +} + +// Load implements [ratelimiter.Storage]. A hit is promoted to most-recently-used. +func (s *LRUStorage[K, V]) Load(key K) (V, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if el, ok := s.items[key]; ok { + s.ll.MoveToFront(el) + return el.Value.(*lruEntry[K, V]).value, true + } + var zero V + return zero, false +} + +// LoadOrStore implements [ratelimiter.Storage]. It is atomic under the mutex, so +// concurrent callers racing on the same new key all receive the same value. +func (s *LRUStorage[K, V]) LoadOrStore(key K, value V) (V, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if el, ok := s.items[key]; ok { + s.ll.MoveToFront(el) + return el.Value.(*lruEntry[K, V]).value, true + } + s.set(key, value) + return value, false +} + +// Delete implements [ratelimiter.Storage]. +func (s *LRUStorage[K, V]) Delete(key K) { + s.mu.Lock() + defer s.mu.Unlock() + if el, ok := s.items[key]; ok { + s.remove(el) + } +} + +// Range implements [ratelimiter.Storage]. Iteration order is most- to +// least-recently-used and does not count as a use. +func (s *LRUStorage[K, V]) Range(f func(K, V) bool) { + s.mu.Lock() + defer s.mu.Unlock() + for el := s.ll.Front(); el != nil; el = el.Next() { + e := el.Value.(*lruEntry[K, V]) + if !f(e.key, e.value) { + return + } + } +} + +// set inserts or updates key and evicts the LRU entry if over capacity. +// Callers must hold s.mu. +func (s *LRUStorage[K, V]) set(key K, value V) { + if el, ok := s.items[key]; ok { + el.Value.(*lruEntry[K, V]).value = value + s.ll.MoveToFront(el) + return + } + el := s.ll.PushFront(&lruEntry[K, V]{key: key, value: value}) + s.items[key] = el + if s.ll.Len() > s.capacity { + if oldest := s.ll.Back(); oldest != nil { + evicted := oldest.Value.(*lruEntry[K, V]).key + s.remove(oldest) + fmt.Printf(" [LRUStorage] capacity %d exceeded, evicted %v\n", s.capacity, evicted) + } + } +} + +// remove deletes an element from both the list and the index. Callers must hold s.mu. +func (s *LRUStorage[K, V]) remove(el *list.Element) { + s.ll.Remove(el) + delete(s.items, el.Value.(*lruEntry[K, V]).key) +} + +func main() { + capacity := flag.Int("cap", 2, "maximum number of distinct keys kept in the store") + flag.Parse() + + // Compile-time check that our type satisfies the interface. + var _ ratelimiter.Storage[string, ratelimiter.Limiter] = NewLRUStorage[string, ratelimiter.Limiter](*capacity) + + // Bring our own store instead of ratelimiter.NewInMemoryStorage(...). + store := NewLRUStorage[string, ratelimiter.Limiter](*capacity) + + // One token per key; a used key needs ~1s to refill, so re-use is observable. + newLimiter := ratelimiter.NewRateLimiterFunc(rate.Every(time.Second), 1) + + // deleteAfter <= 0 disables the manager's time-based eviction so the only + // thing bounding memory here is our LRU cap — keeping the demo focused. + manager := ratelimiter.NewBucketLimiter(newLimiter, 0, store) + defer manager.Close() + + fmt.Printf("store capacity = %d\n\n", *capacity) + + // Touch three distinct keys. With cap=2 the first key is evicted when the + // third arrives; its bucket is then rebuilt fresh on the next request. + for _, key := range []string{"alice", "bob", "carol", "alice"} { + allowed := manager.GetOrAdd(key).Allow() + fmt.Printf("request %-6s -> allowed=%v\n", key, allowed) + } + + fmt.Println("\ncurrent keys (most- to least-recently-used):") + store.Range(func(k string, _ ratelimiter.Limiter) bool { + fmt.Printf(" %s\n", k) + return true + }) +} diff --git a/examples/middleware/main.go b/examples/middleware/main.go index 51023fd..1e8ec54 100644 --- a/examples/middleware/main.go +++ b/examples/middleware/main.go @@ -65,26 +65,33 @@ func IPRateLimiter(manager *ratelimiter.BucketLimiter[string]) Middleware { ip := clientIP(r) limiter := manager.GetOrAdd(ip) - // The manager hands out *rate.Limiter values via NewRateLimiterFunc, - // so we can use the richer Reserve API for accurate headers. - rl, ok := limiter.(*rate.Limiter) + setHeader(w, "RateLimit-Limit", strconv.Itoa(limiter.Burst())) + + // Prefer the backend-agnostic Reserver capability: it yields the exact + // delay until a token is available, which drives accurate Retry-After + // and RateLimit-Reset headers for ANY Limiter implementation — a + // *rate.Limiter, a Redis/Valkey-backed one, etc. Limiters that cannot + // reserve degrade to a plain Allow() decision without timing headers. + reserver, ok := limiter.(ratelimiter.Reserver) if !ok { - // Fallback for custom Limiter implementations: Allow-only. if !limiter.Allow() { - http.Error(w, "too many requests", http.StatusTooManyRequests) + w.Header().Set("Retry-After", "1") + http.Error(w, fmt.Sprintf("too many requests from %s", ip), http.StatusTooManyRequests) return } next.ServeHTTP(w, r) return } - now := time.Now() - res := rl.ReserveN(now, 1) - delay := res.DelayFrom(now) + res := reserver.Reserve() + delay := res.Delay() - setHeader(w, "RateLimit-Limit", strconv.Itoa(rl.Burst())) - setHeader(w, "RateLimit-Remaining", strconv.Itoa(int(rl.TokensAt(now)))) setHeader(w, "RateLimit-Reset", strconv.Itoa(secondsCeil(delay))) + // RateLimit-Remaining is backend-specific; enrich it opportunistically + // when the limiter can report its current token count. + if rl, ok := limiter.(ratelimiter.RateLimiter); ok { + setHeader(w, "RateLimit-Remaining", strconv.Itoa(int(rl.TokensAt(time.Now())))) + } if !res.OK() || delay > 0 { // Not allowed right now: give the token back and reject. diff --git a/limiter.go b/limiter.go index 472b5c8..18539e3 100644 --- a/limiter.go +++ b/limiter.go @@ -6,6 +6,11 @@ import "context" // subset of the methods on *golang.org/x/time/rate.Limiter, so a standard // *rate.Limiter satisfies Limiter directly. // +// A Limiter may optionally also implement [Reserver] to expose token +// reservations (and the delay until a token is available). Callers such as HTTP +// middleware feature-detect that capability with a type assertion; the default +// limiter from [NewRateLimiterFunc] provides it. +// // Implementations MUST be safe for concurrent use by multiple goroutines. type Limiter interface { // Burst returns the maximum number of tokens that can be consumed in a diff --git a/rate.go b/rate.go index 25e3cef..ea347ad 100644 --- a/rate.go +++ b/rate.go @@ -2,8 +2,27 @@ package ratelimiter import "golang.org/x/time/rate" +// RateLimiter adapts a *golang.org/x/time/rate.Limiter to both the [Limiter] +// and [Reserver] interfaces. The embedded *rate.Limiter supplies Allow, Wait, +// and Burst directly (along with its other methods, such as TokensAt); the +// added Reserve method exposes reservations through the backend-agnostic +// [Reservation] interface, so callers like HTTP middleware can compute accurate +// Retry-After headers without depending on the concrete *rate.Limiter type. +// +// [NewRateLimiterFunc] hands these out, so the default limiter is a [Reserver]. +type RateLimiter struct { + *rate.Limiter +} + +// Reserve implements [Reserver]. It reserves one token at the current instant; +// the returned [Reservation]'s Delay reports how long until it becomes valid. +func (r RateLimiter) Reserve() Reservation { + return r.Limiter.Reserve() +} + // NewRateLimiterFunc returns a factory suitable for [NewBucketLimiter] that -// builds an independent *golang.org/x/time/rate.Limiter for each new key. +// builds an independent [RateLimiter] (wrapping a *golang.org/x/time/rate.Limiter) +// for each new key. // // - limit is the sustained refill rate in tokens per second. Use // rate.Every(d) to express "one token every d", or rate.Inf for no limit. @@ -13,8 +32,17 @@ import "golang.org/x/time/rate" // Example: 5 requests per second, absorbing bursts of up to 10: // // newLimiter := ratelimiter.NewRateLimiterFunc(rate.Limit(5), 10) +// +// The returned limiters satisfy [Reserver] in addition to [Limiter]. func NewRateLimiterFunc(limit rate.Limit, burst int) func() Limiter { return func() Limiter { - return rate.NewLimiter(limit, burst) + return RateLimiter{rate.NewLimiter(limit, burst)} } } + +// Compile-time guarantees for the adapter and the interface it mirrors. +var ( + _ Limiter = RateLimiter{} + _ Reserver = RateLimiter{} + _ Reservation = (*rate.Reservation)(nil) +) diff --git a/rate_test.go b/rate_test.go new file mode 100644 index 0000000..8d4a44a --- /dev/null +++ b/rate_test.go @@ -0,0 +1,60 @@ +package ratelimiter + +import ( + "testing" + "time" + + "golang.org/x/time/rate" +) + +func TestNewRateLimiterFuncImplementsReserver(t *testing.T) { + lim := NewRateLimiterFunc(rate.Limit(1), 1)() + + if _, ok := lim.(Reserver); !ok { + t.Fatalf("limiter from NewRateLimiterFunc does not implement Reserver") + } +} + +func TestReserveConsumesAndCancelReturnsToken(t *testing.T) { + // 1 token/sec, burst 1. + lim := NewRateLimiterFunc(rate.Limit(1), 1)() + r := lim.(Reserver) + + r.Reserve() // consume the initial token (delay 0) + + // The next token is a bit under a second out. + pending := r.Reserve() + if !pending.OK() { + t.Fatal("pending reservation should be OK") + } + withoutCancel := pending.Delay() + if withoutCancel <= 0 { + t.Fatalf("pending reservation should have a positive delay, got %v", withoutCancel) + } + + // Returning that reserved token should make the following reservation land + // at roughly the same delay rather than stacking a second token's wait. + pending.Cancel() + next := r.Reserve() + defer next.Cancel() + if got := next.Delay(); got > withoutCancel+500*time.Millisecond { + t.Fatalf("Cancel did not return the token: delay grew from %v to %v", withoutCancel, got) + } +} + +func TestReserveDelayDrivesRetryAfter(t *testing.T) { + // 2 tokens/sec, burst 1: after consuming the token the next is ~500ms out. + lim := NewRateLimiterFunc(rate.Limit(2), 1)() + r := lim.(Reserver) + + first := r.Reserve() + if first.Delay() != 0 { + t.Fatalf("first reservation should be immediate, got %v", first.Delay()) + } + + second := r.Reserve() + defer second.Cancel() + if got := second.Delay(); got <= 0 || got > time.Second { + t.Fatalf("expected a sub-second positive delay, got %v", got) + } +} diff --git a/reservation.go b/reservation.go new file mode 100644 index 0000000..3bd7b6e --- /dev/null +++ b/reservation.go @@ -0,0 +1,53 @@ +package ratelimiter + +import "time" + +// Reservation is a backend-agnostic view of a single token that has been +// reserved for the earliest instant it can be granted. It mirrors the useful +// subset of the methods on *golang.org/x/time/rate.Reservation, so a standard +// *rate.Reservation satisfies Reservation directly. +// +// A reservation holds a token until it is either used (by proceeding after +// Delay) or returned with Cancel. Whichever a caller does, it must do exactly +// one of the two to avoid leaking or double-counting tokens. +type Reservation interface { + // OK reports whether a token can be granted within the limiter's limits. + // If it returns false the caller must not proceed and must not act on + // Delay; the request can never be satisfied (for example a burst larger + // than the bucket capacity). + OK() bool + + // Delay returns how long the caller must wait, from now, before the + // reserved token becomes valid. Zero means it may be used immediately. + Delay() time.Duration + + // Cancel returns the reserved token to the limiter if it has not yet been + // consumed, so it does not count against the limit. Call it when you decide + // not to proceed (for example when rejecting an over-limit request). + Cancel() +} + +// Reserver is an OPTIONAL capability a [Limiter] may implement in addition to +// the core Allow/Wait/Burst methods. When a limiter can pre-reserve a token and +// report the exact delay until it is valid, implementing Reserver lets callers +// such as HTTP middleware emit accurate Retry-After and RateLimit-Reset headers +// for any backend — not only *rate.Limiter. +// +// Feature-detect it with a type assertion and degrade gracefully when it is +// absent: +// +// if r, ok := lim.(ratelimiter.Reserver); ok { +// res := r.Reserve() +// defer res.Cancel() // if you end up not proceeding +// // ... use res.OK() / res.Delay() for headers and the decision +// } else { +// allowed := lim.Allow() // no timing information available +// } +// +// The [RateLimiter] returned by [NewRateLimiterFunc] implements Reserver. +type Reserver interface { + // Reserve reserves one token for the earliest instant it can be granted and + // returns a [Reservation] describing it. The token is held until it is used + // or Reservation.Cancel is called. + Reserve() Reservation +}