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

// loadgen/cmd/mqtt is the M4 MQTT publisher for broad-announce.
// Same data shape as loadgen/cmd/http (severity mix, dedupe
// ratio, company cycling) but posts to an MQTT broker on
// ba/<co>/<src>/incoming. Auth is twofold:
//
//   1. MQTT username/password: prom-prod-acme-001 / s3cret-acme
//      (the EMQX built-in-db row from deploy/emqx/auth-built-in-
//      db-bootstrap.csv). The broker's ACL is keyed on the
//      username, so this is what gates topic access.
//   2. Per-message X-BA-Signature in the JSON envelope's `auth`
//      field, identical to HTTP. This is the per-message HMAC
//      the ingestd mqtt subscriber re-uses verifyHMAC() to check.
//
// Example:
//
//	loadgen-mqtt --broker tcp://localhost:1883 \
//	  --api-key acme-001:prom-prod:s3cret-acme \
//	  --count 10 --rate 5
package main

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

	"git3.techno-world.net/lrosales/broad-announce/internal/mqttclient"
)

func main() {
	var (
		broker    = flag.String("broker", "tcp://localhost:1883", "MQTT broker URL")
		apiKey    = flag.String("api-key", "", "company:source:secret")
		count     = flag.Int("count", 10, "total alerts to send")
		rate      = flag.Int("rate", 10, "target alerts/sec")
		rampUp    = flag.Duration("ramp-up", 0, "linear ramp from 0 to --rate over this duration (default 0 = instant)")
		mode      = flag.String("mode", "normal", "profile: normal|burst")
		dedupePct = flag.Int("dedupe-pct", 30, "percent sharing a dedupe_key (normal)")
		dedupeKey = flag.String("dedupe-key", "", "force a specific dedupe_key on every alert")
		metrics   = flag.String("metrics", "", "Prometheus metrics listen addr (empty to disable)")
		timeout   = flag.Duration("duration", 30*time.Second, "max run time")
		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-mqtt-1)")
	)
	flag.Parse()
	if *apiKey == "" {
		fmt.Fprintln(os.Stderr, "loadgen-mqtt: --api-key is required (company:source:secret)")
		os.Exit(2)
	}
	parts := strings.SplitN(*apiKey, ":", 3)
	if len(parts) != 3 {
		fmt.Fprintln(os.Stderr, "loadgen-mqtt: --api-key must be company:source:secret")
		os.Exit(2)
	}
	company, source, secret := parts[0], parts[1], parts[2]
	username := fmt.Sprintf("%s-%s", source, company) // matches auth csv

	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelInfo}))
	logger.Info("starting",
		"broker", *broker,
		"rate", *rate,
		"ramp_up", *rampUp,
		"count", *count,
		"duration", *timeout,
		"cluster_id", *clusterID,
		"instance", *instance,
	)
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	host, _ := os.Hostname()
	client, err := mqttclient.Connect(ctx, mqttclient.Config{
		Broker:   *broker,
		ClientID: fmt.Sprintf("loadgen-mqtt-%s", host),
		Username: username,
		Password: secret,
		Clean:    true,
	}, logger)
	if err != nil {
		fmt.Fprintln(os.Stderr, "loadgen-mqtt: connect:", err)
		os.Exit(1)
	}
	defer client.Disconnect()

	topic := fmt.Sprintf("ba/%s/%s/incoming", company, source)
	logger.Info("publishing", "topic", topic, "count", *count, "rate", *rate, "mode", *mode)

	var (
		sent   atomic.Uint64
		failed atomic.Uint64
		dupes  atomic.Uint64
	)
	limiter := time.NewTicker(time.Second / time.Duration(*rate))
	defer limiter.Stop()

	var wg sync.WaitGroup
	work := make(chan int, 64)
	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < *count; i++ {
			work <- i
		}
		close(work)
	}()

	deadline := time.Now().Add(*timeout)
	for i := range work {
		if time.Now().After(deadline) {
			logger.Warn("deadline reached, stopping")
			break
		}
		<-limiter.C

		a := makeAlert(i, *mode, *dedupePct, *dedupeKey, company, source)
		ts := time.Now().Unix()
		// HTTP-style signature over the alert body, with t= prefix.
		// The MQTT subscriber's verifyHMAC accepts the same shape.
		body, _ := json.Marshal(a)
		mac := hmac.New(sha256.New, []byte(secret))
		mac.Write([]byte(fmt.Sprintf("%d", ts)))
		mac.Write([]byte("."))
		mac.Write(body)
		sig := hex.EncodeToString(mac.Sum(nil))
		env := map[string]json.RawMessage{
			"alert": body,
		}
		env["auth"] = json.RawMessage(fmt.Sprintf("%q", fmt.Sprintf("t=%d,v1=%s", ts, sig)))
		envelope, _ := json.Marshal(env)

		if err := client.Publish(topic, envelope); err != nil {
			failed.Add(1)
			logger.Warn("publish failed", "err", err, "i", i)
			continue
		}
		if i > 0 && i%(*count/10+1) == 0 {
			logger.Info("progress", "sent", i, "total", *count)
		}
		if !*isUnique(*dedupePct, i) {
			dupes.Add(1)
		}
		sent.Add(1)
	}
	wg.Wait()

	logger.Info("done",
		"sent", sent.Load(),
		"failed", failed.Load(),
		"dupes", dupes.Load(),
	)
	if failed.Load() > 0 {
		os.Exit(1)
	}
	// Metrics endpoint (registered but currently a no-op; wired in M10)
	if *metrics != "" {
		go runMQTTMetrics(*metrics, *clusterID, *instance, &sent, &failed, &dupes)
	}
}

// makeAlert generates one alert. Same shape as loadgen/cmd/http.
func makeAlert(idx int, mode string, dedupePct int, forceKey, company, source string) map[string]any {
	severity := pickSeverity(mode)
	// M6: --dedupe-key forces a specific key on every alert,
	// which is what the ×N smoke test needs. Otherwise we
	// use the M4 distribution: 30% share a key, 70% get a
	// unique one.
	dedupeKey := fmt.Sprintf("lg-m4-%d", idx)
	if forceKey != "" {
		dedupeKey = forceKey
	} else if dedupePct > 0 && idx > 0 && rand.IntN(100) < dedupePct {
		dedupeKey = "lg-m4-shared"
	}
	return map[string]any{
		"company_id": company,
		"source_id":  source,
		"severity":   severity,
		"category":   "loadgen",
		"title":      fmt.Sprintf("LG M4 #%d", idx),
		"body":       "mqtt smoke",
		"data":       map[string]string{"host": "lg-host", "idx": fmt.Sprintf("%d", idx)},
		"dedupe_key": dedupeKey,
	}
}

func pickSeverity(mode string) string {
	r := rand.IntN(100)
	switch mode {
	case "burst":
		// Mostly critical to exercise the bypass.
		switch {
		case r < 70:
			return "critical"
		case r < 95:
			return "inminent_colapse"
		default:
			return "warning"
		}
	default: // "normal"
		switch {
		case r < 70:
			return "info"
		case r < 95:
			return "warning"
		case r < 99:
			return "critical"
		default:
			return "inminent_colapse"
		}
	}
}

func isUnique(dedupePct, i int) *bool {
	b := i == 0 || dedupePct == 0 || rand.IntN(100) >= dedupePct
	return &b
}

func runMQTTMetrics(addr, clusterID, instance string, sent, failed, dupes *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{instance=%q,cluster_id=%q} %d\n",
			instance, clusterID, sent.Load())
		fmt.Fprintf(w, "loadgen_alerts_failed_total{instance=%q,cluster_id=%q} %d\n",
			instance, clusterID, failed.Load())
		fmt.Fprintf(w, "loadgen_dedupe_hits_total{instance=%q,cluster_id=%q} %d\n",
			instance, clusterID, dupes.Load())
	})
	srv := &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	_ = srv.ListenAndServe()
}