learning-go-for-sre-interview/99-experiments/global-ratelimit/server/main.go

// 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)
}