fix: harden websocket heartbeat under load

loadtest/im-pressure/wsclient.go

@@ -40,6 +40,8 @@ type ctrlError struct {
 	Text string
 }
 
+const dataCallbackQueueSize = 512
+
 func (e *ctrlError) Error() string {
 	return fmt.Sprintf("ctrl code=%d text=%s", e.Code, e.Text)
 }
@@ -56,6 +58,7 @@ type WSClient struct {
 	seq    uint64
 	waits  map[string]chan ctrlPacket
 	done   chan struct{}
+	data   chan dataPacket
 	closed atomic.Bool
 }
 
@@ -84,6 +87,7 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap
 		account: account, conn: conn, commandTimeout: timeout, onData: onData, onDisconnect: onDisconnect,
 		waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
 	}
+	client.startDataProcessor()
 	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",
@@ -346,6 +350,9 @@ func (c *WSClient) removeWait(id string) {
 func (c *WSClient) readLoop() {
 	defer func() {
 		_ = c.conn.Close()
+		if c.data != nil {
+			close(c.data)
+		}
 		close(c.done)
 	}()
 	for {
@@ -372,12 +379,28 @@ func (c *WSClient) readLoop() {
 		if raw := envelope["data"]; len(raw) > 0 && c.onData != nil {
 			var data dataPacket
 			if json.Unmarshal(raw, &data) == nil {
-				c.onData(c.account, data)
+				if c.data == nil {
+					c.onData(c.account, data)
+				} else {
+					c.data <- data
+				}
 			}
 		}
 	}
 }
 
+func (c *WSClient) startDataProcessor() {
+	if c.onData == nil || c.data != nil {
+		return
+	}
+	c.data = make(chan dataPacket, dataCallbackQueueSize)
+	go func() {
+		for data := range c.data {
+			c.onData(c.account, data)
+		}
+	}()
+}
+
 func (c *WSClient) startKeepalive(interval time.Duration) {
 	go func() {
 		ticker := time.NewTicker(interval)

loadtest/im-pressure/wsclient_test.go

@@ -84,3 +84,58 @@ func TestWSClientReadLoopClosesConnectionAfterDecodeError(t *testing.T) {
 		t.Fatal("read loop exited without closing the underlying connection")
 	}
 }
+
+func TestWSClientDataCallbackDoesNotBlockPong(t *testing.T) {
+	pong := 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.SetPongHandler(func(string) error {
+			select {
+			case pong <- struct{}{}:
+			default:
+			}
+			return nil
+		})
+		if err = conn.WriteJSON(map[string]any{"data": map[string]any{
+			"topic": "grp-test", "seq": 1, "head": map[string]any{}, "content": map[string]any{},
+		}}); err != nil {
+			return
+		}
+		if err = conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(time.Second)); err != nil {
+			return
+		}
+		conn.SetReadDeadline(time.Now().Add(time.Second))
+		for {
+			if _, _, readErr := conn.ReadMessage(); readErr != nil {
+				return
+			}
+		}
+	}))
+	defer server.Close()
+
+	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	blocked := make(chan struct{})
+	client := &WSClient{
+		conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
+		onData: func(AccountRecord, dataPacket) { <-blocked },
+	}
+	client.startDataProcessor()
+	go client.readLoop()
+	defer func() {
+		close(blocked)
+		client.Close()
+	}()
+
+	select {
+	case <-pong:
+	case <-time.After(2 * time.Second):
+		t.Fatal("pong was blocked by the data callback")
+	}
+}