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:
//
//	loadgen-http --target http://localhost:8080 \
//	  --api-key acme-001:prom-prod:s3cret \
//	  --mode normal --rate 100 --duration 30s
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"
)

func main() {
	var (
		target     = flag.String("target", "http://localhost:8080", "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 per instance")
		duration   = flag.Duration("duration", 30*time.Second, "total run time")
		conc       = flag.Int("concurrency", 16, "concurrent HTTP requests in flight")
		metrics    = flag.String("metrics", ":9091", "Prometheus metrics listen addr (empty to disable)")
		dedupePct  = flag.Int("dedupe-pct", 30, "percent of alerts sharing a dedupe_key (normal mode)")
		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]
	_ = company
	_ = source
	_ = secret

	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo}))

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

	limiter := newRateLimiter(*rate)

	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
				}
				if err := limiter.Wait(ctx); err != nil {
					return
				}
				a := mkAlert(*mode, *companyN, *dedupePct, *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.
	producerCtx, cancelProducer := context.WithTimeout(ctx, *duration)
	defer cancelProducer()
	go func() {
		ticker := time.NewTicker(time.Second / time.Duration(*rate+1))
		defer ticker.Stop()
		for {
			select {
			case <-producerCtx.Done():
				close(jobs)
				return
			case <-ticker.C:
				select {
				case jobs <- struct{}{}:
				default:
					// backpressure: drop rather than block
				}
			}
		}
	}()

	// Metrics endpoint.
	if *metrics != "" {
		go runMetrics(*metrics, &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, payloadB int) alert.Alert {
	// Cycle through a few company_ids; the source is fixed for the
	// run (loadgen is one source).
	companyID := fmt.Sprintf("acme-%03d", (rand.IntN(companies) + 1))
	severity := pickSeverity(mode)
	category := pickCategory(severity)

	// Dedupe: dedupe_pct % of alerts share a fixed "burst" key
	// within a minute. For the simpler "normal" mode, we just
	// include a stable dedupe key on a fraction of alerts.
	dk := ""
	if rand.IntN(100) < dedupePct {
		dk = fmt.Sprintf("burst:%s:probe", category)
	}

	// Synthetic data roughly `payloadB` bytes.
	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 {
	// 70/25/4/1 split for normal, even for stress
	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
}

// rateLimiter is a simple per-second pacer. Not a token bucket; for
// loadgen we want "send at most N per second" averaged, not bursty.
type rateLimiter struct {
	perSec int
	last   time.Time
	mu     sync.Mutex
}

func newRateLimiter(perSec int) *rateLimiter { return &rateLimiter{perSec: perSec} }

func (r *rateLimiter) Wait(ctx context.Context) error {
	if r.perSec <= 0 {
		return nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	now := time.Now()
	interval := time.Second / time.Duration(r.perSec)
	if d := interval - now.Sub(r.last); d > 0 {
		t := time.NewTimer(d)
		defer t.Stop()
		r.mu.Unlock()
		select {
		case <-ctx.Done():
			r.mu.Lock()
			return ctx.Err()
		case <-t.C:
		}
		r.mu.Lock()
	}
	r.last = time.Now()
	return nil
}

func runMetrics(addr 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 %d\n", sent.Load())
		fmt.Fprintf(w, "loadgen_alerts_failed_total %d\n", failed.Load())
		fmt.Fprintf(w, "loadgen_dedupe_hits_total %d\n", dupes.Load())
		fmt.Fprintf(w, "loadgen_rate_limited_total %d\n", 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
}