feat: add dual-host IM pressure test tool

2026-09-26 21:50:43 +08:00
parent f898a7171c
commit 2b90a2e759
21 changed files with 3843 additions and 0 deletions
@@ -0,0 +1,68 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/http"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+func startAgentMetrics(address string, run *agentRun) (func(), error) {
+	if address == "" {
+		return func() {}, nil
+	}
+	listener, err := net.Listen("tcp", address)
+	if err != nil {
+		return nil, fmt.Errorf("listen for agent metrics on %s: %w", address, err)
+	}
+	labels := prometheus.Labels{"machine": run.machine}
+	registry := prometheus.NewRegistry()
+	registry.MustRegister(
+		collectors.NewGoCollector(),
+		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
+		prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: "im_load_agent_ready_connections", Help: "Current ready WebSocket connections.", ConstLabels: labels}, func() float64 {
+			return float64(run.ready.Load())
+		}),
+		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_attempted_total", Help: "Publish commands attempted.", ConstLabels: labels}, func() float64 {
+			return float64(run.attempted.Load())
+		}),
+		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_accepted_total", Help: "Publish commands acknowledged by the server.", ConstLabels: labels}, func() float64 {
+			run.mu.Lock()
+			defer run.mu.Unlock()
+			return float64(len(run.accepted))
+		}),
+		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_publish_rejected_total", Help: "Publish commands rejected or timed out.", ConstLabels: labels}, func() float64 {
+			return float64(run.rejected.Load())
+		}),
+		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_received_total", Help: "Unique matching data deliveries received.", ConstLabels: labels}, func() float64 {
+			received, _, _, _ := run.recorder.metricCounts()
+			return float64(received)
+		}),
+		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_duplicates_total", Help: "Duplicate matching deliveries received.", ConstLabels: labels}, func() float64 {
+			_, duplicates, _, _ := run.recorder.metricCounts()
+			return float64(duplicates)
+		}),
+		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_out_of_order_total", Help: "Per-connection topic sequence regressions.", ConstLabels: labels}, func() float64 {
+			_, _, outOfOrder, _ := run.recorder.metricCounts()
+			return float64(outOfOrder)
+		}),
+		prometheus.NewCounterFunc(prometheus.CounterOpts{Name: "im_load_agent_disconnects_total", Help: "Unexpected WebSocket disconnects.", ConstLabels: labels}, func() float64 {
+			_, _, _, disconnects := run.recorder.metricCounts()
+			return float64(disconnects)
+		}),
+	)
+	mux := http.NewServeMux()
+	mux.Handle("/metrics", promhttp.HandlerFor(registry, promhttp.HandlerOpts{}))
+	server := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second}
+	go func() { _ = server.Serve(listener) }()
+	return func() {
+		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+		_ = server.Shutdown(ctx)
+	}, nil
+}