broad-announce/loadgen/cmd/http/main.go

// loadgen/cmd/http is the HTTP POST traffic generator for M0.
// --mode normal only for now; other profiles land as the protocols
// come online.
//
// Example (single instance):
//
//	loadgen-http --target http://localhost:8800 \
//	  --api-key acme-001:prom-prod:s3cret \
//	  --mode normal --rate 100 --duration 30s --ramp-up 10s
//
// Example (3-instance cluster hitting 5k/s):
//
//	loadgen-http-1 --rate 1700 --duration 10m --ramp-up 30s --cluster-id m10-run
//	loadgen-http-2 --rate 1700 --duration 10m --ramp-up 30s --cluster-id m10-run
//	loadgen-http-3 --rate 1700 --duration 10m --ramp-up 30s --cluster-id m10-run
package main

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"math/rand/v2"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"git3.techno-world.net/lrosales/broad-announce/internal/alert"
	"git3.techno-world.net/lrosales/broad-announce/loadgen/internal/pacer"
)

func main() {
	var (
		target     = flag.String("target", "http://localhost:8800", "ingestd base URL")
		apiKey     = flag.String("api-key", "", "company_id:source_id:secret")
		mode       = flag.String("mode", "normal", "profile: normal|burst|stress")
		rate       = flag.Int("rate", 100, "target alerts/sec (after ramp-up)")
		duration   = flag.Duration("duration", 30*time.Second, "total run time")
		rampUp     = flag.Duration("ramp-up", 0, "linear ramp from 0 to --rate over this duration (default 0 = instant)")
		conc       = flag.Int("concurrency", 16, "concurrent HTTP requests in flight")
		metrics    = flag.String("metrics", ":8891", "Prometheus metrics listen addr (empty to disable)")
		clusterID  = flag.String("cluster-id", "default", "tag added as label on all loadgen metrics")
		instance   = flag.String("instance", "default", "loadgen instance name (e.g. loadgen-http-1)")
		dedupePct  = flag.Int("dedupe-pct", 30, "percent of alerts sharing a dedupe_key (normal mode)")
		dedupeKey  = flag.String("dedupe-key", "", "force a specific dedupe_key on every alert (overrides --dedupe-pct)")
		companyN   = flag.Int("companies", 1, "how many fake companies to cycle through")
		payloadB   = flag.Int("payload-bytes", 256, "approximate payload size in bytes (Data map)")
	)
	flag.Parse()
	if *apiKey == "" {
		fmt.Fprintln(os.Stderr, "loadgen-http: --api-key is required (company:source:secret)")
		os.Exit(2)
	}
	parts := strings.SplitN(*apiKey, ":", 3)
	if len(parts) != 3 {
		fmt.Fprintln(os.Stderr, "loadgen-http: --api-key must be company:source:secret")
		os.Exit(2)
	}
	company, source, secret := parts[0], parts[1], parts[2]

	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo}))
	logger.Info("starting",
		"target", *target,
		"rate", *rate,
		"ramp_up", *rampUp,
		"duration", *duration,
		"concurrency", *conc,
		"cluster_id", *clusterID,
		"instance", *instance,
	)

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	var (
		sent   atomic.Uint64
		failed atomic.Uint64
		dupes  atomic.Uint64
		rlHits atomic.Uint64
	)

	// Spawn workers.
	var wg sync.WaitGroup
	jobs := make(chan struct{}, *conc*4)
	for i := 0; i < *conc; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			client := &http.Client{Timeout: 10 * time.Second}
			for range jobs {
				if ctx.Err() != nil {
					return
				}
				a := mkAlert(*mode, *companyN, *dedupePct, *dedupeKey, *payloadB)
				if err := sendOne(ctx, client, *target, company, source, []byte(secret), a); err != nil {
					failed.Add(1)
					if isRateLimited(err) {
						rlHits.Add(1)
					}
					logger.Debug("send failed", "err", err)
				} else {
					sent.Add(1)
					if a.DedupeCount > 1 {
						dupes.Add(1)
					}
				}
			}
		}(i)
	}

	// Producer with pacer (supports linear ramp-up).
	producerCtx, cancelProducer := context.WithTimeout(ctx, *duration)
	defer cancelProducer()

	p := pacer.New(*rate, *rampUp)
	tickCh, stopPacer := p.Tick(producerCtx)
	defer stopPacer()

	go func() {
		for {
			select {
			case <-producerCtx.Done():
				close(jobs)
				return
			case <-tickCh:
				select {
				case jobs <- struct{}{}:
				default:
					// backpressure: drop rather than block
				}
			}
		}
	}()

	// Metrics endpoint.
	if *metrics != "" {
		go runMetrics(*metrics, *clusterID, *instance, *mode, &sent, &failed, &dupes, &rlHits)
	}

	// Status ticker.
	go func() {
		t := time.NewTicker(2 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				logger.Info("progress",
					"sent", sent.Load(),
					"failed", failed.Load(),
					"dupes", dupes.Load(),
					"rate_limited", rlHits.Load(),
				)
			}
		}
	}()

	wg.Wait()
	logger.Info("done",
		"sent", sent.Load(),
		"failed", failed.Load(),
		"dupes", dupes.Load(),
		"rate_limited", rlHits.Load(),
	)
}

func mkAlert(mode string, companies, dedupePct int, dedupeKey string, payloadB int) alert.Alert {
	companyID := fmt.Sprintf("acme-%03d", (rand.IntN(companies) + 1))
	severity := pickSeverity(mode)
	category := pickCategory(severity)

	dk := ""
	if dedupeKey != "" {
		dk = dedupeKey
	} else if rand.IntN(100) < dedupePct {
		dk = fmt.Sprintf("burst:%s:probe", category)
	}

	data := map[string]string{
		"host":    fmt.Sprintf("host-%d", rand.IntN(100)),
		"probe":   category,
		"raw_msg": strings.Repeat("x", max(0, payloadB-64)),
	}

	return alert.Alert{
		CompanyID: companyID,
		SourceID:  "prom-prod",
		Severity:  severity,
		Category:  category,
		Title:     fmt.Sprintf("%s on %s", category, data["host"]),
		Body:      "synthetic loadgen alert",
		Data:      data,
		DedupeKey: dk,
	}
}

func pickSeverity(mode string) alert.Severity {
	r := rand.IntN(100)
	switch {
	case r < 70:
		return alert.SeverityInfo
	case r < 95:
		return alert.SeverityWarning
	case r < 99:
		return alert.SeverityCritical
	default:
		return alert.SeverityInminentColapse
	}
}

var categoriesBySev = map[alert.Severity][]string{
	alert.SeverityInfo:            {"deploy", "schedule", "audit"},
	alert.SeverityWarning:         {"disk", "memory", "latency", "queue"},
	alert.SeverityCritical:        {"storage", "network", "process"},
	alert.SeverityInminentColapse: {"power", "hvac", "rack"},
}

func pickCategory(s alert.Severity) string {
	opts := categoriesBySev[s]
	return opts[rand.IntN(len(opts))]
}

func sendOne(ctx context.Context, c *http.Client, target, company, source string, secret []byte, a alert.Alert) error {
	body, err := json.Marshal(a)
	if err != nil {
		return err
	}
	ts := strconv.FormatInt(time.Now().Unix(), 10)
	mac := hmac.New(sha256.New, secret)
	mac.Write([]byte(ts))
	mac.Write([]byte("."))
	mac.Write(body)
	sig := "t=" + ts + ",v1=" + hex.EncodeToString(mac.Sum(nil))

	req, err := http.NewRequestWithContext(ctx, "POST", target+"/v1/ingest", bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-BA-Signature", sig)

	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusAccepted {
		var ar struct {
			DedupeCount uint32 `json:"dedupe_count"`
		}
		_ = json.NewDecoder(resp.Body).Decode(&ar)
		a.DedupeCount = ar.DedupeCount
		return nil
	}
	b, _ := io.ReadAll(resp.Body)
	if resp.StatusCode == http.StatusTooManyRequests {
		return &rateLimitedErr{status: resp.StatusCode, body: string(b)}
	}
	return fmt.Errorf("status %d: %s", resp.StatusCode, string(b))
}

type rateLimitedErr struct {
	status int
	body   string
}

func (e *rateLimitedErr) Error() string { return fmt.Sprintf("rate limited: %s", e.body) }
func isRateLimited(err error) bool {
	_, ok := err.(*rateLimitedErr)
	return ok
}

func runMetrics(addr, clusterID, instance, mode string, sent, failed, dupes, rlHits *atomic.Uint64) {
	mux := http.NewServeMux()
	mux.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "# HELP loadgen_alerts_sent_total Alerts successfully accepted.\n")
		fmt.Fprintf(w, "# TYPE loadgen_alerts_sent_total counter\n")
		fmt.Fprintf(w, "loadgen_alerts_sent_total{profile=%q,instance=%q,cluster_id=%q,mode=%q} %d\n",
			mode, instance, clusterID, mode, sent.Load())
		fmt.Fprintf(w, "loadgen_alerts_failed_total{reason=%q,instance=%q,cluster_id=%q} %d\n",
			"send", instance, clusterID, failed.Load())
		fmt.Fprintf(w, "loadgen_dedupe_hits_total{instance=%q,cluster_id=%q} %d\n",
			instance, clusterID, dupes.Load())
		fmt.Fprintf(w, "loadgen_rate_limited_total{instance=%q,cluster_id=%q} %d\n",
			instance, clusterID, rlHits.Load())
	})
	srv := &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	_ = srv.ListenAndServe()
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}