Add global-ratelimit experiment: token-bucket limiter as an HTTP service

limiter/ package (ex03 bucket + ex09 sharded keyed store + idle sweep, race-tested), server/ with /allow /stats /healthz and X-RateLimit headers, client/ demo replica with fail-open/closed and latency percentiles. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GbnJj9ZqEx9nUvzkjdwfMq
2026-08-29 18:53:09 +02:00
parent 9230d0cd3d
commit 1610cd69e0
5 changed files with 757 additions and 0 deletions
@@ -0,0 +1,182 @@
+// Global rate limiter service.
+//
+// One process holds per-key token buckets; any number of downstream service
+// replicas ask it "may I?" before doing work, so the quota is enforced
+// globally rather than per replica.
+//
+//	go run ./99-experiments/global-ratelimit/server -rate 10 -burst 10
+//	curl -i 'localhost:8100/allow?key=tenant-a'
+//	curl -s  'localhost:8100/stats'
+//
+// API
+//
+//	POST|GET /allow?key=K[&n=1]
+//	    200 {"allowed":true,  "remaining":9, "limit":10, "rate":10}
+//	    429 {"allowed":false, "remaining":0.3, "limit":10, "rate":10, "retry_after_ms":70}
+//	    headers: X-RateLimit-Limit, X-RateLimit-Remaining, Retry-After (429 only)
+//	    400 when key is missing or n is not a positive integer
+//	GET /stats   → {"keys":N,"allowed_total":A,"denied_total":D}
+//	GET /healthz → 200 ok
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"flag"
+	"fmt"
+	"log"
+	"math"
+	"net/http"
+	"os"
+	"os/signal"
+	"strconv"
+	"strings"
+	"time"
+
+	"gosre-exercises/99-experiments/global-ratelimit/limiter"
+)
+
+type allowResponse struct {
+	Allowed      bool    `json:"allowed"`
+	Remaining    float64 `json:"remaining"`
+	Limit        int     `json:"limit"`
+	Rate         float64 `json:"rate"`
+	RetryAfterMS int64   `json:"retry_after_ms,omitempty"`
+}
+
+func allowHandler(store *limiter.Store) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		key := r.URL.Query().Get("key")
+		if key == "" {
+			http.Error(w, `missing "key" query parameter`, http.StatusBadRequest)
+			return
+		}
+		n := 1
+		if raw := r.URL.Query().Get("n"); raw != "" {
+			v, err := strconv.Atoi(raw)
+			if err != nil || v <= 0 {
+				http.Error(w, `"n" must be a positive integer`, http.StatusBadRequest)
+				return
+			}
+			n = v
+		}
+
+		d := store.Allow(key, n)
+
+		w.Header().Set("Content-Type", "application/json")
+		w.Header().Set("X-RateLimit-Limit", strconv.Itoa(d.Limit))
+		w.Header().Set("X-RateLimit-Remaining", strconv.Itoa(int(math.Floor(d.Remaining))))
+		resp := allowResponse{Allowed: d.Allowed, Remaining: d.Remaining, Limit: d.Limit, Rate: d.Rate}
+		status := http.StatusOK
+		if !d.Allowed {
+			status = http.StatusTooManyRequests
+			resp.RetryAfterMS = d.RetryAfter.Milliseconds()
+			// Retry-After is whole seconds; round UP so clients don't retry too early.
+			w.Header().Set("Retry-After", strconv.Itoa(int(math.Ceil(d.RetryAfter.Seconds()))))
+		}
+		w.WriteHeader(status)
+		json.NewEncoder(w).Encode(resp)
+	}
+}
+
+func statsHandler(store *limiter.Store) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(store.Stats())
+	}
+}
+
+// parseLimits parses "key=rate:burst,key2=rate:burst".
+func parseLimits(s string) (map[string]limiter.Limit, error) {
+	out := map[string]limiter.Limit{}
+	if strings.TrimSpace(s) == "" {
+		return out, nil
+	}
+	for _, item := range strings.Split(s, ",") {
+		key, spec, ok := strings.Cut(strings.TrimSpace(item), "=")
+		if !ok {
+			return nil, fmt.Errorf("bad limit %q: want key=rate:burst", item)
+		}
+		rateS, burstS, ok := strings.Cut(spec, ":")
+		if !ok {
+			return nil, fmt.Errorf("bad limit %q: want key=rate:burst", item)
+		}
+		rate, err1 := strconv.ParseFloat(rateS, 64)
+		burst, err2 := strconv.Atoi(burstS)
+		if err1 != nil || err2 != nil || rate <= 0 || burst <= 0 {
+			return nil, fmt.Errorf("bad limit %q: rate and burst must be positive", item)
+		}
+		out[key] = limiter.Limit{Rate: rate, Burst: burst}
+	}
+	return out, nil
+}
+
+func main() {
+	addr := flag.String("addr", ":8100", "listen address")
+	rate := flag.Float64("rate", 10, "default tokens/second for keys without an override")
+	burst := flag.Int("burst", 20, "default burst for keys without an override")
+	limits := flag.String("limits", "", `per-key overrides, e.g. "tenant-a=100:200,tenant-b=5:5"`)
+	sweep := flag.Duration("sweep", time.Minute, "evict buckets idle for at least this long (0 = never)")
+	flag.Parse()
+
+	overrides, err := parseLimits(*limits)
+	if err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(2)
+	}
+	store := limiter.NewStore(limiter.Limit{Rate: *rate, Burst: *burst}, overrides)
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/allow", allowHandler(store))
+	mux.HandleFunc("/stats", statsHandler(store))
+	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintln(w, "ok") })
+
+	srv := &http.Server{
+		Addr:    *addr,
+		Handler: mux,
+		// Never run a server with zero timeouts (see 99-experiments/body-drain).
+		ReadHeaderTimeout: 5 * time.Second,
+		WriteTimeout:      5 * time.Second,
+		IdleTimeout:       60 * time.Second,
+	}
+
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
+	defer stop()
+
+	// Idle-key eviction: the "map grows forever" fix. Also a periodic stats line.
+	go func() {
+		if *sweep <= 0 {
+			return
+		}
+		t := time.NewTicker(*sweep)
+		defer t.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-t.C:
+				removed := store.Sweep(*sweep)
+				st := store.Stats()
+				log.Printf("sweep: evicted=%d keys=%d allowed=%d denied=%d", removed, st.Keys, st.Allowed, st.Denied)
+			}
+		}
+	}()
+
+	go func() {
+		log.Printf("global rate limiter listening on %s (default %v/s burst %d, %d overrides)", *addr, *rate, *burst, len(overrides))
+		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Fatal(err)
+		}
+	}()
+
+	<-ctx.Done()
+	log.Println("shutting down")
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := srv.Shutdown(shutdownCtx); err != nil {
+		log.Printf("shutdown: %v", err)
+	}
+	st := store.Stats()
+	log.Printf("final: keys=%d allowed=%d denied=%d", st.Keys, st.Allowed, st.Denied)
+}