Rate Limiting Go APIs per API Key and per Route Cleanly

package main

import (
	"math"
	"strconv"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

// Policy is a token bucket: Rate tokens per second, up to Burst at once. A request spends Cost tokens, so an
// expensive route can draw on the same bucket faster than a cheap one.
type Policy struct {
	Rate  float64
	Burst int
	Cost  int
}

type entry struct {
	lim  *rate.Limiter
	last time.Time
}

// Limiter keeps one bucket per (client, bucket name). Buckets are created on first use and swept when idle.
type Limiter struct {
	mu       sync.Mutex
	policies map[string]Policy // keyed by route name, "default" for the rest
	buckets  map[string]*entry
}

func NewLimiter(policies map[string]Policy) *Limiter {
	return &Limiter{policies: policies, buckets: map[string]*entry{}}
}

type Decision struct {
	Allowed    bool
	Limit      int
	Remaining  int
	RetryAfter time.Duration
}

func (l *Limiter) policy(route string) Policy {
	if p, ok := l.policies[route]; ok {
		return p
	}
	return l.policies["default"]
}

// Allow decides whether client may call route at now. The clock is a parameter so tests do not sleep.
func (l *Limiter) Allow(client, route string, now time.Time) Decision {
	p := l.policy(route)
	cost := max(p.Cost, 1)
	key := client + "|" + route
	l.mu.Lock()
	e, ok := l.buckets[key]
	if !ok {
		e = &entry{lim: rate.NewLimiter(rate.Limit(p.Rate), p.Burst)}
		l.buckets[key] = e
	}
	e.last = now
	l.mu.Unlock()

	if cost > p.Burst { // a request bigger than the bucket could never succeed; refuse it without waiting forever
		return Decision{Limit: p.Burst, RetryAfter: time.Duration(math.MaxInt64)}
	}
	r := e.lim.ReserveN(now, cost)
	if delay := r.DelayFrom(now); delay > 0 {
		r.CancelAt(now) // do not charge a refused request
		return Decision{Limit: p.Burst, RetryAfter: delay}
	}
	return Decision{Allowed: true, Limit: p.Burst, Remaining: int(e.lim.TokensAt(now))}
}

// Sweep removes buckets idle for longer than ttl, so a stream of one-off clients cannot grow the map without bound.
func (l *Limiter) Sweep(now time.Time, ttl time.Duration) (removed int) {
	l.mu.Lock()
	defer l.mu.Unlock()
	for k, e := range l.buckets {
		if now.Sub(e.last) > ttl {
			delete(l.buckets, k)
			removed++
		}
	}
	return removed
}

func (l *Limiter) Size() int {
	l.mu.Lock()
	defer l.mu.Unlock()
	return len(l.buckets)
}

func seconds(d time.Duration) string {
	return strconv.Itoa(int(math.Ceil(d.Seconds())))
}

package main

import (
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

// clientKey identifies who is being limited: the API key id when there is one, otherwise the address.
func clientKey(r *http.Request) string {
	if k := r.Header.Get("X-API-Key"); k != "" {
		if parts := strings.SplitN(k, "_", 3); len(parts) == 3 {
			return "key:" + parts[1] // the id, never the secret
		}
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	return "ip:" + host
}

// Limit wraps a handler under the named policy. The route name is explicit, not the URL path, so ids in paths
// do not create a bucket per resource.
func Limit(l *Limiter, route string, now func() time.Time, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		d := l.Allow(clientKey(r), route, now())
		w.Header().Set("RateLimit-Limit", strconv.Itoa(d.Limit))
		w.Header().Set("RateLimit-Remaining", strconv.Itoa(d.Remaining))
		if !d.Allowed {
			w.Header().Set("Retry-After", seconds(d.RetryAfter))
			w.Header().Set("Content-Type", "application/problem+json")
			w.WriteHeader(http.StatusTooManyRequests)
			w.Write([]byte(`{"title":"Too many requests","status":429}` + "\n"))
			return
		}
		next.ServeHTTP(w, r)
	})
}

package main

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"time"
)

func main() {
	clock := time.Date(2026, 6, 1, 12, 0, 0, 0, time.UTC)
	now := func() time.Time { return clock }
	l := NewLimiter(map[string]Policy{
		"default": {Rate: 5, Burst: 10, Cost: 1},
		"report":  {Rate: 5, Burst: 10, Cost: 5}, // expensive: five requests' worth
	})
	ok := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, "ok") })
	mux := http.NewServeMux()
	mux.Handle("GET /v1/tasks", Limit(l, "default", now, ok))
	mux.Handle("GET /v1/reports", Limit(l, "report", now, ok))

	call := func(key, path string) *httptest.ResponseRecorder {
		req := httptest.NewRequest("GET", path, nil)
		req.Header.Set("X-API-Key", key)
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, req)
		return rec
	}
	line := func(rec *httptest.ResponseRecorder) string {
		s := fmt.Sprint(rec.Code)
		if rec.Code == 429 {
			s += " (retry in " + rec.Header().Get("Retry-After") + "s)"
		}
		return s
	}

	fmt.Println("== 14 quick requests from key A to a cheap route (burst 10, refill 5/s)")
	var out []string
	for i := 0; i < 14; i++ {
		out = append(out, line(call("ghk_aaaaaaaa_secret", "/v1/tasks")))
	}
	fmt.Println(strings.Join(out, " "))
	fmt.Println("key B at the same moment:          ", line(call("ghk_bbbbbbbb_secret", "/v1/tasks")))

	rec := call("ghk_aaaaaaaa_secret", "/v1/tasks")
	fmt.Printf("429 headers: RateLimit-Limit=%s RateLimit-Remaining=%s Retry-After=%s\n", rec.Header().Get("RateLimit-Limit"), rec.Header().Get("RateLimit-Remaining"), rec.Header().Get("Retry-After"))

	clock = clock.Add(1 * time.Second)
	fmt.Println("one second later (5 tokens back):  ", line(call("ghk_aaaaaaaa_secret", "/v1/tasks")))

	fmt.Println("\n== an expensive route draws on its own bucket, five tokens per call")
	out = nil
	for i := 0; i < 4; i++ {
		out = append(out, line(call("ghk_cccccccc_secret", "/v1/reports")))
	}
	fmt.Println(strings.Join(out, " "), " <- burst 10 allows two, then it must wait")
	fmt.Println("same key, cheap route, unaffected: ", line(call("ghk_cccccccc_secret", "/v1/tasks")))

	fmt.Println("\n== memory: one bucket per client and route")
	for i := 0; i < 1000; i++ {
		call(fmt.Sprintf("ghk_%08d_x", i), "/v1/tasks")
	}
	fmt.Println("buckets after 1000 one-off clients:", l.Size())
	clock = clock.Add(11 * time.Minute)
	fmt.Println("swept after 10 idle minutes:       ", l.Sweep(clock, 10*time.Minute), "removed,", l.Size(), "left")
}

== 14 quick requests from key A to a cheap route (burst 10, refill 5/s)
200 200 200 200 200 200 200 200 200 200 429 (retry in 1s) 429 (retry in 1s) 429 (retry in 1s) 429 (retry in 1s)
key B at the same moment:           200
429 headers: RateLimit-Limit=10 RateLimit-Remaining=0 Retry-After=1
one second later (5 tokens back):   200

== an expensive route draws on its own bucket, five tokens per call
200 200 429 (retry in 1s) 429 (retry in 1s)  <- burst 10 allows two, then it must wait
same key, cheap route, unaffected:  200

== memory: one bucket per client and route
buckets after 1000 one-off clients: 1004
swept after 10 idle minutes:        1004 removed, 0 left

package main

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

var t0 = time.Date(2026, 6, 1, 12, 0, 0, 0, time.UTC)

func policies() map[string]Policy {
	return map[string]Policy{"default": {Rate: 5, Burst: 10, Cost: 1}, "report": {Rate: 5, Burst: 10, Cost: 5}}
}

func TestBurstThenRefill(t *testing.T) {
	l := NewLimiter(policies())
	for i := 0; i < 10; i++ {
		if !l.Allow("a", "default", t0).Allowed {
			t.Fatalf("request %d was refused inside the burst", i)
		}
	}
	d := l.Allow("a", "default", t0)
	if d.Allowed || d.RetryAfter <= 0 || d.RetryAfter > 250*time.Millisecond {
		t.Fatalf("%+v", d)
	}
	if !l.Allow("a", "default", t0.Add(time.Second)).Allowed {
		t.Fatal("tokens should have refilled after a second")
	}
}

func TestClientsDoNotShareABucket(t *testing.T) {
	l := NewLimiter(policies())
	for i := 0; i < 20; i++ {
		l.Allow("noisy", "default", t0)
	}
	if !l.Allow("quiet", "default", t0).Allowed {
		t.Fatal("one client's traffic limited another")
	}
}

func TestRefusedRequestsAreNotCharged(t *testing.T) {
	l := NewLimiter(policies())
	for i := 0; i < 10; i++ {
		l.Allow("a", "default", t0)
	}
	for i := 0; i < 1000; i++ { // a client hammering a closed door must not push its own recovery further away
		l.Allow("a", "default", t0)
	}
	if !l.Allow("a", "default", t0.Add(300*time.Millisecond)).Allowed {
		t.Fatal("refused requests were charged against the bucket")
	}
}

func TestACostLargerThanTheBurstNeverSucceeds(t *testing.T) {
	l := NewLimiter(map[string]Policy{"default": {Rate: 1, Burst: 2, Cost: 3}})
	if l.Allow("a", "default", t0.Add(time.Hour)).Allowed {
		t.Fatal("a request costing more than the bucket holds was allowed")
	}
}

func TestConcurrentCallersNeverExceedTheBurst(t *testing.T) {
	l := NewLimiter(policies())
	var allowed atomic.Int64
	var wg sync.WaitGroup
	for i := 0; i < 200; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if l.Allow("a", "default", t0).Allowed {
				allowed.Add(1)
			}
		}()
	}
	wg.Wait()
	if allowed.Load() != 10 {
		t.Fatalf("%d allowed, want exactly the burst of 10", allowed.Load())
	}
}

func TestSweepRemovesOnlyIdleBuckets(t *testing.T) {
	l := NewLimiter(policies())
	l.Allow("old", "default", t0)
	l.Allow("fresh", "default", t0.Add(9*time.Minute))
	if n := l.Sweep(t0.Add(11*time.Minute), 10*time.Minute); n != 1 || l.Size() != 1 {
		t.Fatalf("removed %d, %d left", n, l.Size())
	}
}

=== RUN   TestBurstThenRefill
--- PASS: TestBurstThenRefill (0.00s)
=== RUN   TestClientsDoNotShareABucket
--- PASS: TestClientsDoNotShareABucket (0.00s)
=== RUN   TestRefusedRequestsAreNotCharged
--- PASS: TestRefusedRequestsAreNotCharged (0.00s)
=== RUN   TestACostLargerThanTheBurstNeverSucceeds
--- PASS: TestACostLargerThanTheBurstNeverSucceeds (0.00s)
=== RUN   TestConcurrentCallersNeverExceedTheBurst
--- PASS: TestConcurrentCallersNeverExceedTheBurst (0.00s)
=== RUN   TestSweepRemovesOnlyIdleBuckets
--- PASS: TestSweepRemovesOnlyIdleBuckets (0.00s)
PASS
ok  	example.test/lab	1.294s

package main

import (
	"fmt"
	"testing"
)

func BenchmarkAllowOneClient(b *testing.B) {
	l := NewLimiter(policies())
	b.ReportAllocs()
	for b.Loop() {
		l.Allow("a", "default", t0)
	}
}

func BenchmarkAllowManyClientsParallel(b *testing.B) {
	l := NewLimiter(policies())
	keys := make([]string, 1024)
	for i := range keys {
		keys[i] = fmt.Sprint("client-", i)
	}
	b.ReportAllocs()
	b.RunParallel(func(pb *testing.PB) {
		i := 0
		for pb.Next() {
			l.Allow(keys[i%len(keys)], "default", t0)
			i++
		}
	})
}

goos: darwin
goarch: arm64
pkg: example.test/lab
cpu: Apple M1 Max
BenchmarkAllowOneClient-10              	 8732217	       122.2 ns/op	      16 B/op	       1 allocs/op
BenchmarkAllowManyClientsParallel-10    	 6012790	       198.2 ns/op	      23 B/op	       1 allocs/op
PASS
ok  	example.test/lab	2.726s