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

// loadgen/cmd/grpc is the gRPC (StreamAlerts) traffic generator for M11.
// Each instance opens N parallel bidirectional streams (--workers).
//
// Example (single instance, 10k/s):
//
//	loadgen-grpc --target localhost:9090 \
//	  --api-key acme-001:prom-prod:s3cret \
//	  --rate 10000 --workers 8 --duration 10m
//
// Example (2-instance cluster, 10k/s total):
//
//	loadgen-grpc-1 --rate 5000 --workers 8 --duration 15m --instance loadgen-grpc-1
//	loadgen-grpc-2 --rate 5000 --workers 8 --duration 15m --instance loadgen-grpc-2
package main

import (
	"context"
	"flag"
	"fmt"
	"log/slog"
	"math/rand/v2"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	grpcclient "git3.techno-world.net/lrosales/broad-announce/internal/grpcclient"
	pbv1 "git3.techno-world.net/lrosales/broad-announce/gen/go/broadannounce/v1"
	"git3.techno-world.net/lrosales/broad-announce/internal/alert"
	"git3.techno-world.net/lrosales/broad-announce/loadgen/internal/pacer"
)

type workerResult struct {
	sent, failed, dupes, rlHits uint64
}

func main() {
	var (
		target    = flag.String("target", "localhost:9090", "ingestd gRPC address")
		apiKey    = flag.String("api-key", "", "company_id:source_id:secret")
		rate      = flag.Int("rate", 1000, "target alerts/sec total (across all workers)")
		workers   = flag.Int("workers", 8, "parallel gRPC streams")
		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")
		metrics   = flag.String("metrics", ":8892", "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")
		dedupePct = flag.Int("dedupe-pct", 30, "percent of alerts sharing a dedupe_key")
		dedupeKey = flag.String("dedupe-key", "", "force a specific dedupe_key on every alert")
		payloadB  = flag.Int("payload-bytes", 256, "approximate payload size in bytes (Data map)")
	)

	flag.Parse()
	if *apiKey == "" {
		fmt.Fprintln(os.Stderr, "loadgen-grpc: --api-key is required (company:source:secret)")
		os.Exit(2)
	}
	parts := strings.SplitN(*apiKey, ":", 3)
	if len(parts) != 3 {
		fmt.Fprintln(os.Stderr, "loadgen-grpc: --api-key must be company:source:secret")
		os.Exit(2)
	}
	_, 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,
		"workers", *workers,
		"ramp_up", *rampUp,
		"duration", *duration,
		"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
		rttHist atomicHistogram // lives outside workers
	)


	results := make([]workerResult, *workers)
	var wg sync.WaitGroup

	// Rate per worker.
	ratePerWorker := *rate / *workers
	if ratePerWorker < 1 {
		ratePerWorker = 1
	}

	for wid := 0; wid < *workers; wid++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			res := &results[id]
			runWorker(ctx, logger, *target, *apiKey, source, secret,
				ratePerWorker, *duration, *rampUp, *dedupePct, *dedupeKey, *payloadB, res, &rttHist)
		}(wid)
	}

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

	// Status ticker.
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	// done is closed when every worker has returned (e.g. when
	// --duration expires in runWorker via prodCtx.Done()). Without
	// this, main only exits on SIGTERM and the loadgen hangs past
	// --duration even though the workers themselves shut down cleanly
	// (producer closes ackCh, consumer's `for range ackCh` exits,
	// runWorker returns). Symptom: smoke step 3 times out at 60s
	// with "backpressure loadgen did not exit within 60s".
	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	for {
		select {
		case <-ctx.Done():
			// Signal received: wait for in-flight workers to drain
			// via the done channel, then print and return.
			<-done
			var tsent, tfailed, tdupes, trlHits uint64
			for i := range results {
				tsent += results[i].sent
				tfailed += results[i].failed
				tdupes += results[i].dupes
				trlHits += results[i].rlHits
			}
			printFinal(logger, tsent, tfailed, tdupes, trlHits, &rttHist)
			return
		case <-done:
			// All workers finished naturally (e.g. --duration expired).
			var tsent, tfailed, tdupes, trlHits uint64
			for i := range results {
				tsent += results[i].sent
				tfailed += results[i].failed
				tdupes += results[i].dupes
				trlHits += results[i].rlHits
			}
			printFinal(logger, tsent, tfailed, tdupes, trlHits, &rttHist)
			return
		case <-ticker.C:
			// Aggregate per-worker counters.
			var tsent, tfailed, tdupes, trlHits uint64
			for i := range results {
				tsent += results[i].sent
				tfailed += results[i].failed
				tdupes += results[i].dupes
				trlHits += results[i].rlHits
			}
			logger.Info("progress",
				"sent", tsent,
				"failed", tfailed,
				"dupes", tdupes,
				"rate_limited", trlHits,
			)
		}
	}
}

// runWorker opens one gRPC stream and drives it at the target rate.
func runWorker(
	ctx context.Context,
	logger *slog.Logger,
	target, apiKey, sourceID, secret string,
	rate int,
	duration, rampUp time.Duration,
	dedupePct int,
	dedupeKey string,
	payloadB int,
	res *workerResult,
	rttHist *atomicHistogram,
) {
	// Connect.
	c, err := grpcclient.New(target,
		grpcclient.WithAPIKey(apiKey),
		grpcclient.WithInsecure(), // loadgen — use mTLS in production
	)
	if err != nil {
		logger.Error("grpcclient.New", "worker", 0, "err", err)
		return
	}
	defer c.Close()

	// Open stream.
	stream, err := c.Stream(ctx)
	if err != nil {
		logger.Error("Client.Stream", "err", err)
		return
	}
	defer stream.CloseSend()

	// Producer goroutine.
	prodCtx, cancelProd := context.WithTimeout(ctx, duration)
	defer cancelProd()

	type pendingAck struct {
		sentAt time.Time
		dk     string
	}
	ackCh := make(chan pendingAck, rate*2)

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

	go func() {
		i := 0
		for {
			select {
			case <-prodCtx.Done():
				stream.CloseSend()
				close(ackCh)
				return
			case <-tickCh:
				a := mkAlert("normal", companyFromKey(apiKey), sourceID, dedupePct, dedupeKey, payloadB)
				if err := stream.Send(prodCtx, alertToProto(companyFromKey(apiKey), sourceID, a)); err != nil {
					// stream.Send (loadgen wrapper) already retries 3x on
					// retriable codes. Treat any error here as terminal
					// for this tick. Close the stream and ackCh so the
					// consumer's 'for p := range ackCh' loop exits; wg.Wait
					// in main() then unblocks and the worker can shut down.
					// Without this close, the consumer spins forever on a
					// dead stream and the whole loadgen hangs past --duration.
					logger.Error("send failed", "err", err)
					res.failed++
					stream.CloseSend()
					close(ackCh)
					return
				}
				// Push to ackCh inside the select so prodCtx.Done() can
				// unblock the producer when --duration expires. Without
				// this, a full ackCh blocks the producer indefinitely
				// (the consumer drains at the server's ack rate, which
				// may be much slower than the pacer under backpressure).
				select {
				case <-prodCtx.Done():
					stream.CloseSend()
					close(ackCh)
					return
				case ackCh <- pendingAck{sentAt: time.Now(), dk: a.DedupeKey}:
				}
				i++
			}
		}
	}()

	// Consumer goroutine — reads acks and computes RTT.
	go func() {
		for p := range ackCh {
			ack, err := stream.Recv(prodCtx)
			if err != nil {
				// Stream is dead or ctx cancelled. Exit the consumer so
				// wg.Wait in main() can unblock — otherwise the loadgen
				// hangs past --duration even though the producer exited
				// cleanly via the Send-error or prodCtx.Done() path.
				return
			}
			rttMs := time.Since(p.sentAt).Milliseconds()
			rttHist.Record(rttMs)

			if ack.GetError() != nil {
				if ack.GetError().Code == pbv1.Error_RATE_LIMITED {
					res.rlHits++
				} else {
					res.failed++
				}
			} else {
				res.sent++
				if ack.DedupeCount > 1 {
					res.dupes++
				}
			}
		}
	}()

	<-prodCtx.Done()
}

// companyFromKey extracts company_id from the api-key "company:source:secret".
func companyFromKey(apiKey string) string {
	parts := strings.SplitN(apiKey, ":", 3)
	if len(parts) >= 1 {
		return parts[0]
	}
	return "acme-001"
}

func mkAlert(mode, companyID, sourceID string, dedupePct int, dedupeKey string, payloadB int) alert.Alert {
	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:  sourceID,
		Severity:  severity,
		Category:  category,
		Title:     fmt.Sprintf("%s on %s", category, data["host"]),
		Body:      "synthetic loadgen alert",
		Data:      data,
		DedupeKey: dk,
	}
}

func alertToProto(companyID, sourceID string, a alert.Alert) *pbv1.Alert {
	data := make(map[string]string)
	for k, v := range a.Data {
		data[k] = v
	}
	return &pbv1.Alert{
		CompanyId:  companyID,
		SourceId:   sourceID,
		Severity:   string(a.Severity),
		Category:   a.Category,
		Title:      a.Title,
		Body:       a.Body,
		DedupeKey:  a.DedupeKey,
		ClientTsMs: time.Now().UnixMilli(),
		Data:       data,
	}
}

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))]
}

// atomicHistogram is a lock-free fixed-size histogram for RTT values.
// Buckets: 0-1ms, 1-5ms, 5-10ms, 10-25ms, 25-50ms, 50-100ms, 100-250ms, 250ms+.
type atomicHistogram struct {
	buckets [8]atomic.Uint64
}

func (h *atomicHistogram) Record(ms int64) {
	var bucket int
	switch {
	case ms < 1:
		bucket = 0
	case ms < 5:
		bucket = 1
	case ms < 10:
		bucket = 2
	case ms < 25:
		bucket = 3
	case ms < 50:
		bucket = 4
	case ms < 100:
		bucket = 5
	case ms < 250:
		bucket = 6
	default:
		bucket = 7
	}
	h.buckets[bucket].Add(1)
}

func (h *atomicHistogram) String() string {
	var total uint64
	for i := 0; i < len(h.buckets); i++ {
		total += h.buckets[i].Load()
	}
	if total == 0 {
		return "no data"
	}
	return fmt.Sprintf("total=%d p50=%s p99=%s",
		total, h.percentile(50), h.percentile(99))
}

func (h *atomicHistogram) percentile(p int) string {
	var total uint64
	for i := 0; i < len(h.buckets); i++ {
		total += h.buckets[i].Load()
	}
	threshold := uint64(float64(total) * float64(p) / 100.0)
	var cumulative uint64
	bounds := []string{"<1ms", "1-5ms", "5-10ms", "10-25ms", "25-50ms", "50-100ms", "100-250ms", ">250ms"}
	for i := 0; i < len(h.buckets); i++ {
		cumulative += h.buckets[i].Load()
		if cumulative >= threshold {
			return bounds[i]
		}
	}
	return ">250ms"
}

func runMetrics(addr, clusterID, instance 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{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())
		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 printFinal(logger *slog.Logger, sent, failed, dupes, rlHits uint64, rttHist *atomicHistogram) {
	logger.Info("done",
		"sent", sent,
		"failed", failed,
		"dupes", dupes,
		"rate_limited", rlHits,
		"rtt", rttHist.String(),
	)
}

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

var _ = strconv.FormatInt // for future use