From 88cfa36970379ae82718c2114a36ad3fd749114b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20Gonz=C3=A1lez=20Di=20Antonio?= Date: Tue, 25 Aug 2026 18:17:39 +0200 Subject: [PATCH] feat: build a key's Limiter from its key, for datastore-backed limiters 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 `newLimiter func() Limiter` suffices. For a limiter whose counter lives somewhere else -- Redis, Valkey, any shared store -- it is the whole problem, because such a limiter has to know WHICH remote key is its own, and nothing in the API ever told it. WithLimiterFactoryForKey(func(K) Limiter) closes that gap. 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 entire seam for "shared budget when the datastore is there, in-process when it is not" -- and note the storage in both branches is the ordinary in-memory one, because it only caches handles. The shared state lives in the datastore, inside the Limiter. WHY NOT A DISTRIBUTED Storage, which is what everyone tries first: GetOrAdd hands the caller the Limiter and the caller calls Allow() on it. The manager never writes back -- Store is not called by it at all. So a Storage that serialised bucket state to Redis would deserialise a FULL bucket on every request, let the caller spend a token from an object nobody persists, and drop it. The limit would never be reached, and nothing would look wrong. docs/CUSTOM_STORAGE.md already said so; this commit gives people the thing to do instead of only the thing not to do. The previous advice was to use Storage.LoadOrStore as a key->limiter resolver, because it was the only place that saw both the key and a shared client. That still works and is still documented, but it reinterprets an interface whose job is to hold values and hides construction where nobody looks for it. Option is deliberately NOT made generic. Changing it to Option[K] would force every existing WithClock(now) call to be explicitly instantiated, breaking source compatibility for every current user; the factory is stored as any and recovered with a checked assertion at construction, so a key-type mismatch panics there rather than surfacing as a nil limiter on the first request for a new key. Supplying neither factory now panics at construction too, for the same reason. Tests, each verified to fail by mutating the code it guards -- and two of them were rewritten because the first version could not fail: - the key reaches the factory, and one key's spending does not affect another - state survives eviction of the handle, which is the property that separates a datastore-backed limiter from an in-process one - the key-aware factory wins when both are supplied - no factory at all panics at construction - a mismatched key type panics, and the test asserts WHICH panic: checking only that "a panic happened" passed with the type check removed, because the both-factories-nil guard fired instead - concurrent first use returns ONE instance, asserted by identity: the first version counted tokens against shared state, which passes whether or not creation is atomic, since every instance addresses the same budget examples/keyfactory shows the difference in six lines of output: evict the handle and a shared limiter stays spent while an in-process one comes back full. docs/CUSTOM_STORAGE.md and README updated. The distributed section also stops implying a Lua script is the only way to be atomic -- a token bucket is a read-modify-write and needs EVAL or WATCH/MULTI/EXEC, but a sliding-window counter is a single atomic INCR, and choosing the algorithm is a legitimate alternative to choosing a scripting engine. Coverage stays at 100% of statements. --- README.md | 27 +++- bucket_limiter.go | 116 ++++++++++++++-- docs/CUSTOM_STORAGE.md | 62 ++++++--- examples/keyfactory/main.go | 94 +++++++++++++ key_factory_test.go | 259 ++++++++++++++++++++++++++++++++++++ 5 files changed, 529 insertions(+), 29 deletions(-) create mode 100644 examples/keyfactory/main.go create mode 100644 key_factory_test.go 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) + } + } +}