From 1b1bc0b7563a1d8a99cbb2552b01a8b70947bad1 Mon Sep 17 00:00:00 2001
From: Monet Lee <monet_lee@163.com>
Date: Tue, 9 Jul 2024 18:36:10 +0800
Subject: [PATCH] add context life cycle control.

---
 internal/msggateway/client.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/internal/msggateway/client.go b/internal/msggateway/client.go
index a901ebdfe..9d594c1fa 100644
--- a/internal/msggateway/client.go
+++ b/internal/msggateway/client.go
@@ -138,6 +138,7 @@ func (c *Client) readMessage() {
 		if returnErr != nil {
 			log.ZWarn(c.ctx, "readMessage", returnErr, "messageType", messageType)
 			c.closedErr = returnErr
+			<-c.ctx.Done()
 			return
 		}
 
@@ -145,6 +146,7 @@ func (c *Client) readMessage() {
 		if c.closed.Load() {
 			// The scenario where the connection has just been closed, but the coroutine has not exited
 			c.closedErr = ErrConnClosed
+			<-c.ctx.Done()
 			return
 		}
 
@@ -154,10 +156,12 @@ func (c *Client) readMessage() {
 			parseDataErr := c.handleMessage(message)
 			if parseDataErr != nil {
 				c.closedErr = parseDataErr
+				<-c.ctx.Done()
 				return
 			}
 		case MessageText:
 			c.closedErr = ErrNotSupportMessageProtocol
+			<-c.ctx.Done()
 			return
 
 		case PingMessage:
@@ -166,7 +170,9 @@ func (c *Client) readMessage() {
 
 		case CloseMessage:
 			c.closedErr = ErrClientClosed
+			<-c.ctx.Done()
 			return
+
 		default:
 		}
 	}