job-search/internal/middleware/ratelimit.go

package middleware

import (
	"context"
	"hash/maphash"
	"math"
	"net/http"
	"strconv"
	"sync"
	"time"

	"golang.org/x/time/rate"

	"git.avsdev.uk/Go/job-search/internal/httpx"
)

// Limiter is the rate limiting hook. Implementations decide whether a request
// identified by key may proceed. Swap in a Redis/memcached-backed limiter for
// multi-instance deployments by implementing this interface.
type Limiter interface {
	Allow(key string) (allowed bool, retryAfter time.Duration)
}

// LimiterFunc adapts a function to the Limiter interface.
type LimiterFunc func(key string) (bool, time.Duration)

// Allow implements Limiter by calling f.
func (f LimiterFunc) Allow(key string) (bool, time.Duration) { return f(key) }

// KeyFunc derives the rate limit key from a request.
type KeyFunc func(*http.Request) string

// KeyByIP keys on the real client IP (after proxy processing).
func KeyByIP(r *http.Request) string { return ClientIP(r) }

// RateLimit turns a Limiter into middleware. Rejected requests get 429 with
// a Retry-After header. A nil limiter yields a pass-through middleware.
func RateLimit(limiter Limiter, key KeyFunc) func(http.Handler) http.Handler {
	if limiter == nil {
		return func(next http.Handler) http.Handler {
			return next
		}
	}
	if key == nil {
		key = KeyByIP
	}
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ok, retry := limiter.Allow(key(r))
			if !ok {
				seconds := int(math.Ceil(retry.Seconds()))
				if seconds < 1 {
					seconds = 1
				}
				w.Header().Set("Retry-After", strconv.Itoa(seconds))
				httpx.WriteProblem(
					w, r, http.StatusTooManyRequests, "rate limit exceeded",
				)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// KeyedLimiter is an in-memory token bucket per key, sharded to reduce lock
// contention, with idle-entry eviction.
type KeyedLimiter struct {
	limit  rate.Limit
	burst  int
	ttl    time.Duration
	seed   maphash.Seed
	shards [32]limiterShard
}

type limiterShard struct {
	mu      sync.Mutex
	entries map[string]*limiterEntry
}

type limiterEntry struct {
	limiter *rate.Limiter
	last    time.Time
}

// NewKeyedLimiter creates a limiter allowing requestsPerSecond sustained
// requests per second with the given burst. A janitor goroutine evicts
// entries idle for longer than ttl until ctx is cancelled.
func NewKeyedLimiter(
	ctx context.Context, requestsPerSecond float64, burst int, ttl time.Duration,
) *KeyedLimiter {
	limiter := &KeyedLimiter{
		limit: rate.Limit(requestsPerSecond),
		burst: burst,
		ttl:   ttl,
		seed:  maphash.MakeSeed(),
	}
	for i := range limiter.shards {
		limiter.shards[i].entries = make(map[string]*limiterEntry)
	}

	if ttl > 0 {
		go limiter.janitor(ctx)
	}

	return limiter
}

func (k *KeyedLimiter) shard(key string) *limiterShard {
	return &k.shards[maphash.String(k.seed, key)%uint64(len(k.shards))]
}

// Allow implements Limiter.
func (k *KeyedLimiter) Allow(key string) (bool, time.Duration) {
	now := time.Now()
	shard := k.shard(key)
	shard.mu.Lock()
	entry := shard.entries[key]
	if entry == nil {
		entry = &limiterEntry{
			limiter: rate.NewLimiter(k.limit, k.burst),
		}
		shard.entries[key] = entry
	}
	entry.last = now
	limiter := entry.limiter
	shard.mu.Unlock()

	reservation := limiter.ReserveN(now, 1)
	if !reservation.OK() {
		return false, time.Second
	}
	if delay := reservation.DelayFrom(now); delay > 0 {
		reservation.CancelAt(now)
		return false, delay
	}

	return true, 0
}

func (k *KeyedLimiter) janitor(ctx context.Context) {
	ticker := time.NewTicker(k.ttl / 2)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case now := <-ticker.C:
			for i := range k.shards {
				shard := &k.shards[i]
				shard.mu.Lock()
				for key, entry := range shard.entries {
					if now.Sub(entry.last) > k.ttl {
						delete(shard.entries, key)
					}
				}
				shard.mu.Unlock()
			}
		}
	}
}