From ea8732e2b100e49b011e7980911a26fb9217b74a Mon Sep 17 00:00:00 2001 From: richjohnrambo Date: Sun, 27 Sep 2026 21:41:50 +0800 Subject: [PATCH] fix: keep load test websocket sessions alive --- loadtest/im-pressure/README.md | 1 + loadtest/im-pressure/agent.go | 13 ++++++ loadtest/im-pressure/config.example.json | 1 + loadtest/im-pressure/model.go | 7 ++++ loadtest/im-pressure/model_test.go | 3 +- loadtest/im-pressure/wsclient.go | 37 ++++++++++++++++++ loadtest/im-pressure/wsclient_test.go | 50 ++++++++++++++++++++++++ 7 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 loadtest/im-pressure/wsclient_test.go diff --git a/loadtest/im-pressure/README.md b/loadtest/im-pressure/README.md index 7dd32c1..13bef67 100644 --- a/loadtest/im-pressure/README.md +++ b/loadtest/im-pressure/README.md @@ -41,6 +41,7 @@ - 两台机器都必须能访问配置中的 HTTP API 和各自的 WS 地址。 - `8090` 是无鉴权的压测协调端口，只允许压测内网访问，不要暴露到公网。 - 两台机器必须使用 NTP/chrony 同步时间。报告会分别输出同机时钟延迟和跨机器单向延迟；跨机器出现负延迟样本时，后者仅用于诊断，不参与门禁。 +- agent 默认每 20 秒发送 WebSocket Ping，覆盖等待另一台机器、发压、drain 和历史核验阶段；`ws_ping_interval` 必须小于公网入口的空闲连接超时。 ## 统计口径 diff --git a/loadtest/im-pressure/agent.go b/loadtest/im-pressure/agent.go index c187c1d..a58006c 100644 --- a/loadtest/im-pressure/agent.go +++ b/loadtest/im-pressure/agent.go @@ -141,6 +141,9 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo if err != nil { return err } + if alive := run.aliveConnections(); alive != len(run.clients) { + return fmt.Errorf("connections alive after coordinator wait=%d expected=%d", alive, len(run.clients)) + } if !runIDPattern.MatchString(start.RunID) || !validScenario(start.Scenario) { return fmt.Errorf("invalid coordinator start: %+v", start) } @@ -586,6 +589,16 @@ func (r *agentRun) acceptedCopy() []AcceptedMessage { return out } +func (r *agentRun) aliveConnections() int { + alive := 0 + for _, client := range r.clients { + if client.alive() { + alive++ + } + } + return alive +} + func (r *agentRun) drainDuration() (time.Duration, error) { if r.scenario == "p2p" { return configDuration(r.cfg.DrainP2P) diff --git a/loadtest/im-pressure/config.example.json b/loadtest/im-pressure/config.example.json index da71974..0e53dd1 100644 --- a/loadtest/im-pressure/config.example.json +++ b/loadtest/im-pressure/config.example.json @@ -14,6 +14,7 @@ "prepare_workers": 20, "connect_ramp": "20s", "command_timeout": "20s", + "ws_ping_interval": "20s", "preflight_workers": 32, "skip_p2p_preflight": false, "history_workers": 12, diff --git a/loadtest/im-pressure/model.go b/loadtest/im-pressure/model.go index d9fa54a..9e4f689 100644 --- a/loadtest/im-pressure/model.go +++ b/loadtest/im-pressure/model.go @@ -25,6 +25,7 @@ type Config struct { PrepareWorkers int `json:"prepare_workers"` ConnectRamp string `json:"connect_ramp"` CommandTimeout string `json:"command_timeout"` + WSPingInterval string `json:"ws_ping_interval"` P2PMessages int `json:"p2p_messages_per_user"` P2PInterval string `json:"p2p_interval"` DrainP2P string `json:"drain_p2p"` @@ -62,6 +63,12 @@ func loadConfig(path string) (*Config, error) { if cfg.CommandTimeout == "" { cfg.CommandTimeout = "20s" } + if cfg.WSPingInterval == "" { + cfg.WSPingInterval = "20s" + } + if _, err = configDuration(cfg.WSPingInterval); err != nil { + return nil, fmt.Errorf("ws_ping_interval: %w", err) + } if cfg.P2PMessages <= 0 { cfg.P2PMessages = 10 } diff --git a/loadtest/im-pressure/model_test.go b/loadtest/im-pressure/model_test.go index e15fb30..db30a50 100644 --- a/loadtest/im-pressure/model_test.go +++ b/loadtest/im-pressure/model_test.go @@ -16,7 +16,8 @@ func TestLoadConfigAppliesBoundedHistoryDefaults(t *testing.T) { if err != nil { t.Fatal(err) } - if cfg.HistoryWorkers != 12 || cfg.HistorySettle != "5m" || cfg.HistoryTimeout != "3m" || cfg.PreflightWorkers != 32 { + if cfg.HistoryWorkers != 12 || cfg.HistorySettle != "5m" || cfg.HistoryTimeout != "3m" || + cfg.PreflightWorkers != 32 || cfg.WSPingInterval != "20s" { t.Fatalf("unexpected defaults: %+v", cfg) } } diff --git a/loadtest/im-pressure/wsclient.go b/loadtest/im-pressure/wsclient.go index 558b7d9..8508523 100644 --- a/loadtest/im-pressure/wsclient.go +++ b/loadtest/im-pressure/wsclient.go @@ -116,6 +116,12 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap return nil, fmt.Errorf("login %s: %w", account.Username, err) } client.applyLogin(ctrl) + pingInterval, err := configDuration(cfg.WSPingInterval) + if err != nil { + client.Close() + return nil, err + } + client.startKeepalive(pingInterval) return client, nil } @@ -369,6 +375,37 @@ func (c *WSClient) readLoop() { } } +func (c *WSClient) startKeepalive(interval time.Duration) { + go func() { + ticker := time.NewTicker(interval) + defer ticker.Stop() + for { + select { + case <-ticker.C: + deadline := time.Now().Add(interval / 2) + if err := c.conn.WriteControl(websocket.PingMessage, nil, deadline); err != nil { + _ = c.conn.Close() + return + } + case <-c.done: + return + } + } + }() +} + +func (c *WSClient) alive() bool { + if c == nil || c.closed.Load() { + return false + } + select { + case <-c.done: + return false + default: + return true + } +} + func (c *WSClient) Close() { if c == nil || !c.closed.CompareAndSwap(false, true) { return diff --git a/loadtest/im-pressure/wsclient_test.go b/loadtest/im-pressure/wsclient_test.go new file mode 100644 index 0000000..d86ee68 --- /dev/null +++ b/loadtest/im-pressure/wsclient_test.go @@ -0,0 +1,50 @@ +package main + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +func TestWSClientKeepaliveSendsPing(t *testing.T) { + ping := make(chan struct{}, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + conn.SetPingHandler(func(string) error { + select { + case ping <- struct{}{}: + default: + } + return conn.WriteControl(websocket.PongMessage, nil, time.Now().Add(time.Second)) + }) + for { + if _, _, err = conn.ReadMessage(); err != nil { + return + } + } + })) + defer server.Close() + + conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil) + if err != nil { + t.Fatal(err) + } + client := &WSClient{conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})} + go client.readLoop() + client.startKeepalive(10 * time.Millisecond) + defer client.Close() + + select { + case <-ping: + case <-time.After(time.Second): + t.Fatal("websocket keepalive ping was not received") + } +}