broad-announce/cmd/example-grpc-producer/main.go

// example-grpc-producer demonstrates how a peer Go service publishes
// alerts to ingestd via the gRPC StreamAlerts bidi-stream.
//
// Usage:
//
//	BA_API_KEY="company:source:secret" \
//	  go run ./cmd/example-grpc-producer \
//	    --addr localhost:9090 \
//	    --rate 100 \
//	    --duration 30s
//
// This is the reference implementation for the "How to publish to broad-announce"
// story documented in SPEC §19.
package main

import (
	"context"
	"flag"
	"fmt"
	"log/slog"
	"math/rand"
	"os"
	"sync/atomic"
	"time"

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

func main() {
	addr := flag.String("addr", "localhost:9090", "ingestd gRPC address")
	rate := flag.Int("rate", 100, "alerts per second")
	duration := flag.Duration("duration", 30*time.Second, "how long to run")
	flag.Parse()

	apiKey := os.Getenv("BA_API_KEY")
	if apiKey == "" {
		fmt.Fprintln(os.Stderr, "BA_API_KEY env var is required")
		os.Exit(1)
	}

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

	ctx, cancel := context.WithTimeout(context.Background(), *duration)
	defer cancel()

	// Connect to ingestd.
	c, err := grpcclient.New(*addr,
		grpcclient.WithAPIKey(apiKey),
		grpcclient.WithInsecure(), // remove in production (use TLS/mTLS)
	)
	if err != nil {
		logger.Error("grpcclient.New", "addr", *addr, "err", err)
		os.Exit(1)
	}
	defer c.Close()

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

	var sent int64

	// Producer goroutine: generates alerts at --rate.
	alertCh := make(chan *pbv1.Alert, *rate*2)
	go func() {
		ticker := time.NewTicker(time.Second / time.Duration(*rate))
		defer ticker.Stop()
		defer close(alertCh)
		i := 0
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				alertCh <- randomAlert(i)
				i++
			}
		}
	}()

	// Send loop: pulls from alertCh and sends to ingestd.
	go func() {
		for alert := range alertCh {
			if err := stream.Send(ctx, alert); err != nil {
				logger.Warn("stream.Send", "err", err)
				continue
			}
			atomic.AddInt64(&sent, 1)
			if sent%100 == 0 {
				logger.Info("sent", "count", sent)
			}
		}
	}()

	// Recv loop: reads Acks and logs every 100.
	for {
		ack, err := stream.Recv(ctx)
		if err != nil {
			if ctx.Err() != nil {
				return // normal shutdown
			}
			logger.Error("stream.Recv", "err", err)
			return
		}
		if atomic.LoadInt64(&sent)%100 == 0 {
			logger.Info("ack", "dedupe_key", ack.DedupeKey)
		}
	}
}

func randomAlert(i int) *pbv1.Alert {
	severities := []string{"critical", "warning", "info"}
	return &pbv1.Alert{
		CompanyId:  "acme-001",
		SourceId:   "prom-prod",
		Severity:   severities[rand.Intn(len(severities))],
		Category:   "monitoring",
		Title:      fmt.Sprintf("example alert %d", i),
		Body:       "This is a demo alert from example-grpc-producer.",
		DedupeKey:  fmt.Sprintf("example-%d", i),
		ClientTsMs: time.Now().UnixMilli(),
		Data: map[string]string{
			"host":  fmt.Sprintf("host-%d", rand.Intn(10)),
			"value": fmt.Sprintf("%d", rand.Intn(100)),
		},
	}
}