From 98e216841cf7eb5952bb9eb2c96579febc727ca4 Mon Sep 17 00:00:00 2001 From: jacentsao Date: Mon, 14 Sep 2026 15:58:06 +0800 Subject: [PATCH] fix(im): recover stale Yunzhijia WebSocket connections --- internal/im/yunzhijia/factory.go | 2 +- internal/im/yunzhijia/websocket.go | 163 ++++++++++++++++++------ internal/im/yunzhijia/websocket_test.go | 119 +++++++++++++++-- 3 files changed, 236 insertions(+), 48 deletions(-) diff --git a/internal/im/yunzhijia/factory.go b/internal/im/yunzhijia/factory.go index 40252c363..e9e4a341e 100644 --- a/internal/im/yunzhijia/factory.go +++ b/internal/im/yunzhijia/factory.go @@ -44,7 +44,7 @@ func NewFactory() im.AdapterFactory { if err != nil { return nil, nil, err } - client := NewLongConnClient(webSocketURL, msgHandler) + client := NewLongConnClient(channel.ID, webSocketURL, msgHandler) wsCtx, wsCancel := context.WithCancel(context.Background()) go func() { if err := client.Start(wsCtx); err != nil && wsCtx.Err() == nil { diff --git a/internal/im/yunzhijia/websocket.go b/internal/im/yunzhijia/websocket.go index f43213d05..234ded326 100644 --- a/internal/im/yunzhijia/websocket.go +++ b/internal/im/yunzhijia/websocket.go @@ -22,6 +22,7 @@ const ( webSocketMaxMessageSize = 1 << 20 webSocketMaxInvalidFrames = 3 webSocketMessageQueueSize = 64 + webSocketMaxConnectionAge = 6 * time.Hour ) var webSocketReconnectDelays = [...]time.Duration{ @@ -34,37 +35,55 @@ var webSocketReconnectDelays = [...]time.Duration{ } type webSocketFrame struct { - message *callbackMessage - ack []byte - control string + message *callbackMessage + ack []byte + cmd string + typeName string + event string + seq int64 + hasSeq bool + needAck bool + control string } type LongConnClient struct { - url string - handler func(context.Context, *im.IncomingMessage) error - dialer *ws.Dialer - messages chan *im.IncomingMessage + channelID string + url string + handler func(context.Context, *im.IncomingMessage) error + dialer *ws.Dialer + messages chan *im.IncomingMessage + maxConnectionAge time.Duration - mu sync.Mutex - conn *ws.Conn - closed atomic.Bool + mu sync.Mutex + conn *ws.Conn + closed atomic.Bool + connectedAt atomic.Int64 + lastFrameAt atomic.Int64 + lastPongAt atomic.Int64 + lastBusinessMessageAt atomic.Int64 + reconnectCount atomic.Int64 } -func NewLongConnClient(webSocketURL string, handler func(context.Context, *im.IncomingMessage) error) *LongConnClient { +func NewLongConnClient( + channelID, webSocketURL string, + handler func(context.Context, *im.IncomingMessage) error, +) *LongConnClient { dialer := *ws.DefaultDialer dialer.HandshakeTimeout = webSocketHandshakeTimeout dialer.Proxy = nil dialer.NetDialContext = safeDialContext return &LongConnClient{ - url: webSocketURL, - handler: handler, - dialer: &dialer, - messages: make(chan *im.IncomingMessage, webSocketMessageQueueSize), + channelID: channelID, + url: webSocketURL, + handler: handler, + dialer: &dialer, + messages: make(chan *im.IncomingMessage, webSocketMessageQueueSize), + maxConnectionAge: webSocketMaxConnectionAge, } } func (c *LongConnClient) Start(ctx context.Context) error { - logger.Infof(ctx, "[IM] Yunzhijia WebSocket connecting") + logger.Infof(ctx, "[IM] Yunzhijia WebSocket connecting channel_id=%s", c.channelID) go c.handleMessages(ctx) attempt := 0 for { @@ -72,18 +91,22 @@ func (c *LongConnClient) Start(ctx context.Context) error { return nil } - connectedAt := time.Now() + startedAt := time.Now() err := c.connectAndRun(ctx) if ctx.Err() != nil || c.closed.Load() { return nil } - if time.Since(connectedAt) >= webSocketReconnectDelays[len(webSocketReconnectDelays)-1] { + if time.Since(startedAt) >= webSocketReconnectDelays[len(webSocketReconnectDelays)-1] { attempt = 0 } delay := webSocketReconnectDelay(attempt) attempt++ - logger.Warnf(ctx, "[Yunzhijia] WebSocket connection lost: %v; reconnecting in %v", err, delay) + reconnectCount := c.reconnectCount.Add(1) + logger.Warnf(ctx, + "[Yunzhijia] WebSocket connection lost channel_id=%s reconnect_count=%d reason=%v; reconnecting in %v; %s", + c.channelID, reconnectCount, err, delay, c.healthStatus(), + ) select { case <-ctx.Done(): return nil @@ -118,21 +141,39 @@ func (c *LongConnClient) connectAndRun(ctx context.Context) error { conn.SetReadLimit(webSocketMaxMessageSize) _ = conn.SetReadDeadline(time.Now().Add(webSocketReadTimeout)) conn.SetPongHandler(func(string) error { + c.lastPongAt.Store(time.Now().UnixNano()) return conn.SetReadDeadline(time.Now().Add(webSocketReadTimeout)) }) - heartbeatCtx, cancelHeartbeat := context.WithCancel(ctx) - defer cancelHeartbeat() - go c.heartbeatLoop(heartbeatCtx, conn) + now := time.Now() + c.connectedAt.Store(now.UnixNano()) + c.lastFrameAt.Store(0) + c.lastPongAt.Store(0) + c.lastBusinessMessageAt.Store(0) + logger.Infof(ctx, + "[IM] Yunzhijia WebSocket connected channel_id=%s connected_at=%s", + c.channelID, now.Format(time.RFC3339Nano), + ) - logger.Infof(ctx, "[IM] Yunzhijia WebSocket connected") + connectionCtx, cancelConnection := context.WithCancel(ctx) + defer cancelConnection() + maxAgeReached := make(chan struct{}) + go c.heartbeatLoop(connectionCtx, conn) + go c.maxConnectionAgeLoop(connectionCtx, conn, maxAgeReached) invalidFrames := 0 for { messageType, data, err := conn.ReadMessage() if err != nil { + select { + case <-maxAgeReached: + return fmt.Errorf("websocket maximum connection age %s exceeded", c.maxConnectionAge) + default: + } return fmt.Errorf("read websocket message: %w", err) } - _ = conn.SetReadDeadline(time.Now().Add(webSocketReadTimeout)) + now := time.Now() + c.lastFrameAt.Store(now.UnixNano()) + _ = conn.SetReadDeadline(now.Add(webSocketReadTimeout)) if messageType != ws.TextMessage { invalidFrames++ @@ -145,13 +186,14 @@ func (c *LongConnClient) connectAndRun(ctx context.Context) error { frame, err := parseWebSocketFrame(data) if err != nil { invalidFrames++ - logger.Warnf(ctx, "[Yunzhijia] Invalid WebSocket frame: %v", err) + logger.Warnf(ctx, "[Yunzhijia] Invalid WebSocket frame channel_id=%s: %v", c.channelID, err) if invalidFrames >= webSocketMaxInvalidFrames { return fmt.Errorf("too many invalid websocket frames") } continue } invalidFrames = 0 + c.logFrame(ctx, frame) if len(frame.ack) > 0 { if err := c.writeText(conn, frame.ack); err != nil { @@ -161,6 +203,7 @@ func (c *LongConnClient) connectAndRun(ctx context.Context) error { if frame.message == nil { continue } + c.lastBusinessMessageAt.Store(time.Now().UnixNano()) incoming := toIncomingMessage(ctx, frame.message) if incoming == nil { @@ -174,6 +217,46 @@ func (c *LongConnClient) connectAndRun(ctx context.Context) error { } } +func (c *LongConnClient) maxConnectionAgeLoop(ctx context.Context, conn *ws.Conn, maxAgeReached chan<- struct{}) { + timer := time.NewTimer(c.maxConnectionAge) + defer timer.Stop() + select { + case <-ctx.Done(): + return + case <-timer.C: + close(maxAgeReached) + logger.Infof(ctx, + "[Yunzhijia] WebSocket rotating after maximum connection age channel_id=%s max_connection_age=%s %s", + c.channelID, c.maxConnectionAge, c.healthStatus(), + ) + _ = conn.Close() + } +} + +func (c *LongConnClient) logFrame(ctx context.Context, frame *webSocketFrame) { + logger.Debugf(ctx, + "[Yunzhijia] WebSocket frame channel_id=%s cmd=%s type=%s event=%s seq=%d has_seq=%t need_ack=%t has_message=%t", + c.channelID, frame.cmd, frame.typeName, frame.event, frame.seq, frame.hasSeq, frame.needAck, frame.message != nil, + ) +} + +func (c *LongConnClient) healthStatus() string { + return fmt.Sprintf( + "connected_at=%s last_frame_at=%s last_pong_at=%s last_business_message_at=%s", + c.healthTime(c.connectedAt.Load()), + c.healthTime(c.lastFrameAt.Load()), + c.healthTime(c.lastPongAt.Load()), + c.healthTime(c.lastBusinessMessageAt.Load()), + ) +} + +func (c *LongConnClient) healthTime(unixNano int64) string { + if unixNano == 0 { + return "" + } + return time.Unix(0, unixNano).UTC().Format(time.RFC3339Nano) +} + func (c *LongConnClient) handleMessages(ctx context.Context) { for { select { @@ -196,7 +279,10 @@ func (c *LongConnClient) heartbeatLoop(ctx context.Context, conn *ws.Conn) { return case <-ticker.C: if err := conn.WriteControl(ws.PingMessage, nil, time.Now().Add(5*time.Second)); err != nil { - logger.Warnf(ctx, "[Yunzhijia] WebSocket heartbeat failed: %v", err) + logger.Warnf(ctx, + "[Yunzhijia] WebSocket heartbeat failed channel_id=%s reason=%v; %s", + c.channelID, err, c.healthStatus(), + ) _ = conn.Close() return } @@ -257,21 +343,27 @@ func parseWebSocketFrame(data []byte) (*webSocketFrame, error) { return nil, fmt.Errorf("decode frame: %w", err) } + typeName := rawString(fields["type"]) + cmd := strings.ToLower(strings.TrimSpace(rawString(fields["cmd"]))) + event := strings.ToLower(strings.TrimSpace(rawString(fields["event"]))) + frame := &webSocketFrame{cmd: cmd, typeName: typeName, event: event} + _ = json.Unmarshal(fields["needAck"], &frame.needAck) + frame.hasSeq = json.Unmarshal(fields["seq"], &frame.seq) == nil if msg := decodeBusinessMessage(trimmed); msg != nil { - return &webSocketFrame{message: msg}, nil + frame.message = msg + return frame, nil } - typeName := rawString(fields["type"]) if strings.EqualFold(typeName, "robotMessage") { if msg := decodeBusinessMessage(fields["msg"]); msg != nil { - return &webSocketFrame{message: msg}, nil + frame.message = msg + return frame, nil } return nil, fmt.Errorf("robotMessage envelope has invalid msg") } - cmd := strings.ToLower(strings.TrimSpace(rawString(fields["cmd"]))) typeName = strings.ToLower(strings.TrimSpace(typeName)) - event := strings.ToLower(strings.TrimSpace(rawString(fields["event"]))) + frame.typeName = typeName control := cmd if control == "" { control = typeName @@ -280,16 +372,13 @@ func parseWebSocketFrame(data []byte) (*webSocketFrame, error) { control = event } - frame := &webSocketFrame{control: control} + frame.control = control if cmd == "directpush" || typeName == "msgchg" { - var needAck bool - _ = json.Unmarshal(fields["needAck"], &needAck) - var seq int64 - if needAck && json.Unmarshal(fields["seq"], &seq) == nil { + if frame.needAck && frame.hasSeq { frame.ack, _ = json.Marshal(struct { Cmd string `json:"cmd"` Seq int64 `json:"seq"` - }{Cmd: "ack", Seq: seq}) + }{Cmd: "ack", Seq: frame.seq}) } return frame, nil } diff --git a/internal/im/yunzhijia/websocket_test.go b/internal/im/yunzhijia/websocket_test.go index 5b9b23273..b2cbf8929 100644 --- a/internal/im/yunzhijia/websocket_test.go +++ b/internal/im/yunzhijia/websocket_test.go @@ -64,6 +64,47 @@ func TestParseWebSocketFrameBuildsDirectPushACK(t *testing.T) { } } +func TestParseWebSocketFrameAcknowledgesControlFramesWithPayloadFields(t *testing.T) { + testCases := []struct { + name string + input string + seq int64 + }{ + { + name: "directPush with msg", + input: `{"cmd":"directPush","needAck":true,"seq":43,"msg":{"content":"not logged"}}`, + seq: 43, + }, + { + name: "msgChg with data", + input: `{"type":"msgChg","needAck":true,"seq":44,"data":{"content":"not logged"}}`, + seq: 44, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + frame, err := parseWebSocketFrame([]byte(tc.input)) + if err != nil { + t.Fatalf("parseWebSocketFrame() error = %v", err) + } + if frame.message != nil { + t.Fatalf("control frame unexpectedly produced message: %#v", frame.message) + } + var ack struct { + Cmd string `json:"cmd"` + Seq int64 `json:"seq"` + } + if err := json.Unmarshal(frame.ack, &ack); err != nil { + t.Fatalf("unmarshal ack: %v", err) + } + if ack.Cmd != "ack" || ack.Seq != tc.seq { + t.Fatalf("ack = %#v, want cmd=ack seq=%d", ack, tc.seq) + } + }) + } +} + func TestParseWebSocketFrameControlAndInvalid(t *testing.T) { frame, err := parseWebSocketFrame([]byte(`{"event":"pong"}`)) if err != nil || frame.control != "pong" { @@ -104,13 +145,15 @@ func TestWebSocketReconnectDelayCaps(t *testing.T) { func TestLongConnClientProcessesMessagesInOrder(t *testing.T) { started := make(chan string, 2) releaseFirst := make(chan struct{}) - client := NewLongConnClient("wss://example.com/ws", func(_ context.Context, msg *im.IncomingMessage) error { - started <- msg.MessageID - if msg.MessageID == "first" { - <-releaseFirst - } - return nil - }) + client := NewLongConnClient( + "channel-1", "wss://example.com/ws", func(_ context.Context, msg *im.IncomingMessage) error { + started <- msg.MessageID + if msg.MessageID == "first" { + <-releaseFirst + } + return nil + }, + ) ctx, cancel := context.WithCancel(t.Context()) defer cancel() @@ -163,9 +206,11 @@ func TestLongConnClientAcknowledgesFrameAndStops(t *testing.T) { })) defer server.Close() - client := NewLongConnClient("ws"+strings.TrimPrefix(server.URL, "http"), func(context.Context, *im.IncomingMessage) error { - return nil - }) + client := NewLongConnClient( + "channel-1", "ws"+strings.TrimPrefix(server.URL, "http"), func(context.Context, *im.IncomingMessage) error { + return nil + }, + ) dialer := &net.Dialer{} client.dialer.NetDialContext = dialer.DialContext @@ -189,3 +234,57 @@ func TestLongConnClientAcknowledgesFrameAndStops(t *testing.T) { t.Fatal("websocket client did not stop promptly") } } + +func TestLongConnClientReconnectsAfterMaxConnectionAge(t *testing.T) { + connections := make(chan struct{}, 2) + upgrader := ws.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + connections <- struct{}{} + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + })) + defer server.Close() + + client := NewLongConnClient( + "channel-1", "ws"+strings.TrimPrefix(server.URL, "http"), func(context.Context, *im.IncomingMessage) error { + return nil + }, + ) + client.maxConnectionAge = 25 * time.Millisecond + dialer := &net.Dialer{} + client.dialer.NetDialContext = dialer.DialContext + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + done := make(chan error, 1) + go func() { done <- client.Start(ctx) }() + + for i := 0; i < 2; i++ { + select { + case <-connections: + case <-time.After(2 * time.Second): + t.Fatalf("connection %d was not established", i+1) + } + } + if got := client.reconnectCount.Load(); got != 1 { + t.Fatalf("reconnect count = %d, want 1 after max-age rotation", got) + } + + client.Stop() + select { + case err := <-done: + if err != nil { + t.Fatalf("Start() error after stop = %v", err) + } + case <-time.After(time.Second): + t.Fatal("websocket client did not stop promptly") + } +}