fix: keep load test websocket sessions alive

This commit is contained in:
richjohnrambo
2026-09-27 21:41:50 +08:00
parent bde02e8eac
commit ea8732e2b1
7 changed files with 111 additions and 1 deletions
+1
View File
@@ -41,6 +41,7 @@
- 两台机器都必须能访问配置中的 HTTP API 和各自的 WS 地址。 - 两台机器都必须能访问配置中的 HTTP API 和各自的 WS 地址。
- `8090` 是无鉴权的压测协调端口,只允许压测内网访问,不要暴露到公网。 - `8090` 是无鉴权的压测协调端口,只允许压测内网访问,不要暴露到公网。
- 两台机器必须使用 NTP/chrony 同步时间。报告会分别输出同机时钟延迟和跨机器单向延迟;跨机器出现负延迟样本时,后者仅用于诊断,不参与门禁。 - 两台机器必须使用 NTP/chrony 同步时间。报告会分别输出同机时钟延迟和跨机器单向延迟;跨机器出现负延迟样本时,后者仅用于诊断,不参与门禁。
- agent 默认每 20 秒发送 WebSocket Ping,覆盖等待另一台机器、发压、drain 和历史核验阶段;`ws_ping_interval` 必须小于公网入口的空闲连接超时。
## 统计口径 ## 统计口径
+13
View File
@@ -141,6 +141,9 @@ func runAgent(ctx context.Context, cfg *Config, manifest *Manifest, machine, coo
if err != nil { if err != nil {
return err 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) { if !runIDPattern.MatchString(start.RunID) || !validScenario(start.Scenario) {
return fmt.Errorf("invalid coordinator start: %+v", start) return fmt.Errorf("invalid coordinator start: %+v", start)
} }
@@ -586,6 +589,16 @@ func (r *agentRun) acceptedCopy() []AcceptedMessage {
return out 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) { func (r *agentRun) drainDuration() (time.Duration, error) {
if r.scenario == "p2p" { if r.scenario == "p2p" {
return configDuration(r.cfg.DrainP2P) return configDuration(r.cfg.DrainP2P)
+1
View File
@@ -14,6 +14,7 @@
"prepare_workers": 20, "prepare_workers": 20,
"connect_ramp": "20s", "connect_ramp": "20s",
"command_timeout": "20s", "command_timeout": "20s",
"ws_ping_interval": "20s",
"preflight_workers": 32, "preflight_workers": 32,
"skip_p2p_preflight": false, "skip_p2p_preflight": false,
"history_workers": 12, "history_workers": 12,
+7
View File
@@ -25,6 +25,7 @@ type Config struct {
PrepareWorkers int `json:"prepare_workers"` PrepareWorkers int `json:"prepare_workers"`
ConnectRamp string `json:"connect_ramp"` ConnectRamp string `json:"connect_ramp"`
CommandTimeout string `json:"command_timeout"` CommandTimeout string `json:"command_timeout"`
WSPingInterval string `json:"ws_ping_interval"`
P2PMessages int `json:"p2p_messages_per_user"` P2PMessages int `json:"p2p_messages_per_user"`
P2PInterval string `json:"p2p_interval"` P2PInterval string `json:"p2p_interval"`
DrainP2P string `json:"drain_p2p"` DrainP2P string `json:"drain_p2p"`
@@ -62,6 +63,12 @@ func loadConfig(path string) (*Config, error) {
if cfg.CommandTimeout == "" { if cfg.CommandTimeout == "" {
cfg.CommandTimeout = "20s" 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 { if cfg.P2PMessages <= 0 {
cfg.P2PMessages = 10 cfg.P2PMessages = 10
} }
+2 -1
View File
@@ -16,7 +16,8 @@ func TestLoadConfigAppliesBoundedHistoryDefaults(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) 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) t.Fatalf("unexpected defaults: %+v", cfg)
} }
} }
+37
View File
@@ -116,6 +116,12 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap
return nil, fmt.Errorf("login %s: %w", account.Username, err) return nil, fmt.Errorf("login %s: %w", account.Username, err)
} }
client.applyLogin(ctrl) client.applyLogin(ctrl)
pingInterval, err := configDuration(cfg.WSPingInterval)
if err != nil {
client.Close()
return nil, err
}
client.startKeepalive(pingInterval)
return client, nil 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() { func (c *WSClient) Close() {
if c == nil || !c.closed.CompareAndSwap(false, true) { if c == nil || !c.closed.CompareAndSwap(false, true) {
return return
+50
View File
@@ -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")
}
}