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

// loadgen/cmd/ws is the M5 WebSocket publisher for
// broad-announce. Same data shape as loadgen/cmd/http and
// loadgen/cmd/mqtt, but talks the WS transport at
// ws://ingestd:8800/v1/ingest/ws. The auth model is the same
// as the HTTP path: the api-key is `company:source:secret`,
// sent as a JSON frame on connect. The per-message HMAC is the
// same X-BA-Signature as the HTTP path, embedded in the
// envelope's `auth` field.
//
// Example:
//
//	loadgen-ws --target ws://localhost:8800/v1/ingest/ws \
//	  --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/wsclient"
)

func main() {
	var (
		target    = flag.String("target", "ws://localhost:8800/v1/ingest/ws", "WebSocket endpoint 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 (M6 ×N smoke test)")
		timeout   = flag.Duration("duration", 30*time.Second, "max run time")
		concurrencyFlag = flag.Int("concurrency", 1, "parallel WS connections (each one is one source)")
		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-ws-1)")
		metricsAddr = flag.String("metrics", "", "Prometheus metrics listen addr (empty to disable)")
	)
	flag.Parse()
	if *apiKey == "" {
		fmt.Fprintln(os.Stderr, "loadgen-ws: --api-key is required (company:source:secret)")
		os.Exit(2)
	}
	parts := strings.SplitN(*apiKey, ":", 3)
	if len(parts) != 3 {
		fmt.Fprintln(os.Stderr, "loadgen-ws: --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,
		"count", *count,
		"concurrency", *concurrencyFlag,
		"cluster_id", *clusterID,
		"instance", *instance,
	)
	_, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	logger.Info("publishing",
		"target", *target,
		"count", *count, "rate", *rate, "mode", *mode,
		"concurrency", *concurrencyFlag,
	)

	var (
		sent    atomic.Uint64
		failed  atomic.Uint64
		dupes   atomic.Uint64
		idxCh   = make(chan int, *count)
	)
	for i := 0; i < *count; i++ {
		idxCh <- i
	}
	close(idxCh)

	limiter := time.NewTicker(time.Second / time.Duration(*rate))
	defer limiter.Stop()
	deadline := time.Now().Add(*timeout)

	var wg sync.WaitGroup
	for w := 0; w < *concurrencyFlag; w++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			c, err := wsclient.Connect(wsclient.Config{
				URL:    *target,
				APIKey: *apiKey,
			})
			if err != nil {
				logger.Error("ws connect", "worker", workerID, "err", err)
				failed.Add(1)
				return
			}
			defer c.Close()
			logger.Info("ws connected", "worker", workerID, "auth_reply", string(c.AuthReply()))
			for i := range idxCh {
				if time.Now().After(deadline) {
					return
				}
				<-limiter.C
				a := makeAlert(i, *mode, *dedupePct, *dedupeKey, company, source)
				body, _ := json.Marshal(a)
				ts := time.Now().Unix()
				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)
				ack, err := c.SendAlert(envelope)
				if err != nil {
					failed.Add(1)
					logger.Warn("ws send", "err", err, "i", i)
					continue
				}
				if !isUnique(*dedupePct, i) {
					dupes.Add(1)
				}
				sent.Add(1)
				if i == 0 || (i+1)%(*count/10+1) == 0 {
					logger.Info("progress", "worker", workerID, "sent", i+1, "total", *count, "ack", string(ack))
				}
			}
		}(w)
	}
	wg.Wait()
	logger.Info("done",
		"sent", sent.Load(),
		"failed", failed.Load(),
		"dupes", dupes.Load(),
	)
	if failed.Load() > 0 {
		os.Exit(1)
	}
	if *metricsAddr != "" {
		go runWSMetrics(*metricsAddr, *clusterID, *instance, &sent, &failed, &dupes)
	}
}

// makeAlert is the same shape as loadgen/cmd/http and
// loadgen/cmd/mqtt.
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 M5 distribution: 30% share a key, 70% get a
	// unique one.
	dedupeKey := fmt.Sprintf("lg-m5-%d", idx)
	if forceKey != "" {
		dedupeKey = forceKey
	} else if dedupePct > 0 && idx > 0 && rand.IntN(100) < dedupePct {
		dedupeKey = "lg-m5-shared"
	}
	title := fmt.Sprintf("LG M5 #%d", idx)
	if forceKey != "" {
		// M6.5 smoke: include the forced key in the title so
		// the tail subscriber can filter on a unique string
		// (the tail event includes title but not dedupe_key).
		title = fmt.Sprintf("LG M5 burst %s #%d", forceKey, idx)
	}
	return map[string]any{
		"company_id": company,
		"source_id":  source,
		"severity":   severity,
		"category":   "loadgen",
		"title":      title,
		"body":       "ws 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":
		switch {
		case r < 70:
			return "critical"
		case r < 95:
			return "inminent_colapse"
		default:
			return "warning"
		}
	default:
		switch {
		case r < 70:
			return "info"
		case r < 95:
			return "warning"
		case r < 99:
			return "critical"
		default:
			return "inminent_colapse"
		}
	}
}

func isUnique(dedupePct, idx int) bool {
	if dedupePct == 0 || idx == 0 {
		return true
	}
	return rand.IntN(100) >= dedupePct
}

func runWSMetrics(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()
}