feat: build a key's Limiter from its key, for datastore-backed limiters

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
 

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 {

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
 

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())
+}