From 6b58b7f233911a0a5b05528fb819962e770f97cf Mon Sep 17 00:00:00 2001
From: Roger Ferdinan <rogerferdian@gmail.com>
Date: Thu, 25 Sep 2025 07:24:57 +0700
Subject: [PATCH] fix: data race hotfix

---
 internal/hub.go     | 48 ++++++++++++++++++++++++++-------------------
 v1/server/server.go |  1 -
 2 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/internal/hub.go b/internal/hub.go
index 74d3294..597a30b 100644
--- a/internal/hub.go
+++ b/internal/hub.go
@@ -122,14 +122,6 @@ func WritePump(c *Client, h *Hub) {
 			c.Conn.Close()
 		}()
 
-		c.Conn.SetReadLimit(1024)
-		c.Conn.SetPongHandler(func(string) error {
-			if err := c.Conn.WriteMessage(websocket.PongMessage, []byte{}); err != nil {
-				return fmt.Errorf("failed to send pong: %v", err)
-			}
-			return nil
-		})
-
 		ticker := time.NewTicker(30 * time.Second)
 		defer ticker.Stop()
 
@@ -142,19 +134,13 @@ func WritePump(c *Client, h *Hub) {
 						return fmt.Errorf("message not ok")
 					}
 
-					w, err := c.Conn.NextWriter(websocket.TextMessage)
-					if err != nil {
-						return fmt.Errorf("failed to get writer: %q", err)
-					}
-
-					w.Write(message)
-
-					if err := w.Close(); err != nil {
+					if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
 						return err
 					}
+
 					return nil
-				}); err == nil {
-					continue
+				}); err != nil {
+					return
 				}
 			case <-ticker.C:
 				if err := c.mu.WriteHandler(func() error {
@@ -172,8 +158,28 @@ func WritePump(c *Client, h *Hub) {
 
 func ReadPump(c *Client) {
 	go func() {
+		defer func() {
+			c.Conn.Close()
+		}()
+
+		c.Conn.SetReadLimit(1024)
+		c.Conn.SetPongHandler(func(string) error {
+			return c.mu.WriteHandler(func() error {
+				if err := c.Conn.WriteMessage(websocket.PongMessage, []byte{}); err != nil {
+					return fmt.Errorf("failed to send pong: %v", err)
+				}
+				return nil
+			})
+		})
 		for {
-			_, message, err := c.Conn.ReadMessage()
+			var messageType int
+			var message []byte
+			var err error
+			// c.mu.ReadHandler(func() error {
+			messageType, message, err = c.Conn.ReadMessage()
+			// 	return nil
+			// })
+
 			if err != nil {
 				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
 					log.Printf("WebSocket error: %v", err)
@@ -181,7 +187,9 @@ func ReadPump(c *Client) {
 				break
 			}
 
-			fmt.Println(string(message))
+			if messageType == websocket.TextMessage {
+				fmt.Printf("Received: %s\n", message)
+			}
 		}
 	}()
 
diff --git a/v1/server/server.go b/v1/server/server.go
index d8b70f1..c83aa78 100644
--- a/v1/server/server.go
+++ b/v1/server/server.go
@@ -64,7 +64,6 @@ func (b *SafeWebsocketServerBuilder) HandleFuncWebsocket(
 		}
 
 		subscribedPath := strings.TrimPrefix(r.URL.Path, pattern)
-		fmt.Println(subscribedPath)
 		if subscribedPath == "" {
 			http.Error(w, "invalid path", http.StatusBadRequest)
 			return