pro/im_server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: trace delayed websocket acknowledgements

Browse Source
This commit is contained in:
richjohnrambo
2026-10-01 15:19:57 +08:00
parent 19584d0939
commit 73a594146f
19 changed files with 788 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
loadtest/im-pressure/wsclient.go
+107 -19
View File
@@ -7,6 +7,7 @@ import (
"encoding/json"
"errors"
"fmt"
"net"
"net/http"
"net/url"
"strconv"
@@ -40,12 +41,36 @@ type queuedDataPacket struct {
queuedAt time.Time
}
type ctrlWait struct {
ch chan ctrlPacket
kind string
sentAt time.Time
}
type timedOutCtrl struct {
kind string
sentAt time.Time
}
type ctrlObservation struct {
ID string
Kind string
Code int
SentAt time.Time
ReceivedAt time.Time
Latency time.Duration
Late bool
}
type ctrlError struct {
Code int
Text string
}
const dataCallbackQueueSize = 512
const (
dataCallbackQueueSize = 512
lateCtrlRetention = 10 * time.Minute
)
func (e *ctrlError) Error() string {
return fmt.Sprintf("ctrl code=%d text=%s", e.Code, e.Text)
@@ -57,19 +82,22 @@ type WSClient struct {
commandTimeout time.Duration
onData func(AccountRecord, dataPacket)
onDisconnect func(AccountRecord, error)
onCtrlObserved func(AccountRecord, ctrlObservation)
callbackQueue *callbackQueueMetrics
sendMu sync.Mutex
mu sync.Mutex
seq uint64
waits map[string]chan ctrlPacket
done chan struct{}
data chan queuedDataPacket
closed atomic.Bool
sendMu sync.Mutex
mu sync.Mutex
seq uint64
waits map[string]ctrlWait
timedOutCtrls map[string]timedOutCtrl
done chan struct{}
data chan queuedDataPacket
closed atomic.Bool
}
func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, apiKey, password string,
onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error), callbackQueue *callbackQueueMetrics) (*WSClient, error) {
onData func(AccountRecord, dataPacket), onDisconnect func(AccountRecord, error),
onCtrlObserved func(AccountRecord, ctrlObservation), callbackQueue *callbackQueueMetrics) (*WSClient, error) {
timeout, err := configDuration(cfg.CommandTimeout)
if err != nil {
return nil, err
@@ -79,6 +107,17 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap
return nil, err
}
dialer := *websocket.DefaultDialer
if account.DialIP != "" {
targetHost := endpointHostname(endpoint)
var netDialer net.Dialer
dialer.NetDialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
host, port, splitErr := net.SplitHostPort(address)
if splitErr == nil && strings.EqualFold(host, targetHost) {
address = net.JoinHostPort(account.DialIP, port)
}
return netDialer.DialContext(ctx, network, address)
}
}
if cfg.InsecureTLS {
dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} // #nosec G402: explicit load-test option.
}
@@ -91,8 +130,9 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap
}
client := &WSClient{
account: account, conn: conn, commandTimeout: timeout, onData: onData, onDisconnect: onDisconnect,
callbackQueue: callbackQueue,
waits: make(map[string]chan ctrlPacket), done: make(chan struct{}),
onCtrlObserved: onCtrlObserved,
callbackQueue: callbackQueue,
waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), done: make(chan struct{}),
}
client.startDataProcessor()
go client.readLoop()
@@ -136,6 +176,14 @@ func connectWSClient(ctx context.Context, cfg *Config, account AccountRecord, ap
return client, nil
}
func endpointHostname(raw string) string {
u, err := url.Parse(raw)
if err != nil {
return ""
}
return u.Hostname()
}
func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord, api *apiClient, password, verificationCode string) (*WSClient, bool, error) {
timeout, err := configDuration(cfg.CommandTimeout)
if err != nil {
@@ -153,7 +201,8 @@ func connectForPrepare(ctx context.Context, cfg *Config, account AccountRecord,
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{})}
client := &WSClient{account: account, conn: conn, commandTimeout: timeout,
waits: make(map[string]ctrlWait), timedOutCtrls: make(map[string]timedOutCtrl), 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,
@@ -296,14 +345,15 @@ func ctrlInt64(raw json.RawMessage) (int64, error) {
return strconv.ParseInt(text, 10, 64)
}
func (c *WSClient) command(ctx context.Context, _ string, message map[string]any) (ctrlPacket, error) {
func (c *WSClient) command(ctx context.Context, kind 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)
sentAt := time.Now()
c.mu.Lock()
c.waits[id] = wait
c.waits[id] = ctrlWait{ch: wait, kind: kind, sentAt: sentAt}
c.mu.Unlock()
c.sendMu.Lock()
err := c.conn.WriteJSON(message)
@@ -321,7 +371,33 @@ func (c *WSClient) command(ctx context.Context, _ string, message map[string]any
}
return ctrl, nil
case <-timer.C:
c.removeWait(id)
c.mu.Lock()
entry, pending := c.waits[id]
if pending {
delete(c.waits, id)
if entry.kind == "pub" {
c.timedOutCtrls[id] = timedOutCtrl{kind: entry.kind, sentAt: entry.sentAt}
}
}
c.mu.Unlock()
if !pending {
// The read loop won the timer race and placed the response in the
// buffered channel before removing the waiter.
ctrl := <-wait
if ctrl.Code < 200 || ctrl.Code >= 300 {
return ctrl, &ctrlError{Code: ctrl.Code, Text: ctrl.Text}
}
return ctrl, nil
}
if entry.kind == "pub" {
time.AfterFunc(lateCtrlRetention, func() {
c.mu.Lock()
if current, ok := c.timedOutCtrls[id]; ok && current.sentAt.Equal(sentAt) {
delete(c.timedOutCtrls, id)
}
c.mu.Unlock()
})
}
return ctrlPacket{}, fmt.Errorf("ctrl timeout id=%s", id)
case <-ctx.Done():
c.removeWait(id)
@@ -374,11 +450,23 @@ func (c *WSClient) readLoop() {
var ctrl ctrlPacket
if json.Unmarshal(raw, &ctrl) == nil && ctrl.ID != "" {
c.mu.Lock()
wait := c.waits[ctrl.ID]
delete(c.waits, ctrl.ID)
wait, pending := c.waits[ctrl.ID]
late, timedOut := c.timedOutCtrls[ctrl.ID]
receivedAt := time.Now()
if pending {
wait.ch <- ctrl
delete(c.waits, ctrl.ID)
} else if timedOut {
delete(c.timedOutCtrls, ctrl.ID)
}
c.mu.Unlock()
if wait != nil {
wait <- ctrl
if c.onCtrlObserved != nil && (pending && wait.kind == "pub" || timedOut && late.kind == "pub") {
kind, sentAt, isLate := wait.kind, wait.sentAt, false
if timedOut {
kind, sentAt, isLate = late.kind, late.sentAt, true
}
c.onCtrlObserved(c.account, ctrlObservation{ID: ctrl.ID, Kind: kind, Code: ctrl.Code,
SentAt: sentAt, ReceivedAt: receivedAt, Latency: receivedAt.Sub(sentAt), Late: isLate})
}
}
continue