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

websocket.go

@@ -51,57 +51,56 @@ func newWsConfig(endpoint string) *WsConfig {
 	}
 }
 
-var wsServe = func(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (doneCh, stopCh chan struct{}, err error) {
-	Dialer := websocket.Dialer{
-		Proxy:             http.ProxyFromEnvironment,
-		HandshakeTimeout:  45 * time.Second,
-		EnableCompression: false,
-	}
-	headers := http.Header{}
-	headers.Add("User-Agent", fmt.Sprintf("%s/%s", Name, Version))
-	c, _, err := Dialer.Dial(cfg.Endpoint, headers)
-	if err != nil {
-		return nil, nil, err
-	}
-	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)
-		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)
-			}
-		}()
+func wsServe(cfg *WsConfig, handler WsHandler, errHandler ErrHandler) (doneCh, stopCh chan struct{}, err error) {
+    Dialer := websocket.Dialer{
+        Proxy:             http.ProxyFromEnvironment,
+        HandshakeTimeout:  45 * time.Second,
+        EnableCompression: false,
+    }
+    headers := http.Header{}
+    headers.Add("User-Agent", fmt.Sprintf("%s/%s", Name, Version))
+    c, _, err := Dialer.Dial(cfg.Endpoint, headers)
+    if err != nil {
+        return nil, nil, err
+    }
+    c.SetReadLimit(655350)
+    doneCh = make(chan struct{})
+    stopCh = make(chan struct{})
+
+    go func() {
+        defer close(doneCh)
+        defer close(stopCh)
+        if WebsocketKeepalive {
+            keepAlive(c, WebsocketTimeout)
+        }
+        // Separate goroutine to handle message reading
+        go readMessages(c, handler, errHandler, stopCh)
+
+        for {
+            select {
+            case <-stopCh:
+                return
+            case <-doneCh:
+            }
+        }
+    }()
+    return
+}
 
-		for {
-			select {
-			case <-stopCh:
-				silent = true
+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
-			case <-doneCh:
 			}
+			handler(message)
 		}
-	}()
-	return
+	}
 }
 
 func keepAlive(c *websocket.Conn, timeout time.Duration) {