fix the issue with the 'ReadMessage' thread not exiting when 'stopCh' is triggered

Author: jackjohn

websocket.go

@@ -66,36 +66,19 @@ var wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (don
 	c.SetReadLimit(655350)
 	doneCh = make(chan struct{})
 	stopCh = make(chan struct{})
+
 	go func() {
-		// This function will exit either on error from
-		// websocket.Conn.ReadMessage or when the stopC channel is
-		// closed by the client.
 		defer close(doneCh)
+		defer close(stopCh)
 		if WebsocketKeepalive {
 			keepAlive(c, WebsocketTimeout)
 		}
-		// Wait for the stopC channel to be closed.  We do that in a
-		// separate goroutine because ReadMessage is a blocking
-		// operation.
-		silent := false
-		go func() {
-			for {
-				_, message, err := c.ReadMessage()
-				if err != nil {
-					if !silent {
-						errHandler(err)
-					}
-					stopCh <- struct{}{}
-					return
-				}
-				handler(message)
-			}
-		}()
+		// Separate goroutine to handle message reading
+		go readMessages(c, handler, errHandler, stopCh)
 
 		for {
 			select {
 			case <-stopCh:
-				silent = true
 				return
 			case <-doneCh:
 			}
@@ -104,6 +87,22 @@ var wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (don
 	return
 }
 
+func readMessages(c *websocket.Conn, handler WsHandler, errHandler ErrHandler, stopCh chan struct{}) {
+	for {
+		select {
+		case <-stopCh:
+			return
+		default:
+			_, message, err := c.ReadMessage()
+			if err != nil {
+				errHandler(err)
+				return
+			}
+			handler(message)
+		}
+	}
+}
+
 func keepAlive(c *websocket.Conn, timeout time.Duration) {
 	ticker := time.NewTicker(timeout)