fix: trace delayed websocket acknowledgements

loadtest/im-pressure/wsclient_test.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -10,6 +11,49 @@ import (
 	"github.com/gorilla/websocket"
 )
 
+func TestWSClientObservesCtrlAfterCommandTimeout(t *testing.T) {
+	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()
+		var command map[string]any
+		if err = conn.ReadJSON(&command); err != nil {
+			return
+		}
+		time.Sleep(40 * time.Millisecond)
+		_ = conn.WriteJSON(map[string]any{"ctrl": map[string]any{"id": "pub-a-0001", "code": 202, "text": "accepted"}})
+	}))
+	defer server.Close()
+
+	conn, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(server.URL, "http"), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	observed := make(chan ctrlObservation, 1)
+	client := &WSClient{
+		account: AccountRecord{Machine: "a", Index: 1}, conn: conn, commandTimeout: 10 * time.Millisecond,
+		waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{}),
+		onCtrlObserved: func(_ AccountRecord, observation ctrlObservation) { observed <- observation },
+	}
+	go client.readLoop()
+	defer client.Close()
+
+	_, err = client.command(context.Background(), "pub", map[string]any{"pub": map[string]any{"id": "pub-a-0001"}})
+	if err == nil || !strings.Contains(err.Error(), "ctrl timeout") {
+		t.Fatalf("command error=%v, want ctrl timeout", err)
+	}
+	select {
+	case observation := <-observed:
+		if !observation.Late || observation.Kind != "pub" || observation.Code != 202 || observation.Latency < 10*time.Millisecond {
+			t.Fatalf("unexpected observation: %+v", observation)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("late ctrl response was not observed")
+	}
+}
+
 func TestWSClientKeepaliveSendsPing(t *testing.T) {
 	ping := make(chan struct{}, 1)
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -37,7 +81,7 @@ func TestWSClientKeepaliveSendsPing(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	client := &WSClient{conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
+	client := &WSClient{conn: conn, waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{})}
 	go client.readLoop()
 	client.startKeepalive(10 * time.Millisecond)
 	defer client.Close()
@@ -70,7 +114,7 @@ func TestWSClientReadLoopClosesConnectionAfterDecodeError(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	client := &WSClient{conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{})}
+	client := &WSClient{conn: conn, waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{})}
 	go client.readLoop()
 
 	select {
@@ -123,7 +167,7 @@ func TestWSClientDataCallbackDoesNotBlockPong(t *testing.T) {
 	}
 	blocked := make(chan struct{})
 	client := &WSClient{
-		conn: conn, waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
+		conn: conn, waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{}),
 		onData: func(AccountRecord, dataPacket) { <-blocked },
 	}
 	client.startDataProcessor()