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

// Demo downstream service replica that consults the global rate limiter
// before each unit of work.
//
// Run several in parallel with different -id against one server and watch
// the SUM of their allowed rates stay at the key's limit:
//
//	go run ./99-experiments/global-ratelimit/client -id inst-1 -rps 20
//	go run ./99-experiments/global-ratelimit/client -id inst-2 -rps 20
//	go run ./99-experiments/global-ratelimit/client -id inst-3 -rps 20
//
// Then Ctrl-C the server mid-run and compare -fail open vs -fail closed.
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"sort"
	"time"
)

type allowResponse struct {
	Allowed      bool  `json:"allowed"`
	RetryAfterMS int64 `json:"retry_after_ms"`
}

type outcome int

const (
	allowed outcome = iota
	denied
	errored // limiter unreachable / timed out / bad response
)

// ask performs one /allow call. It is the ONLY place that talks to the limiter.
func ask(ctx context.Context, client *http.Client, endpoint string) (outcome, time.Duration) {
	start := time.Now()
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, nil)
	if err != nil {
		return errored, time.Since(start)
	}
	resp, err := client.Do(req)
	if err != nil {
		return errored, time.Since(start)
	}
	defer resp.Body.Close()
	// Decode fully (which also drains the body → connection reuse).
	var ar allowResponse
	if err := json.NewDecoder(resp.Body).Decode(&ar); err != nil {
		io.Copy(io.Discard, resp.Body)
		return errored, time.Since(start)
	}
	switch resp.StatusCode {
	case http.StatusOK:
		return allowed, time.Since(start)
	case http.StatusTooManyRequests:
		return denied, time.Since(start)
	default:
		return errored, time.Since(start)
	}
}

func percentile(sorted []time.Duration, p float64) time.Duration {
	if len(sorted) == 0 {
		return 0
	}
	i := int(float64(len(sorted)-1) * p)
	return sorted[i]
}

func main() {
	id := flag.String("id", "inst-1", "this replica's name (for output only)")
	limiterURL := flag.String("limiter", "http://127.0.0.1:8100", "limiter service base URL")
	key := flag.String("key", "tenant-a", "rate-limit key to consume from")
	rps := flag.Float64("rps", 20, "how many units of work this replica ATTEMPTS per second")
	duration := flag.Duration("duration", 30*time.Second, "how long to run")
	failMode := flag.String("fail", "open", "what to do when the limiter is unreachable: open (proceed) | closed (deny)")
	timeout := flag.Duration("timeout", 200*time.Millisecond, "per-call timeout to the limiter")
	flag.Parse()

	if *failMode != "open" && *failMode != "closed" {
		fmt.Fprintln(os.Stderr, "-fail must be open or closed")
		os.Exit(2)
	}

	endpoint := *limiterURL + "/allow?" + url.Values{"key": {*key}}.Encode()

	// One shared client, with a timeout: a limiter that is slow must not make
	// US slow beyond `timeout`. Keep-alive means one TCP conn for the whole run.
	client := &http.Client{Timeout: *timeout}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	ctx, cancel := context.WithTimeout(ctx, *duration)
	defer cancel()

	tick := time.NewTicker(time.Duration(float64(time.Second) / *rps))
	defer tick.Stop()
	report := time.NewTicker(time.Second)
	defer report.Stop()

	var sec, total [3]int // per-second and cumulative counts by outcome
	var latencies []time.Duration
	workDone := 0 // units of work actually performed (allowed + fail-open errors)

	fmt.Printf("%s: attempting %.0f/s on key %q via %s, fail-%s, timeout %v\n",
		*id, *rps, *key, *limiterURL, *failMode, *timeout)

loop:
	for {
		select {
		case <-ctx.Done():
			break loop
		case <-report.C:
			fmt.Printf("%s  allowed=%-3d denied=%-3d errors=%-3d  (last 1s)\n", *id, sec[allowed], sec[denied], sec[errored])
			sec = [3]int{}
		case <-tick.C:
			o, lat := ask(ctx, client, endpoint)
			if ctx.Err() != nil {
				break loop // aborted by our own shutdown, not a limiter failure
			}
			sec[o]++
			total[o]++
			latencies = append(latencies, lat)
			// Do the actual "work" only if allowed — or on error, if fail-open.
			if o == allowed || (o == errored && *failMode == "open") {
				workDone++
			}
		}
	}

	sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] })
	n := total[allowed] + total[denied] + total[errored]
	fmt.Printf("\n%s summary: attempts=%d allowed=%d denied=%d errors=%d → work done=%d (fail-%s)\n",
		*id, n, total[allowed], total[denied], total[errored], workDone, *failMode)
	fmt.Printf("%s /allow latency: p50=%v p99=%v max=%v\n",
		*id, percentile(latencies, 0.50), percentile(latencies, 0.99), percentile(latencies, 1))
}