feat: 测试交付 Q1 混沌压测骨架与 Q4 镜像骨架

2026-09-30 06:59:44 +08:00
parent 0e068c6ade
commit 8dde564517
21 changed files with 1268 additions and 16 deletions
@@ -0,0 +1,40 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"git.asio.asia/nixevol/NixMsg/test/load"
+)
+
+// 压测客户端骨架：连上 N 个 MQTT 连接并打印连接数。完整 1000 连接 / 10 分钟压测留给 Q3。
+func main() {
+	addr := flag.String("addr", "", "MQTT broker host:port（必填，不要写死业务端口）")
+	n := flag.Int("n", 10, "连接数")
+	prefix := flag.String("prefix", "q-bench-", "client id 前缀")
+	hold := flag.Duration("hold", 3*time.Second, "保持连接时长")
+	flag.Parse()
+	if *addr == "" {
+		fmt.Fprintln(os.Stderr, "usage: mqttbench -addr host:port [-n 10]")
+		os.Exit(2)
+	}
+
+	b := &load.MQTTBench{Addr: *addr, ClientIDPrefix: *prefix, Timeout: 5 * time.Second}
+	defer b.Close()
+	if err := b.ConnectN(*n, os.Stdout); err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+	fmt.Printf("held %d connections for %s\n", b.Alive(), hold.String())
+
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
+	select {
+	case <-time.After(*hold):
+	case <-ch:
+	}
+}
@@ -0,0 +1,145 @@
+// Package load 提供 MQTT 压测客户端骨架：先连上并统计连接数。
+package load
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// MQTTBench 保持多个 MQTT 3.1.1 TCP 连接（仅 CONNECT/CONNACK）。
+type MQTTBench struct {
+	Addr           string
+	ClientIDPrefix string
+	Timeout        time.Duration
+
+	mu    sync.Mutex
+	conns []net.Conn
+	alive atomic.Int64
+}
+
+// ConnectN 建立 n 条连接；成功一条 alive+1，并打印当前连接数到 w（可为 nil）。
+func (b *MQTTBench) ConnectN(n int, w io.Writer) error {
+	if b.Addr == "" {
+		return fmt.Errorf("addr required")
+	}
+	if n <= 0 {
+		return fmt.Errorf("n must be > 0")
+	}
+	timeout := b.Timeout
+	if timeout <= 0 {
+		timeout = 5 * time.Second
+	}
+	prefix := b.ClientIDPrefix
+	if prefix == "" {
+		prefix = "q-bench-"
+	}
+
+	for i := 0; i < n; i++ {
+		conn, err := net.DialTimeout("tcp", b.Addr, timeout)
+		if err != nil {
+			return fmt.Errorf("dial %d: %w", i, err)
+		}
+		_ = conn.SetDeadline(time.Now().Add(timeout))
+		cid := fmt.Sprintf("%s%d", prefix, i)
+		if err := mqttConnect(conn, cid); err != nil {
+			_ = conn.Close()
+			return fmt.Errorf("connect %d: %w", i, err)
+		}
+		_ = conn.SetDeadline(time.Time{})
+		b.mu.Lock()
+		b.conns = append(b.conns, conn)
+		b.mu.Unlock()
+		cur := b.alive.Add(1)
+		if w != nil {
+			_, _ = fmt.Fprintf(w, "mqtt connections: %d\n", cur)
+		}
+	}
+	return nil
+}
+
+// Alive 当前仍打开的连接数。
+func (b *MQTTBench) Alive() int64 {
+	return b.alive.Load()
+}
+
+// Close 关闭全部连接。
+func (b *MQTTBench) Close() {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	for _, c := range b.conns {
+		_ = c.Close()
+	}
+	b.conns = nil
+	b.alive.Store(0)
+}
+
+func mqttConnect(conn net.Conn, clientID string) error {
+	pkt := buildConnect(clientID)
+	if _, err := conn.Write(pkt); err != nil {
+		return err
+	}
+	header := make([]byte, 4)
+	if _, err := io.ReadFull(conn, header[:2]); err != nil {
+		return err
+	}
+	if header[0] != 0x20 {
+		return fmt.Errorf("unexpected packet type 0x%02x", header[0])
+	}
+	// remaining length 对 CONNACK 固定为 2
+	if header[1] != 2 {
+		return fmt.Errorf("unexpected remaining length %d", header[1])
+	}
+	if _, err := io.ReadFull(conn, header[2:4]); err != nil {
+		return err
+	}
+	if header[3] != 0 {
+		return fmt.Errorf("CONNACK rc=%d", header[3])
+	}
+	return nil
+}
+
+func buildConnect(clientID string) []byte {
+	// Variable header: protocol name MQTT, level 4, flags 0, keepalive 60
+	vh := []byte{
+		0x00, 0x04, 'M', 'Q', 'T', 'T',
+		0x04,
+		0x00, // clean session=0 flags for skeleton; brokers may still accept
+		0x00, 0x3c,
+	}
+	// Actually clean session bit should be set for simple benches
+	vh[7] = 0x02 // Clean Session
+
+	id := []byte(clientID)
+	payload := make([]byte, 2+len(id))
+	binary.BigEndian.PutUint16(payload[0:2], uint16(len(id)))
+	copy(payload[2:], id)
+
+	remaining := len(vh) + len(payload)
+	pkt := make([]byte, 0, 2+remaining)
+	pkt = append(pkt, 0x10)
+	pkt = append(pkt, encodeRemainingLength(remaining)...)
+	pkt = append(pkt, vh...)
+	pkt = append(pkt, payload...)
+	return pkt
+}
+
+func encodeRemainingLength(n int) []byte {
+	var out []byte
+	for {
+		encoded := byte(n % 128)
+		n /= 128
+		if n > 0 {
+			encoded |= 0x80
+		}
+		out = append(out, encoded)
+		if n == 0 {
+			break
+		}
+	}
+	return out
+}
@@ -0,0 +1,75 @@
+package load
+
+import (
+	"bytes"
+	"io"
+	"net"
+	"testing"
+	"time"
+)
+
+// 极简 MQTT broker：读 CONNECT，回 CONNACK accepted，保持连接。
+func startFakeMQTTBroker(t *testing.T) (addr string, closeFn func()) {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	done := make(chan struct{})
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				select {
+				case <-done:
+					return
+				default:
+					return
+				}
+			}
+			go func(c net.Conn) {
+				defer func() { _ = c.Close() }()
+				_ = c.SetDeadline(time.Now().Add(5 * time.Second))
+				buf := make([]byte, 256)
+				n, err := c.Read(buf)
+				if err != nil || n < 2 || buf[0] != 0x10 {
+					return
+				}
+				// CONNACK: type 0x20, remaining 2, flags 0, rc 0
+				_, _ = c.Write([]byte{0x20, 0x02, 0x00, 0x00})
+				_ = c.SetDeadline(time.Time{})
+				_, _ = io.Copy(io.Discard, c)
+			}(conn)
+		}
+	}()
+	return ln.Addr().String(), func() {
+		close(done)
+		_ = ln.Close()
+	}
+}
+
+func TestMQTTBenchConnectN(t *testing.T) {
+	addr, stop := startFakeMQTTBroker(t)
+	defer stop()
+
+	var out bytes.Buffer
+	b := &MQTTBench{Addr: addr, ClientIDPrefix: "q-t-", Timeout: 2 * time.Second}
+	defer b.Close()
+	if err := b.ConnectN(3, &out); err != nil {
+		t.Fatalf("ConnectN: %v", err)
+	}
+	if b.Alive() != 3 {
+		t.Fatalf("alive=%d", b.Alive())
+	}
+	got := out.String()
+	if !bytes.Contains(out.Bytes(), []byte("mqtt connections: 3")) {
+		t.Fatalf("output missing count: %q", got)
+	}
+}
+
+func TestMQTTBenchRequiresAddr(t *testing.T) {
+	b := &MQTTBench{}
+	if err := b.ConnectN(1, nil); err == nil {
+		t.Fatal("expected error")
+	}
+}