add context life cycle control.

openimsdk · Jul 9, 2024 · 1b1bc0b · 1b1bc0b
1 parent 557a5a3
commit 1b1bc0b
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/internal/msggateway/client.go b/internal/msggateway/client.go
@@ -138,13 +138,15 @@ func (c *Client) readMessage() {
 		if returnErr != nil {
 			log.ZWarn(c.ctx, "readMessage", returnErr, "messageType", messageType)
 			c.closedErr = returnErr
+			<-c.ctx.Done()
 			return
 		}
 
 		log.ZDebug(c.ctx, "readMessage", "messageType", messageType)
 		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:
 		}
 	}