Token Bucket Rate-Limiting Middleware for a Go net/http API
package ratelimit
import (
"sync"
"time"
)
type Bucket struct {
mu sync.Mutex
capacity float64
tokens float64
refillPerSec float64
last time.Time
}
func NewBucket(capacity, refillPerSec float64) *Bucket {
return &Bucket{
capacity: capacity,
tokens: capacity,
refillPerSec: refillPerSec,
last: time.Now(),
}
}
func (b *Bucket) refill(now time.Time) {
elapsed := now.Sub(b.last).Seconds()
if elapsed <= 0 {
return
}
b.tokens += elapsed * b.refillPerSec
if b.tokens > b.capacity {
b.tokens = b.capacity
}
b.last = now
}
func (b *Bucket) Allow() (bool, float64) {
b.mu.Lock()
defer b.mu.Unlock()
b.refill(time.Now())
if b.tokens < 1 {
return false, b.tokens
}
b.tokens--
return true, b.tokens
}
func (b *Bucket) isFull() bool {
b.mu.Lock()
defer b.mu.Unlock()
return b.tokens >= b.capacity
}package ratelimit
import (
"sync"
"time"
)
type Limiter struct {
mu sync.Mutex
buckets map[string]*Bucket
capacity float64
refillPerSec float64
}
func NewLimiter(capacity, refillPerSec float64) *Limiter {
l := &Limiter{
buckets: make(map[string]*Bucket),
capacity: capacity,
refillPerSec: refillPerSec,
}
go l.cleanup(5 * time.Minute)
return l
}
func (l *Limiter) getBucket(key string) *Bucket {
l.mu.Lock()
defer l.mu.Unlock()
b, ok := l.buckets[key]
if !ok {
b = NewBucket(l.capacity, l.refillPerSec)
l.buckets[key] = b
}
return b
}
func (l *Limiter) Allow(key string) (bool, float64) {
return l.getBucket(key).Allow()
}
func (l *Limiter) cleanup(interval time.Duration) {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for range ticker.C {
l.mu.Lock()
for key, b := range l.buckets {
if b.isFull() {
delete(l.buckets, key)
}
}
l.mu.Unlock()
}
}package ratelimit
import (
"net"
"net/http"
"strconv"
)
type KeyFunc func(r *http.Request) string
func ClientIP(r *http.Request) string {
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
return fwd
}
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
return r.RemoteAddr
}
return host
}
func RateLimit(l *Limiter, keyFunc KeyFunc) func(http.Handler) http.Handler {
if keyFunc == nil {
keyFunc = ClientIP
}
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
key := keyFunc(r)
allowed, remaining := l.Allow(key)
w.Header().Set("X-RateLimit-Remaining", strconv.Itoa(int(remaining)))
if !allowed {
w.Header().Set("Retry-After", "1")
http.Error(w, "rate limit exceeded", http.StatusTooManyRequests)
return
}
next.ServeHTTP(w, r)
})
}
}package main
import (
"log"
"net/http"
"example.com/api/ratelimit"
)
func main() {
// 20 token burst, refilling at 5 req/sec per client
limiter := ratelimit.NewLimiter(20, 5)
mux := http.NewServeMux()
mux.HandleFunc("/v1/search", func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte(`{"results":[]}`))
})
guarded := ratelimit.RateLimit(limiter, ratelimit.ClientIP)(mux)
log.Println("listening on :8080")
if err := http.ListenAndServe(":8080", guarded); err != nil {
log.Fatal(err)
}
}This snippet implements per-client rate limiting for a Go HTTP API using the classic token bucket algorithm, wrapped as standard net/http middleware. A token bucket models a fixed-capacity bucket that refills at a steady rate: every request tries to remove one token, and when the bucket is empty the request is rejected. Because the bucket has capacity, it tolerates short bursts while still enforcing an average rate over time, which is usually what an API wants — smooth throughput without punishing the occasional spike.
The tokenbucket.go tab holds the pure algorithm, deliberately free of any HTTP concerns so it can be unit tested in isolation. Bucket stores its state as a floating-point tokens count rather than integers, which lets the refill be computed lazily: instead of running a background goroutine that ticks tokens in, refill looks at the elapsed time since the last call and adds elapsed * refillPerSec tokens, clamped to capacity. This lazy, on-demand refill is the key trade-off — it avoids one timer per client and scales to many buckets cheaply, at the cost of doing a little arithmetic on every Allow call. The mutex makes a single bucket safe under concurrent requests, and Allow returns false without blocking when fewer than one token remains.
The limiter.go tab manages a map of buckets keyed by client identity so that one noisy client cannot exhaust another's budget. Limiter guards the map with its own sync.Mutex and lazily creates a bucket the first time a key is seen in getBucket. A real deployment would also evict idle buckets to bound memory; the cleanup loop shows the shape of that, sweeping entries whose bucket has been full and untouched for a while. Keying by IP is simple but imperfect behind proxies, so the middleware allows the key function to be injected.
The middleware.go tab bridges the limiter to net/http. RateLimit is a higher-order function returning a func(http.Handler) http.Handler, the idiomatic middleware signature that composes cleanly with routers like chi or the standard ServeMux. It derives a key via keyFunc, calls Allow, and on rejection writes standards-friendly headers: Retry-After tells clients when to try again and X-RateLimit-Remaining reports the budget, then returns 429 Too Many Requests. When a request is allowed it simply calls next.ServeHTTP, so allowed traffic pays almost no overhead.
A subtle pitfall worth noting is clock choice: the code uses time.Now() for elapsed time, which is fine in practice but sensitive to wall-clock jumps; a monotonic source avoids that. Another is that this limiter is per-process — behind a load balancer each instance keeps its own buckets, so the effective global rate is multiplied by the instance count. For a single service or a sidecar this in-memory approach is fast and dependency-free; for a fleet a shared store like Redis with the same token-bucket math is the natural next step.
Share this code
Here's the card — post it anywhere.
