descrybe/scripts/api-load-smoke/main.go

// Command api-load-smoke: short, low-concurrency HTTP smoke against critical API paths.
//
// Default profile is intentionally gentle (local laptop / shared Postgres).
// It is NOT a capacity test — see README.md for larger profiles and RPS assumptions.
//
// Usage (API listening — default HTTP_ADDR :28471):
//
//	cd scripts/api-load-smoke && go run .
//	go run . -base http://127.0.0.1:28471 -c 2 -d 5s -rate 6
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const demoAPIKey = "dk_demo_local_descrybe_test_key_v1"

// Smoke defaults — capped aggregate RPS so local Postgres is not hammered.
const (
	defaultConcurrency = 2
	defaultDuration    = 5 * time.Second
	defaultTimeout     = 10 * time.Second
	defaultRateRPS     = 6.0 // aggregate across all workers/endpoints
	defaultLimitQuery  = "limit=10&offset=0"
)

type endpoint struct {
	name string
	path string
	auth bool
}

type stats struct {
	ok       atomic.Int64
	fail     atomic.Int64
	latMu    sync.Mutex
	latNanos []int64
}

func (s *stats) record(ok bool, d time.Duration) {
	if ok {
		s.ok.Add(1)
	} else {
		s.fail.Add(1)
	}
	s.latMu.Lock()
	s.latNanos = append(s.latNanos, d.Nanoseconds())
	s.latMu.Unlock()
}

func (s *stats) snapshot() (ok, fail int64, min, avg, p95, max time.Duration) {
	ok, fail = s.ok.Load(), s.fail.Load()
	s.latMu.Lock()
	defer s.latMu.Unlock()
	n := len(s.latNanos)
	if n == 0 {
		return ok, fail, 0, 0, 0, 0
	}
	cp := append([]int64(nil), s.latNanos...)
	sort.Slice(cp, func(i, j int) bool { return cp[i] < cp[j] })
	var sum int64
	for _, v := range cp {
		sum += v
	}
	min = time.Duration(cp[0])
	max = time.Duration(cp[n-1])
	avg = time.Duration(sum / int64(n))
	idx := (n * 95) / 100
	if idx >= n {
		idx = n - 1
	}
	p95 = time.Duration(cp[idx])
	return ok, fail, min, avg, p95, max
}

func main() {
	base := flag.String("base", envOr("API_BASE", "http://127.0.0.1:28471"), "API base URL (no trailing slash)")
	key := flag.String("key", envOr("API_KEY", demoAPIKey), "API key for /api/v1 feeds+products (Bearer)")
	c := flag.Int("c", defaultConcurrency, "worker concurrency (smoke: 2; load: raise carefully)")
	d := flag.Duration("d", defaultDuration, "test duration (smoke: 5s)")
	rate := flag.Float64("rate", defaultRateRPS, "max aggregate requests/sec (0 = unlimited; smoke default 6)")
	timeout := flag.Duration("timeout", defaultTimeout, "per-request timeout")
	mode := flag.String("mode", "smoke", "smoke|load (banner + safety hint; still honor -c/-d/-rate)")
	failOnErr := flag.Bool("fail-on-error", true, "exit 1 if any request fails")
	flag.Parse()

	if *c < 1 {
		fmt.Fprintln(os.Stderr, "-c must be >= 1")
		os.Exit(2)
	}
	if *d <= 0 {
		fmt.Fprintln(os.Stderr, "-d must be > 0")
		os.Exit(2)
	}
	if *rate < 0 {
		fmt.Fprintln(os.Stderr, "-rate must be >= 0")
		os.Exit(2)
	}
	if strings.EqualFold(*mode, "smoke") && *rate == 0 {
		fmt.Fprintln(os.Stderr, "refusing -rate 0 in smoke mode (would DDOS local stack); use -mode load -rate 0 intentionally")
		os.Exit(2)
	}

	baseURL := strings.TrimRight(strings.TrimSpace(*base), "/")
	endpoints := []endpoint{
		{name: "healthz", path: "/healthz", auth: false},
		{name: "feeds", path: "/api/v1/feeds?" + defaultLimitQuery, auth: true},
		{name: "products", path: "/api/v1/products?" + defaultLimitQuery, auth: true},
	}

	client := &http.Client{Timeout: *timeout}
	st := make([]*stats, len(endpoints))
	for i := range st {
		st[i] = &stats{}
	}

	ctx, cancel := context.WithTimeout(context.Background(), *d)
	defer cancel()

	tokens := startRateGate(ctx, *rate)

	fmt.Printf("api-load-smoke mode=%s base=%s c=%d d=%s rate=%.1f\n", *mode, baseURL, *c, d.String(), *rate)
	fmt.Println("ASSUMPTION: local smoke defaults to ~6 aggregate RPS (c=2, 5s); raise -rate/-c only on dedicated/staging.")
	if strings.EqualFold(*mode, "load") {
		fmt.Println("NOTE: prefer staging. GET lists are not under RateLimitV1Process (mutations only); DB is still the bottleneck.")
	}

	var wg sync.WaitGroup
	start := time.Now()
	for w := 0; w < *c; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			i := 0
			for {
				if !waitToken(ctx, tokens) {
					return
				}
				ep := endpoints[i%len(endpoints)]
				si := st[i%len(endpoints)]
				i++
				ok, lat, cancelled := doRequest(ctx, client, baseURL, *key, ep)
				if cancelled {
					return
				}
				si.record(ok, lat)
			}
		}()
	}
	wg.Wait()
	elapsed := time.Since(start)

	var totalOK, totalFail int64
	fmt.Println()
	fmt.Printf("%-10s %8s %8s %10s %10s %10s %10s\n", "endpoint", "ok", "fail", "min", "avg", "p95", "max")
	for i, ep := range endpoints {
		ok, fail, min, avg, p95, max := st[i].snapshot()
		totalOK += ok
		totalFail += fail
		fmt.Printf("%-10s %8d %8d %10s %10s %10s %10s\n",
			ep.name, ok, fail, fmtDur(min), fmtDur(avg), fmtDur(p95), fmtDur(max))
	}
	total := totalOK + totalFail
	rps := 0.0
	if elapsed > 0 {
		rps = float64(total) / elapsed.Seconds()
	}
	fmt.Printf("\ntotal_ok=%d total_fail=%d elapsed=%s approx_rps=%.1f\n", totalOK, totalFail, elapsed.Round(time.Millisecond), rps)

	if totalFail > 0 && *failOnErr {
		os.Exit(1)
	}
}

// startRateGate returns nil for unlimited; otherwise a token channel paced at rate RPS.
func startRateGate(ctx context.Context, rate float64) <-chan struct{} {
	if rate <= 0 {
		return nil
	}
	ch := make(chan struct{})
	interval := time.Duration(float64(time.Second) / rate)
	if interval < time.Millisecond {
		interval = time.Millisecond
	}
	go func() {
		t := time.NewTicker(interval)
		defer t.Stop()
		// First token immediately so smoke does not wait a full interval.
		select {
		case <-ctx.Done():
			return
		case ch <- struct{}{}:
		}
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				select {
				case <-ctx.Done():
					return
				case ch <- struct{}{}:
				}
			}
		}
	}()
	return ch
}

func waitToken(ctx context.Context, tokens <-chan struct{}) bool {
	if tokens == nil {
		return ctx.Err() == nil
	}
	select {
	case <-ctx.Done():
		return false
	case <-tokens:
		return true
	}
}

func doRequest(ctx context.Context, client *http.Client, base, key string, ep endpoint) (ok bool, lat time.Duration, cancelled bool) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+ep.path, nil)
	if err != nil {
		return false, 0, ctx.Err() != nil
	}
	req.Header.Set("Accept", "application/json")
	if ep.auth {
		req.Header.Set("Authorization", "Bearer "+key)
	}
	t0 := time.Now()
	res, err := client.Do(req)
	lat = time.Since(t0)
	if err != nil {
		if ctx.Err() != nil {
			return false, lat, true
		}
		fmt.Fprintf(os.Stderr, "%s error: %v\n", ep.name, err)
		return false, lat, false
	}
	defer res.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(res.Body, 1<<20))
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		fmt.Fprintf(os.Stderr, "%s status=%d\n", ep.name, res.StatusCode)
		return false, lat, false
	}
	return true, lat, false
}

func envOr(key, fallback string) string {
	if v := strings.TrimSpace(os.Getenv(key)); v != "" {
		return v
	}
	return fallback
}

func fmtDur(d time.Duration) string {
	if d <= 0 {
		return "-"
	}
	if d < time.Millisecond {
		return d.Round(time.Microsecond).String()
	}
	return d.Round(time.Millisecond).String()
}