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,365 @@
+package main
+
+import (
+	"context"
+	"crypto/tls"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+type ctrlPacket struct {
+	ID     string                     `json:"id"`
+	Topic  string                     `json:"topic"`
+	Code   int                        `json:"code"`
+	Text   string                     `json:"text"`
+	Params map[string]json.RawMessage `json:"params"`
+}
+
+type dataPacket struct {
+	Topic     string          `json:"topic"`
+	From      string          `json:"from"`
+	Timestamp time.Time       `json:"ts"`
+	Seq       int64           `json:"seq"`
+	Head      map[string]any  `json:"head"`
+	Content   json.RawMessage `json:"content"`
+}
+
+type ctrlError struct {
+	Code int
+	Text string
+}
+
+func (e *ctrlError) Error() string {
+	return fmt.Sprintf("ctrl code=%d text=%s", e.Code, e.Text)
+}
+
+type WSClient struct {
+	account        AccountRecord
+	conn           *websocket.Conn
+	commandTimeout time.Duration
+	onData         func(AccountRecord, dataPacket)
+	onDisconnect   func(AccountRecord, error)
+
+	sendMu sync.Mutex
+	mu     sync.Mutex
+	seq    uint64
+	waits  map[string]chan ctrlPacket
+	done   chan struct{}
+	closed atomic.Bool
+}
+
+func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string,
+	onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error)) (*WSClient, error) {
+	timeout, err := configDuration(cfg.CommandTimeout)
+	if err != nil {
+		return nil, err
+	}
+	endpoint, err := websocketEndpoint(account.Endpoint, apiKey)
+	if err != nil {
+		return nil, err
+	}
+	dialer := *websocket.DefaultDialer
+	if cfg.InsecureTLS {
+		dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
+	}
+	conn, response, err := dialer.DialContext(ctx, endpoint, http.Header{})
+	if err != nil {
+		if response != nil {
+			return nil, fmt.Errorf("websocket dial status=%d: %w", response.StatusCode, err)
+		}
+		return nil, err
+	}
+	client := &WSClient{
+		account: account, conn: conn, commandTimeout: timeout, onData: onData, onDisconnect: onDisconnect,
+		waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
+	}
+	go client.readLoop()
+	if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
+		"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure/1.0",
+		"platf": "loadtest", "tenant": cfg.Tenant,
+	}}); err != nil {
+		client.Close()
+		return nil, fmt.Errorf("hi %s: %w", account.Username, err)
+	}
+	login := map[string]any{"id": client.nextID("login")}
+	if account.Token != "" {
+		login["scheme"] = "token"
+		login["secret"] = account.Token
+	} else {
+		login["scheme"] = "basic"
+		login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
+		login["login_type"] = 0
+		login["tenant_code"] = cfg.Tenant
+	}
+	ctrl, err := client.command(ctx, "login", map[string]any{"login": login})
+	if err != nil && account.Token != "" && password != "" {
+		login["id"] = client.nextID("login-basic")
+		login["scheme"] = "basic"
+		login["secret"] = base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
+		login["login_type"] = 0
+		login["tenant_code"] = cfg.Tenant
+		ctrl, err = client.command(ctx, "login", map[string]any{"login": login})
+	}
+	if err != nil {
+		client.Close()
+		return nil, fmt.Errorf("login %s: %w", account.Username, err)
+	}
+	client.applyLogin(ctrl)
+	return client, nil
+}
+
+func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string) (*WSClient, bool, error) {
+	timeout, err := configDuration(cfg.CommandTimeout)
+	if err != nil {
+		return nil, false, err
+	}
+	endpoint, err := websocketEndpoint(account.Endpoint, apiKey)
+	if err != nil {
+		return nil, false, err
+	}
+	dialer := *websocket.DefaultDialer
+	if cfg.InsecureTLS {
+		dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
+	}
+	conn, _, err := dialer.DialContext(ctx, endpoint, http.Header{})
+	if err != nil {
+		return nil, false, err
+	}
+	client := &WSClient{account: account, conn: conn, commandTimeout: timeout, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
+	go client.readLoop()
+	if _, err = client.command(ctx, "hi", map[string]any{"hi": map[string]any{
+		"id": client.nextID("hi"), "ver": "0.25", "ua": "im-pressure-prepare/1.0", "tenant": cfg.Tenant,
+	}}); err != nil {
+		client.Close()
+		return nil, false, err
+	}
+	secret := base64.StdEncoding.EncodeToString([]byte(account.Username + ":" + password))
+	ctrl, loginErr := client.command(ctx, "login", map[string]any{"login": map[string]any{
+		"id": client.nextID("login"), "scheme": "basic", "secret": secret,
+		"login_type": 0, "tenant_code": cfg.Tenant,
+	}})
+	created := false
+	if loginErr != nil {
+		var ce *ctrlError
+		if !errors.As(loginErr, &ce) || ce.Code != http.StatusUnauthorized {
+			client.Close()
+			return nil, false, loginErr
+		}
+		ctrl, err = client.command(ctx, "acc", map[string]any{"acc": map[string]any{
+			"id": client.nextID("acc"), "user": "new", "scheme": "basic", "username": account.Username,
+			"registration_type": 0, "tenant_code": cfg.Tenant, "secret": secret, "login": true,
+			"desc": map[string]any{"public": map[string]any{"fn": strings.ToUpper(account.Username)}},
+		}})
+		if err != nil {
+			client.Close()
+			return nil, false, err
+		}
+		created = true
+	}
+	client.applyLogin(ctrl)
+	return client, created, nil
+}
+
+func websocketEndpoint(raw, apiKey string) (string, error) {
+	u, err := url.Parse(strings.TrimSpace(raw))
+	if err != nil {
+		return "", err
+	}
+	if u.Scheme != "ws" && u.Scheme != "wss" {
+		return "", fmt.Errorf("unsupported websocket scheme %q", u.Scheme)
+	}
+	q := u.Query()
+	if q.Get("apikey") == "" {
+		q.Set("apikey", apiKey)
+	}
+	u.RawQuery = q.Encode()
+	return u.String(), nil
+}
+
+func (c *WSClient) applyLogin(ctrl ctrlPacket) {
+	if raw := ctrl.Params["user"]; len(raw) > 0 {
+		_ = json.Unmarshal(raw, &c.account.UserID)
+	}
+	if raw := ctrl.Params["token"]; len(raw) > 0 {
+		_ = json.Unmarshal(raw, &c.account.Token)
+	}
+	if raw := ctrl.Params["expires"]; len(raw) > 0 {
+		var value string
+		if json.Unmarshal(raw, &value) == nil {
+			c.account.TokenExpires, _ = time.Parse(time.RFC3339Nano, value)
+		}
+	}
+}
+
+func (c *WSClient) Account() AccountRecord { return c.account }
+
+func (c *WSClient) Subscribe(ctx context.Context, topic string) error {
+	if topic == "" {
+		return errors.New("empty topic")
+	}
+	id := c.nextID("sub")
+	_, err := c.command(ctx, "sub", map[string]any{"sub": map[string]any{
+		"id": id, "topic": topic, "get": map[string]any{"what": "desc"},
+	}})
+	return err
+}
+
+func (c *WSClient) GetData(ctx context.Context, topic string, since, before int64, limit int) error {
+	if topic == "" || since <= 0 || before <= since || limit <= 0 {
+		return errors.New("invalid get data range")
+	}
+	id := c.nextID("get-data")
+	_, err := c.command(ctx, "get", map[string]any{"get": map[string]any{
+		"id": id, "topic": topic, "what": "data",
+		"data": map[string]any{"since": since, "before": before, "limit": limit},
+	}})
+	return err
+}
+
+func (c *WSClient) Publish(ctx context.Context, topic, clientMid, runID, scenario string, messageIndex int) (AcceptedMessage, error) {
+	id := c.nextID("pub")
+	sentAt := time.Now().UTC()
+	ctrl, err := c.command(ctx, "pub", map[string]any{"pub": map[string]any{
+		"id": id, "topic": topic, "noecho": false,
+		"head": map[string]any{"mime": "application/json", "x-im-schema": 1, "x-im-type": "text", "x-im-client-mid": clientMid},
+		"content": map[string]any{"text": "loadtest", "run_id": runID, "scenario": scenario,
+			"sender_index": c.account.Index, "message_index": messageIndex, "sent_at_unix_nano": sentAt.UnixNano()},
+	}})
+	if err != nil {
+		return AcceptedMessage{}, err
+	}
+	seq, err := ctrlInt64(ctrl.Params["seq"])
+	if err != nil || seq <= 0 {
+		return AcceptedMessage{}, fmt.Errorf("publish ACK has invalid seq: %w", err)
+	}
+	kind := "p2p"
+	if strings.HasPrefix(topic, "grp") {
+		kind = "group"
+	}
+	return AcceptedMessage{RunID: runID, Scenario: scenario, Kind: kind, ClientMid: clientMid,
+		Topic: topic, Seq: seq, SenderMachine: c.account.Machine, SenderIndex: c.account.Index,
+		SentAt: sentAt, AckAt: time.Now().UTC()}, nil
+}
+
+func ctrlInt64(raw json.RawMessage) (int64, error) {
+	if len(raw) == 0 {
+		return 0, errors.New("missing value")
+	}
+	var number json.Number
+	if err := json.Unmarshal(raw, &number); err == nil {
+		return strconv.ParseInt(number.String(), 10, 64)
+	}
+	var text string
+	if err := json.Unmarshal(raw, &text); err != nil {
+		return 0, err
+	}
+	return strconv.ParseInt(text, 10, 64)
+}
+
+func (c *WSClient) command(ctx context.Context, _ string, message map[string]any) (ctrlPacket, error) {
+	id := nestedMessageID(message)
+	if id == "" {
+		return ctrlPacket{}, errors.New("command id is required")
+	}
+	wait := make(chan ctrlPacket, 1)
+	c.mu.Lock()
+	c.waits[id] = wait
+	c.mu.Unlock()
+	c.sendMu.Lock()
+	err := c.conn.WriteJSON(message)
+	c.sendMu.Unlock()
+	if err != nil {
+		c.removeWait(id)
+		return ctrlPacket{}, err
+	}
+	timer := time.NewTimer(c.commandTimeout)
+	defer timer.Stop()
+	select {
+	case ctrl := <-wait:
+		if ctrl.Code < 200 || ctrl.Code >= 300 {
+			return ctrl, &ctrlError{Code: ctrl.Code, Text: ctrl.Text}
+		}
+		return ctrl, nil
+	case <-timer.C:
+		c.removeWait(id)
+		return ctrlPacket{}, fmt.Errorf("ctrl timeout id=%s", id)
+	case <-ctx.Done():
+		c.removeWait(id)
+		return ctrlPacket{}, ctx.Err()
+	case <-c.done:
+		c.removeWait(id)
+		return ctrlPacket{}, errors.New("websocket closed")
+	}
+}
+
+func nestedMessageID(message map[string]any) string {
+	for _, value := range message {
+		if body, ok := value.(map[string]any); ok {
+			if id, ok := body["id"].(string); ok {
+				return id
+			}
+		}
+	}
+	return ""
+}
+
+func (c *WSClient) nextID(prefix string) string {
+	n := atomic.AddUint64(&c.seq, 1)
+	return fmt.Sprintf("%s-%s-%04d-%d", prefix, c.account.Machine, c.account.Index, n)
+}
+
+func (c *WSClient) removeWait(id string) {
+	c.mu.Lock()
+	delete(c.waits, id)
+	c.mu.Unlock()
+}
+
+func (c *WSClient) readLoop() {
+	defer close(c.done)
+	for {
+		var envelope map[string]json.RawMessage
+		if err := c.conn.ReadJSON(&envelope); err != nil {
+			if !c.closed.Load() && c.onDisconnect != nil {
+				c.onDisconnect(c.account, err)
+			}
+			return
+		}
+		if raw := envelope["ctrl"]; len(raw) > 0 {
+			var ctrl ctrlPacket
+			if json.Unmarshal(raw, &ctrl) == nil && ctrl.ID != "" {
+				c.mu.Lock()
+				wait := c.waits[ctrl.ID]
+				delete(c.waits, ctrl.ID)
+				c.mu.Unlock()
+				if wait != nil {
+					wait <- ctrl
+				}
+			}
+			continue
+		}
+		if raw := envelope["data"]; len(raw) > 0 && c.onData != nil {
+			var data dataPacket
+			if json.Unmarshal(raw, &data) == nil {
+				c.onData(c.account, data)
+			}
+		}
+	}
+}
+
+func (c *WSClient) Close() {
+	if c == nil || !c.closed.CompareAndSwap(false, true) {
+		return
+	}
+	_ = c.conn.Close()
+}