fix: readPump and writePump hotfix

internal/hub.go

@@ -8,6 +8,12 @@ import (
 	"github.com/gorilla/websocket"
 )
 
+const (
+	writeWait  = 10 * time.Second
+	pongWait   = 60 * time.Second
+	pingPeriod = 55 * time.Second
+)
+
 type Client struct {
 	Conn           *websocket.Conn
 	Send           chan []byte
@@ -19,7 +25,7 @@ type Client struct {
 func NewClient(conn *websocket.Conn, subscribedPath string) *Client {
 	return &Client{
 		Conn:           conn,
-		Send:           make(chan []byte, 1024),
+		Send:           make(chan []byte, 64),
 		SubscribedPath: subscribedPath,
 		done:           make(chan struct{}),
 		mu:             NewCustomRwMutex(),
@@ -31,7 +37,6 @@ type Hub struct {
 	Broadcast  chan []byte
 	Register   chan *Client
 	Unregister chan *Client
-	ClientData map[string]chan []byte
 	writeMu    *CustomRwMutex
 	readMu     *CustomRwMutex
 }
@@ -42,72 +47,29 @@ func NewHub() *Hub {
 		Register:   make(chan *Client),
 		Unregister: make(chan *Client),
 		Clients:    make(map[*Client]bool),
-		ClientData: make(map[string]chan []byte),
 		writeMu:    NewCustomRwMutex(),
 	}
 }
 
-func (h *Hub) AddDataChannel(dataID string) chan []byte {
-	ch := make(chan []byte, 256)
-	h.writeMu.WriteHandler(func() error {
-		if innerCh, ok := h.ClientData[dataID]; ok {
-			ch = innerCh
-			return nil
-		}
-
-		h.ClientData[dataID] = ch
-		log.Printf("Created data channel for: %s\n", dataID)
-		return nil
-	})
-
-	return ch
-}
-
-func (h *Hub) GetDataChannel(dataID string) (chan []byte, bool) {
-	var ch chan []byte
-	var ok bool
-	h.writeMu.ReadHandler(func() error {
-		innerCh, innerOk := h.ClientData[dataID]
-		ch = innerCh
-		ok = innerOk
-		return nil
-	})
-
-	return ch, ok
-}
-
-func (h *Hub) RemoveDataChannel(dataID string) {
-	h.writeMu.WriteHandler(func() error {
-		if ch, ok := h.ClientData[dataID]; ok {
-			close(ch)
-			delete(h.ClientData, dataID)
-			log.Printf("Removed data channel for: %s\n", dataID)
-		}
-		return nil
-	})
-}
-
 func (h *Hub) Run() {
 	go func() {
 		for {
 			select {
-			case c := <-h.Register:
-				h.Clients[c] = true
+			case client := <-h.Register:
+				h.Clients[client] = true
 				log.Println("Client registered")
 			case c := <-h.Unregister:
 				if _, ok := h.Clients[c]; ok {
 					delete(h.Clients, c)
 					close(c.Send)
-					c.Conn.Close()
-					log.Println("Client unregistered")
 				}
 			case message := <-h.Broadcast:
-				for c := range h.Clients {
+				for client := range h.Clients {
 					select {
-					case c.Send <- message:
+					case client.Send <- message:
 					default:
-						close(c.Send)
-						delete(h.Clients, c)
+						close(client.Send)
+						delete(h.Clients, client)
 					}
 				}
 			}
@@ -116,81 +78,58 @@ func (h *Hub) Run() {
 }
 
 func WritePump(c *Client, h *Hub) {
-	go func() {
-		defer func() {
-			h.Unregister <- c
-			c.Conn.Close()
-		}()
-
-		ticker := time.NewTicker(30 * time.Second)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case message, ok := <-c.Send:
-				if err := c.mu.WriteHandler(func() error {
-					if !ok {
-						c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
-						return fmt.Errorf("message not ok")
-					}
-
-					if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
-						return err
-					}
-
-					return nil
-				}); err != nil {
-					return
-				}
-			case <-ticker.C:
-				if err := c.mu.WriteHandler(func() error {
-					if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
-						return err
-					}
-					return nil
-				}); err != nil {
-					return
-				}
-			}
-		}
-	}()
-}
-
-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 {
-			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)
-				}
-				break
-			}
-
-			if messageType == websocket.TextMessage {
-				fmt.Printf("Received: %s\n", message)
-			}
-		}
+	pingTicker := time.NewTicker(pingPeriod)
+	defer func() {
+		h.Unregister <- c
+		pingTicker.Stop()
+		c.Conn.Close()
 	}()
 
+	for {
+		select {
+		case message, ok := <-c.Send:
+			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if !ok {
+				c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
+				return
+			}
+		case <-pingTicker.C:
+			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
+			}
+		}
+	}
+}
+
+func ReadPump(c *Client, h *Hub) {
+	defer func() {
+		h.Unregister <- c
+		c.Conn.Close()
+	}()
+
+	c.Conn.SetReadLimit(512)
+	c.Conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.Conn.SetPongHandler(func(string) error {
+		c.Conn.SetReadDeadline(time.Now().Add(pongWait))
+		return nil
+	})
+
+	for {
+		messageType, message, err := c.Conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
+				log.Printf("WebSocket error: %v", err)
+			}
+			break
+		}
+
+		if messageType == websocket.TextMessage {
+			fmt.Printf("Received: %s\n", message)
+		}
+	}
 }