feat(go): add 100k ingest hardening baseline

go/vehicle-gateway/internal/loadsim/runner.go

@@ -0,0 +1,103 @@
+package loadsim
+
+import (
+	"context"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)
+
+type Runner struct {
+	Dial DialFunc
+}
+
+type Stats struct {
+	ConnectionsOpened int64
+	ConnectionsFailed int64
+	FramesWritten     int64
+	WriteErrors       int64
+}
+
+func (r Runner) Run(ctx context.Context, cfg Config) (Stats, error) {
+	if _, err := (&FlagConfig{
+		protocol:     string(cfg.Protocol),
+		addr:         cfg.Addr,
+		connections:  cfg.Connections,
+		connectRate:  cfg.ConnectRatePerSecond,
+		sendInterval: cfg.SendInterval,
+		duration:     cfg.Duration,
+		template:     cfg.Template,
+	}).Build(); err != nil {
+		return Stats{}, err
+	}
+	payload, err := FrameTemplate(cfg.Protocol, cfg.Template)
+	if err != nil {
+		return Stats{}, err
+	}
+	dial := r.Dial
+	if dial == nil {
+		dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
+		dial = dialer.DialContext
+	}
+
+	runCtx, cancel := context.WithTimeout(ctx, cfg.Duration)
+	defer cancel()
+
+	var stats Stats
+	var wg sync.WaitGroup
+	for i, batch := range ConnectionBatches(cfg.Connections, cfg.ConnectRatePerSecond) {
+		if i > 0 {
+			if !sleepOrDone(runCtx, time.Second) {
+				break
+			}
+		}
+		for range batch {
+			if runCtx.Err() != nil {
+				break
+			}
+			conn, err := dial(runCtx, "tcp", cfg.Addr)
+			if err != nil {
+				atomic.AddInt64(&stats.ConnectionsFailed, 1)
+				continue
+			}
+			atomic.AddInt64(&stats.ConnectionsOpened, 1)
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				defer conn.Close()
+				writeLoop(runCtx, conn, payload, cfg.SendInterval, &stats)
+			}()
+		}
+	}
+	<-runCtx.Done()
+	wg.Wait()
+	return stats, nil
+}
+
+func writeLoop(ctx context.Context, conn net.Conn, payload []byte, interval time.Duration, stats *Stats) {
+	for {
+		_ = conn.SetWriteDeadline(time.Now().Add(3 * time.Second))
+		if _, err := conn.Write(payload); err != nil {
+			atomic.AddInt64(&stats.WriteErrors, 1)
+			return
+		}
+		atomic.AddInt64(&stats.FramesWritten, 1)
+		if !sleepOrDone(ctx, interval) {
+			return
+		}
+	}
+}
+
+func sleepOrDone(ctx context.Context, duration time.Duration) bool {
+	timer := time.NewTimer(duration)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return false
+	case <-timer.C:
+		return true
+	}
+}