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"
	"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")
		mode      = flag.String("mode", "normal", "profile: normal|burst")
		dedupePct = flag.Int("dedupe-pct", 30, "percent sharing a dedupe_key (normal)")
		timeout   = flag.Duration("duration", 30*time.Second, "max run time")
		concurrencyFlag = flag.Int("concurrency", 1, "parallel WS connections (each one is one source)")
	)
	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}))
	_, 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, 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)
	}
}

// makeAlert is the same shape as loadgen/cmd/http and
// loadgen/cmd/mqtt.
func makeAlert(idx int, mode string, dedupePct int, company, source string) map[string]any {
	severity := pickSeverity(mode)
	dedupeKey := fmt.Sprintf("lg-m5-%d", idx)
	if dedupePct > 0 && idx > 0 && rand.IntN(100) < dedupePct {
		dedupeKey = "lg-m5-shared"
	}
	return map[string]any{
		"company_id": company,
		"source_id":  source,
		"severity":   severity,
		"category":   "loadgen",
		"title":      fmt.Sprintf("LG M5 #%d", idx),
		"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
}