broad-announce/loadgen/cmd/m5drivers/tail/main.go

// m5-tail-test connects to the ingestd live-tail endpoint and
// prints every alert frame it receives until --duration
// elapses. This is the smoke driver used in scripts/m5_smoke.sh
// to assert that the tail hub actually fans out accepted alerts.
package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	target := flag.String("target", "ws://localhost:8800/v1/tail/ws", "tail ws URL")
	token := flag.String("token", "tail-dev-token-please-change-in-prod", "tail token")
	companyID := flag.String("company", "", "company filter (empty=all)")
	duration := flag.Duration("duration", 8*time.Second, "max listen time")
	flag.Parse()

	url := *target
	if *companyID != "" {
		url += "?company_id=" + *companyID + "&token=" + *token
	} else {
		url += "?token=" + *token
	}

	dialer := *websocket.DefaultDialer
	dialer.HandshakeTimeout = 5 * time.Second
	hdr := http.Header{}
	hdr.Set("X-BA-Tail-Token", *token)
	conn, resp, err := dialer.Dial(url, hdr)
	if err != nil {
		fmt.Println("dial err:", err, "resp:", resp)
		os.Exit(1)
	}
	defer conn.Close()
	fmt.Println("TAIL CONNECTED", url)

	var frames atomic.Int64
	deadline := time.Now().Add(*duration)
	// Set a single, long read deadline. We don't poll. We block
	// on the first frame; if it never arrives we let the
	// overall duration kill us.
	_ = conn.SetReadDeadline(deadline)
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			fmt.Println("read err:", err)
			break
		}
		frames.Add(1)
		fmt.Println("FRAME:", string(msg))
		// Move the deadline forward after each successful
		// read so a busy tail keeps the test alive.
		_ = conn.SetReadDeadline(time.Now().Add(*duration))
	}
	fmt.Printf("done: %d frames\n", frames.Load())
}