From a91a8b9c8b8fd5b8b8926bab15ebd8733ad98fe6 Mon Sep 17 00:00:00 2001
From: Monet Lee <monet_lee@163.com>
Date: Wed, 10 Jul 2024 10:44:06 +0800
Subject: [PATCH] update heartbeat logic.

---
 internal/msggateway/client.go | 31 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/internal/msggateway/client.go b/internal/msggateway/client.go
index ca72bfb84..7a4303410 100644
--- a/internal/msggateway/client.go
+++ b/internal/msggateway/client.go
@@ -90,7 +90,7 @@ func (c *Client) ResetClient(ctx *UserConnContext, conn LongConn, longConnServer
 	c.closed.Store(false)
 	c.closedErr = nil
 	c.token = ctx.GetToken()
-	c.hbCtx, _ = context.WithTimeout(c.ctx, pongWait*2)
+	c.hbCtx, _ = context.WithCancel(c.ctx)
 }
 
 func (c *Client) pingHandler(_ string) error {
@@ -114,10 +114,7 @@ func (c *Client) readMessage() {
 	c.conn.SetReadLimit(maxMessageSize)
 	_ = c.conn.SetReadDeadline(pongWait)
 	c.conn.SetPingHandler(c.pingHandler)
-
-	if c.PlatformID == constant.WebPlatformID {
-		go c.heartbeat(c.hbCtx)
-	}
+	go c.heartbeat(c.hbCtx)
 
 	for {
 		log.ZDebug(c.ctx, "readMessage")
@@ -331,19 +328,21 @@ func (c *Client) writeBinaryMsg(resp Resp) error {
 }
 
 func (c *Client) heartbeat(ctx context.Context) {
-	log.ZDebug(ctx, "server initiative send heartbeat start.")
-	ticker := time.NewTicker(pingPeriod)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-ticker.C:
-			if err := c.conn.WriteMessage(PingMessage, nil); err != nil {
-				log.ZError(c.ctx, "send Ping Message error.", err)
+	if c.PlatformID == constant.WebPlatformID {
+		log.ZDebug(ctx, "server initiative send heartbeat start.")
+		ticker := time.NewTicker(pingPeriod)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				if err := c.conn.WriteMessage(PingMessage, nil); err != nil {
+					log.ZError(c.ctx, "send Ping Message error.", err)
+					return
+				}
+			case <-c.hbCtx.Done():
 				return
 			}
-		case <-c.hbCtx.Done():
-			return
 		}
 	}
 }