diff --git a/README.md b/README.md index bf811a6..8167078 100644 --- a/README.md +++ b/README.md @@ -268,8 +268,31 @@ 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). +Redis / [Valkey](https://github.com/valkey-io/valkey-go)**: a datastore-backed +`Limiter`, given its key by +[`WithLimiterFactoryForKey`](#per-key-limiters-and-shared-backends). + +### Per-key limiters and shared backends + +A `Limiter` is bound to exactly one key — `Allow()` takes no arguments, so the +instance *is* the bucket. `WithLimiterFactoryForKey` builds each one **from its +key**, which is what a limiter needs when its counter lives somewhere else: + +```go +bl := ratelimiter.NewBucketLimiter(nil, time.Minute, + ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), + ratelimiter.WithLimiterFactoryForKey(func(key string) ratelimiter.Limiter { + if shared != nil { + return sharedLimiter{client: shared, key: "rl:" + key} + } + return ratelimiter.RateLimiter{Limiter: rate.NewLimiter(limit, burst)} + }), +) +``` + +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. ## Scope: single-process only diff --git a/bucket_limiter.go b/bucket_limiter.go index d88f770..0ce3ff1 100644 --- a/bucket_limiter.go +++ b/bucket_limiter.go @@ -1,6 +1,7 @@ package ratelimiter import ( + "fmt" "sync" "sync/atomic" "time" @@ -24,7 +25,15 @@ const defaultSweepDivisor = 2 // // The zero value is not usable; construct one with [NewBucketLimiter]. type BucketLimiter[K comparable] struct { - newLimiter func() Limiter + newLimiter func() Limiter + + // newLimiterForKey, when set by [WithLimiterFactoryForKey], is used in + // place of newLimiter and receives the key the limiter will govern. It is + // what makes a limiter backed by a shared datastore possible: such a + // limiter has to know which remote key holds its counter, and Allow/Wait + // take no arguments. + newLimiterForKey func(K) Limiter + deleteAfter time.Duration interval time.Duration now func() time.Time @@ -46,6 +55,14 @@ type Option func(*config) type config struct { now func() time.Time interval time.Duration + + // keyFactory holds a func(K) Limiter. It is stored as any because Option + // is deliberately not generic: making it Option[K] would force every + // existing call such as WithClock(now) to be explicitly instantiated, + // which would break source compatibility for every current user. The type + // is recovered with a checked assertion in NewBucketLimiter, so a mismatch + // is caught at construction rather than at first use. + keyFactory any } // WithClock overrides the time source used for idle tracking and eviction. @@ -58,6 +75,53 @@ func WithClock(now func() time.Time) Option { } } +// WithLimiterFactoryForKey builds each key's [Limiter] from the key itself, +// instead of from the argument-less factory passed to [NewBucketLimiter]. +// +// # Why this exists +// +// A [Limiter] is bound to exactly one key: Allow and Wait take no arguments, so +// the instance IS the bucket. For an in-process limiter that is invisible — +// every bucket is equivalent, so an argument-less factory suffices. For a +// limiter whose state lives somewhere else — Redis, Valkey, any shared store — +// it is the whole problem: the limiter must know WHICH remote key holds its +// counter, and nothing in the old API ever told it. +// +// Without this option the only injection point that sees both the key and a +// place to hold a shared client is [Storage.LoadOrStore], which meant using the +// store as a factory rather than as a value container. That works, but it +// reinterprets an interface whose documented job is to hold values, and it puts +// construction logic in a place nobody looks for it. This option is the +// first-class version: the store goes back to storing, and the factory does the +// building. +// +// newLimiter := func(key string) ratelimiter.Limiter { +// if sharedStoreAvailable { +// return valkeylimiter.New(client, "rl:"+key, limit, burst) +// } +// return ratelimiter.RateLimiter{rate.NewLimiter(limit, burst)} +// } +// +// bl := ratelimiter.NewBucketLimiter(nil, time.Minute, +// ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), +// ratelimiter.WithLimiterFactoryForKey(newLimiter), +// ) +// +// Note what the storage is in that example: the ordinary in-memory one, in +// BOTH branches. The shared state lives in the datastore, inside the limiter; +// the map only caches one lightweight handle per key. A [Storage] is always an +// in-process container — see docs/CUSTOM_STORAGE.md. +// +// When this option is supplied the newLimiter argument to [NewBucketLimiter] is +// ignored and may be nil. Supplying neither panics at construction. +func WithLimiterFactoryForKey[K comparable](newLimiter func(K) Limiter) Option { + return func(c *config) { + if newLimiter != nil { + c.keyFactory = newLimiter + } + } +} + // WithSweepInterval overrides how often the background eviction goroutine runs. // When unset, it defaults to deleteAfter/2. Ignored when deleteAfter <= 0. func WithSweepInterval(d time.Duration) Option { @@ -72,6 +136,8 @@ func WithSweepInterval(d time.Duration) Option { // // - newLimiter is called once per new key to build that key's independent // [Limiter]. Use [NewRateLimiterFunc] for the common *rate.Limiter case. +// It may be nil when [WithLimiterFactoryForKey] supplies a key-aware +// factory instead; supplying neither panics. // - deleteAfter is the idle duration after which an unused key is evicted. // A value <= 0 disables eviction (limiters live until [BucketLimiter.Remove] // or [BucketLimiter.Close]); prefer this only for bounded key spaces. @@ -99,14 +165,37 @@ func NewBucketLimiter[K comparable]( } } + // Recover the key-aware factory's real type. A mismatch means the caller + // wrote WithLimiterFactoryForKey with a different key type than the + // BucketLimiter's, which is a programming error worth reporting here + // rather than as a nil limiter on the first request for a new key. + var newLimiterForKey func(K) Limiter + + if cfg.keyFactory != nil { + typed, ok := cfg.keyFactory.(func(K) Limiter) + if !ok { + panic(fmt.Sprintf( + "ratelimiter: WithLimiterFactoryForKey was given a %T, but this BucketLimiter's key type is %T", + cfg.keyFactory, *new(K), + )) + } + + newLimiterForKey = typed + } + + if newLimiter == nil && newLimiterForKey == nil { + panic("ratelimiter: NewBucketLimiter needs either a newLimiter factory or WithLimiterFactoryForKey; both are nil, so no limiter could ever be built") + } + b := &BucketLimiter[K]{ - newLimiter: newLimiter, - deleteAfter: deleteAfter, - interval: interval, - now: cfg.now, - storage: storage, - stop: make(chan struct{}), - done: make(chan struct{}), + newLimiter: newLimiter, + newLimiterForKey: newLimiterForKey, + deleteAfter: deleteAfter, + interval: interval, + now: cfg.now, + storage: storage, + stop: make(chan struct{}), + done: make(chan struct{}), } if deleteAfter > 0 { @@ -127,13 +216,22 @@ func (b *BucketLimiter[K]) GetOrAdd(key K) Limiter { if !ok { // LoadOrStore makes creation atomic: if another goroutine wins the // race, we discard our fresh limiter and use the stored one. - limiter, _ = b.storage.LoadOrStore(key, b.newLimiter()) + limiter, _ = b.storage.LoadOrStore(key, b.build(key)) } b.touch(key) return limiter } +// build constructs the Limiter for key, preferring the key-aware factory. +func (b *BucketLimiter[K]) build(key K) Limiter { + if b.newLimiterForKey != nil { + return b.newLimiterForKey(key) + } + + return b.newLimiter() +} + // touch records the current time as key's last-use time. func (b *BucketLimiter[K]) touch(key K) { if b.deleteAfter <= 0 { diff --git a/docs/CUSTOM_STORAGE.md b/docs/CUSTOM_STORAGE.md index 4425e44..06e3961 100644 --- a/docs/CUSTOM_STORAGE.md +++ b/docs/CUSTOM_STORAGE.md @@ -34,8 +34,9 @@ to **hold the per-key limiters in memory**. Common reasons: - **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). + key-aware limiters. Since `WithLimiterFactoryForKey` this is rarely the right + tool: that option gives a factory the key directly. See the + [Valkey pattern](#distributed-limiting-with-redis--valkey). If none of those apply, use `NewInMemoryStorage` — it is correct, atomic, and fast. @@ -199,22 +200,47 @@ 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.) +state lives* — in the datastore, **atomically**. So you implement the +[`Limiter`](../limiter.go) interface (`Allow`, `Wait`, `Burst`) backed by +[valkey-go](https://github.com/valkey-io/valkey-go), and you give it its key +with [`WithLimiterFactoryForKey`](#distributed-limiting-with-redis--valkey). + +"Atomically" does not have to mean a server-side script. A token bucket is a +read-modify-write — read `(tokens, ts)`, refill, compare, write — which is not +one command, so it needs `EVAL` (below) or `WATCH`/`MULTI`/`EXEC`. If you would +rather not run a script, **change the algorithm rather than the atomicity**: a +sliding-window counter is a single atomic `INCR` plus a `PEXPIRE` on the first +hit of each window, at the cost of stepwise rather than continuous refill and a +small over-admission at a window boundary. Both are valid `Limiter` +implementations; the library does not care which you pick. + +> **Use `WithLimiterFactoryForKey`.** A `Limiter` is bound to exactly one key — +> `Allow()` and `Wait()` take no arguments, so the instance *is* the bucket — and +> a datastore-backed limiter must know **which** remote key holds its counter. +> `WithLimiterFactoryForKey(func(K) Limiter)` gives the factory the key: +> +> ```go +> bl := ratelimiter.NewBucketLimiter(nil, time.Minute, +> ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), +> ratelimiter.WithLimiterFactoryForKey(func(key string) ratelimiter.Limiter { +> return valkeylimiter.New(client, "rl:"+key, limit, burst) +> }), +> ) +> ``` +> +> Note the storage in that example: the **ordinary in-memory one**. The shared +> state lives in the datastore, inside the limiter; the map only caches one +> lightweight handle per key. That is the division of labour — `Storage` holds +> values in this process, `Limiter` decides, and only the `Limiter` knows where +> its state lives. +> +> **Before this option existed** the only injection point that saw both the key +> and a place to hold a shared client was `Storage.LoadOrStore`, so the store had +> to be used as a factory/cache rather than as a value container. That still +> works and is still supported, but it reinterprets an interface whose documented +> job is to hold values, and it hides construction somewhere nobody looks for it. +> Prefer the option. (If you would rather not reuse `BucketLimiter` at all, a +> ~30-line standalone manager over the same `valkeyLimiter` works too.) ### The flow diff --git a/examples/keyfactory/main.go b/examples/keyfactory/main.go new file mode 100644 index 0000000..d8bdad0 --- /dev/null +++ b/examples/keyfactory/main.go @@ -0,0 +1,94 @@ +// Command keyfactory demonstrates WithLimiterFactoryForKey: building each key's +// Limiter from the key itself, which is what a limiter needs when its state +// lives outside the process. +// +// The "shared" limiter here is a map behind a mutex rather than Valkey, so the +// example runs with no dependencies — but it has the property that matters: +// several limiter instances built for the same key address ONE budget, and +// evicting an instance does not reset it. +// +// go run ./examples/keyfactory +// go run ./examples/keyfactory -shared=false +package main + +import ( + "context" + "flag" + "fmt" + "sync" + "time" + + "golang.org/x/time/rate" + + "github.com/slashdevops/ratelimiter" +) + +// sharedStore stands in for Valkey: state addressed by key, outliving any +// particular limiter object. +type sharedStore struct { + mu sync.Mutex + tokens map[string]int +} + +func (s *sharedStore) take(key string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + if s.tokens[key] <= 0 { + return false + } + + s.tokens[key]-- + + return true +} + +// sharedLimiter is a ratelimiter.Limiter whose Allow consults the shared store +// for ITS key. Building one requires knowing the key — which is exactly what +// the argument-less factory could not provide. +type sharedLimiter struct { + store *sharedStore + key string + burst int +} + +func (l sharedLimiter) Burst() int { return l.burst } +func (l sharedLimiter) Allow() bool { + return l.store.take(l.key) +} + +func (l sharedLimiter) Wait(_ context.Context) error { return nil } + +func main() { + shared := flag.Bool("shared", true, "use the shared (out-of-process) limiter") + flag.Parse() + + store := &sharedStore{tokens: map[string]int{"alice": 2, "bob": 2}} + + // The one seam. Swap the branch and everything above it stays the same — + // including the storage, which is the ordinary in-memory one either way: + // it caches handles, it does not hold the budget. + newLimiter := func(key string) ratelimiter.Limiter { + if *shared { + return sharedLimiter{store: store, key: key, burst: 2} + } + + return ratelimiter.RateLimiter{Limiter: rate.NewLimiter(rate.Every(time.Hour), 2)} + } + + bl := ratelimiter.NewBucketLimiter(nil, time.Minute, + ratelimiter.NewInMemoryStorage[string, ratelimiter.Limiter](), + ratelimiter.WithLimiterFactoryForKey(newLimiter), + ) + defer bl.Close() + + for _, key := range []string{"alice", "alice", "alice", "bob"} { + fmt.Printf("%-6s allow=%v\n", key, bl.GetOrAdd(key).Allow()) + } + + // Evict alice's handle. With a shared limiter the budget stays spent, + // because it was never in the handle. With the in-process one it comes + // back full — run with -shared=false to see the difference. + bl.Remove("alice") + fmt.Printf("\nafter evicting alice's handle:\n%-6s allow=%v\n", "alice", bl.GetOrAdd("alice").Allow()) +} diff --git a/key_factory_test.go b/key_factory_test.go new file mode 100644 index 0000000..1f37079 --- /dev/null +++ b/key_factory_test.go @@ -0,0 +1,259 @@ +package ratelimiter + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + "time" + + "golang.org/x/time/rate" +) + +// sharedLimiter stands in for a limiter whose state lives somewhere other than +// this object — Redis, Valkey, a database. What matters for these tests is the +// property that makes such a limiter different from *rate.Limiter: several +// instances built for the SAME key share one budget, and instances built for +// different keys do not. +type sharedLimiter struct { + key string + state *sharedState +} + +type sharedState struct { + mu sync.Mutex + remaining map[string]int +} + +func newSharedState(perKey int, keys ...string) *sharedState { + s := &sharedState{remaining: make(map[string]int, len(keys))} + for _, k := range keys { + s.remaining[k] = perKey + } + + return s +} + +func (l *sharedLimiter) Burst() int { return 0 } + +func (l *sharedLimiter) Allow() bool { + l.state.mu.Lock() + defer l.state.mu.Unlock() + + if l.state.remaining[l.key] <= 0 { + return false + } + + l.state.remaining[l.key]-- + + return true +} + +func (l *sharedLimiter) Wait(context.Context) error { return nil } + +// TestWithLimiterFactoryForKeyBindsTheKey is the reason the option exists. +// +// A Limiter is bound to exactly one key — Allow and Wait take no arguments, so +// the instance IS the bucket. An argument-less factory can therefore only build +// limiters that keep their state in themselves. A limiter whose counter lives +// in a shared datastore has to know WHICH remote key is its own, and before +// this option nothing in the API ever told it. +func TestWithLimiterFactoryForKeyBindsTheKey(t *testing.T) { + t.Parallel() + + var built []string + + var mu sync.Mutex + + state := newSharedState(2, "alice", "bob") + + bl := NewBucketLimiter(nil, time.Minute, + NewInMemoryStorage[string, Limiter](), + WithLimiterFactoryForKey(func(key string) Limiter { + mu.Lock() + built = append(built, key) + mu.Unlock() + + return &sharedLimiter{key: key, state: state} + }), + ) + defer bl.Close() + + // Two tokens each, independently. + for range 2 { + if !bl.GetOrAdd("alice").Allow() { + t.Fatal("alice should have budget") + } + + if !bl.GetOrAdd("bob").Allow() { + t.Fatal("bob should have budget") + } + } + + if bl.GetOrAdd("alice").Allow() { + t.Error("alice's budget should be spent") + } + + if bl.GetOrAdd("bob").Allow() { + t.Error("bob's budget should be spent; one key's spending must not affect another's") + } + + mu.Lock() + defer mu.Unlock() + + if len(built) != 2 { + t.Errorf("factory called %d times for 2 keys: %v — limiters must be built once per key and cached", len(built), built) + } +} + +// TestKeyFactoryStateSurvivesEviction pins the property that makes a shared +// limiter useful: because the state is not in the returned object, evicting the +// object does not reset the budget. An in-process *rate.Limiter would come back +// with a full bucket here; a datastore-backed one must not. +func TestKeyFactoryStateSurvivesEviction(t *testing.T) { + t.Parallel() + + state := newSharedState(1, "k") + + bl := NewBucketLimiter(nil, time.Minute, + NewInMemoryStorage[string, Limiter](), + WithLimiterFactoryForKey(func(key string) Limiter { + return &sharedLimiter{key: key, state: state} + }), + ) + defer bl.Close() + + if !bl.GetOrAdd("k").Allow() { + t.Fatal("the first call should be allowed") + } + + // Drop the cached handle, exactly as idle eviction would. + bl.Remove("k") + + // A fresh handle is built, but it addresses the same shared state. + if bl.GetOrAdd("k").Allow() { + t.Error("budget came back after the handle was evicted; the limiter is keeping state in the object, not in the shared store") + } +} + +// TestNewBucketLimiterPrefersTheKeyFactory: when both are supplied the +// key-aware one wins, because it is strictly more informed. +func TestNewBucketLimiterPrefersTheKeyFactory(t *testing.T) { + t.Parallel() + + bl := NewBucketLimiter( + func() Limiter { return RateLimiter{rate.NewLimiter(rate.Inf, 1)} }, + time.Minute, + NewInMemoryStorage[string, Limiter](), + WithLimiterFactoryForKey(func(key string) Limiter { + return &sharedLimiter{key: key, state: newSharedState(0, key)} + }), + ) + defer bl.Close() + + if _, ok := bl.GetOrAdd("k").(*sharedLimiter); !ok { + t.Error("the argument-less factory was used even though a key-aware one was supplied") + } +} + +// TestNewBucketLimiterWithNoFactoryPanics: nothing could ever be built, and +// saying so at construction beats a nil-limiter panic on the first request for +// a new key, which is a different place and a much worse moment. +func TestNewBucketLimiterWithNoFactoryPanics(t *testing.T) { + t.Parallel() + + defer func() { + r := recover() + if r == nil { + t.Fatal("expected a panic when neither factory is supplied") + } + + if msg := fmt.Sprint(r); msg == "" { + t.Error("panic carried no message") + } + }() + + _ = NewBucketLimiter[string](nil, time.Minute, NewInMemoryStorage[string, Limiter]()) +} + +// TestWithLimiterFactoryForKeyRejectsAMismatchedKeyType: Option is not generic +// (making it so would break every existing WithClock call), so the factory's +// key type is checked at construction instead. A mismatch is a programming +// error and must not surface as a nil limiter later. +func TestWithLimiterFactoryForKeyRejectsAMismatchedKeyType(t *testing.T) { + t.Parallel() + + defer func() { + r := recover() + if r == nil { + t.Fatal("expected a panic when the factory's key type does not match the limiter's") + } + + // Assert WHICH panic. Checking only that "a panic happened" is too + // weak here and was: with the type check removed the assertion simply + // yields nil, the factory ends up unset, and the both-factories-nil + // guard panics instead — so the test passed while the behaviour it + // names was gone. Found by mutating the check away. + if msg := fmt.Sprint(r); !strings.Contains(msg, "key type") { + t.Errorf("panic was %q, which does not report a key-type mismatch", msg) + } + }() + + // int factory, string limiter. A non-nil newLimiter is supplied so that + // the both-factories-nil guard cannot be what fires. + _ = NewBucketLimiter[string]( + func() Limiter { return RateLimiter{rate.NewLimiter(rate.Inf, 1)} }, + time.Minute, + NewInMemoryStorage[string, Limiter](), + WithLimiterFactoryForKey(func(int) Limiter { return nil }), + ) +} + +// TestKeyFactoryIsSafeUnderConcurrentFirstUse: GetOrAdd's contract is that +// concurrent callers racing on a brand-new key all receive the SAME instance. +// That has to keep holding when the instance is built from the key. +// +// The assertion is instance IDENTITY, deliberately, and an earlier version of +// this test got it wrong in a way worth recording: it counted tokens spent +// against the shared state. That passes whether or not creation is atomic — +// every instance addresses the same shared budget, so building ten of them +// still spends exactly ten tokens. It asserted something that could not fail. +// +// Note what is NOT asserted: that the factory runs once. GetOrAdd evaluates +// build(key) as the argument to LoadOrStore, so under a race each goroutine +// legitimately builds a candidate and all but one are discarded. Requiring a +// single call would pin an implementation detail the library does not promise. +func TestKeyFactoryIsSafeUnderConcurrentFirstUse(t *testing.T) { + t.Parallel() + + bl := NewBucketLimiter(nil, time.Minute, + NewInMemoryStorage[string, Limiter](), + WithLimiterFactoryForKey(func(key string) Limiter { + // A pointer, so identity is observable. + return &sharedLimiter{key: key, state: newSharedState(1<<30, key)} + }), + ) + defer bl.Close() + + const racers = 100 + + got := make([]Limiter, racers) + + var wg sync.WaitGroup + + for i := range racers { + wg.Go(func() { + got[i] = bl.GetOrAdd("hot") + }) + } + + wg.Wait() + + first := got[0] + for i, l := range got { + if l != first { + t.Fatalf("racer %d received a different limiter instance; concurrent callers on one key must share a bucket, and two instances means two budgets", i) + } + } +}