improvement websocket according to gorilla example

websocket/models/client.go

@@ -13,17 +13,29 @@ import (
 
 var Origins []string = []string{"*"}
 
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 10 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 10 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+)
+
 type Client struct {
-	Id        string
-	Connected bool            `json:"connected"`
-	Conn      *websocket.Conn `json:"-"`
-	OnOpen    func()
-	OnMessage func(data []byte)
-	OnClose   func(code int, reason string)
-	OnError   func(err error)
-	OnPing    func()
-	OnPong    func()
-	timeout   time.Duration
+	Id         string
+	Connected  bool            `json:"connected"`
+	Conn       *websocket.Conn `json:"-"`
+	OnOpen     func()
+	OnMessage  func(data []byte)
+	OnClose    func(code int, reason string)
+	OnError    func(err error)
+	OnPing     func()
+	OnPong     func()
+	send       chan []byte
+	unregister chan []byte
 }
 
 var upgrader = websocket.Upgrader{
@@ -52,44 +64,42 @@ func ConnectNewClient(id string, c *gin.Context) (*Client, error) {
 	}
 
 	client := &Client{
-		Id:        id,
-		Connected: true,
-		Conn:      conn,
-		timeout:   5,
+		Id:         id,
+		Connected:  true,
+		Conn:       conn,
+		send:       make(chan []byte),
+		unregister: make(chan []byte),
 	}
 
 	conn.SetPingHandler(func(appData string) error {
 		if client.OnPing != nil {
 			client.OnPing()
 		}
-		conn.SetWriteDeadline(time.Now().Add(client.timeout))
-		conn.SetReadDeadline(time.Now().Add(client.timeout))
+		conn.SetWriteDeadline(time.Now().Add(writeWait))
+		conn.SetReadDeadline(time.Now().Add(writeWait))
 		return conn.WriteMessage(websocket.PongMessage, []byte(appData))
 	})
 
-	conn.SetPongHandler(func(appData string) error {
+	conn.SetPongHandler(func(string) error {
+		conn.SetReadDeadline(time.Now().Add(pongWait))
 		if client.OnPong != nil {
 			client.OnPong()
 		}
-		conn.SetReadDeadline(time.Now().Add(client.timeout))
 		return nil
 	})
 
 	// Start reading messages from client
-	go client.Listen(7)
+	go client.Read()
+	go client.Write()
 	return client, nil
 }
 
-func (c *Client) Listen(timeout uint) {
-	if timeout > 0 {
-		c.timeout = time.Duration(timeout) * time.Second
-	}
-
+func (c *Client) Read() {
 	if c.OnOpen != nil {
 		c.OnOpen()
 	}
 
-	c.Conn.SetReadDeadline(time.Now().Add(c.timeout))
+	c.Conn.SetReadDeadline(time.Now().Add(writeWait))
 	for c.Connected {
 		msgType, msg, err := c.Conn.ReadMessage()
 		if err != nil {
@@ -98,16 +108,15 @@ func (c *Client) Listen(timeout uint) {
 		}
 		switch msgType {
 		case websocket.CloseMessage:
-			c.handleClose(1000, "Client closed")
+			c.Close(websocket.CloseNormalClosure, "Client closed")
 			return
 		case websocket.TextMessage:
-			if isPing := c.handleJsonPing(msg); isPing {
-				continue
-			}
-			if c.OnMessage != nil {
-				c.OnMessage(msg)
-			} else {
-				log.Printf("Received message but no handler set (id:%s): %s", c.Id, string(msg))
+			if isPing := c.handleJsonPing(msg); !isPing {
+				if c.OnMessage != nil {
+					c.OnMessage(msg)
+				} else {
+					log.Printf("Received message but no handler set (id:%s): %s", c.Id, string(msg))
+				}
 			}
 		default:
 			log.Printf("Unhandled message type %d (id:%s)", msgType, c.Id)
@@ -115,14 +124,66 @@ func (c *Client) Listen(timeout uint) {
 	}
 }
 
+func (c *Client) Write() {
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		ticker.Stop()
+		c.Conn.Close()
+	}()
+	for {
+
+		select {
+		case message, ok := <-c.send:
+			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if !ok {
+				// The hub closed the channel.
+				if err := c.Conn.WriteMessage(websocket.CloseMessage, []byte{}); err != nil {
+					c.handleError(err)
+					return
+				}
+				c.handleError(fmt.Errorf("server %s closed channel", c.Id))
+				return
+			} else {
+				if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
+					c.handleError(err)
+					return
+				}
+			}
+		case <-ticker.C:
+			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				c.handleError(err)
+				return
+			}
+
+			if c.OnPing != nil {
+				c.OnPing()
+			}
+		case message := <-c.unregister:
+			c.Conn.WriteMessage(websocket.CloseMessage, message)
+			c.Connected = false
+			close(c.send)
+			close(c.unregister)
+			return
+		}
+
+	}
+}
+
 func (c *Client) handleJsonPing(msg []byte) (isPing bool) {
 	var wsMsg WSMessage
 	err := json.Unmarshal(msg, &wsMsg)
+
 	if err == nil && wsMsg.IsPing() {
-		c.Conn.SetReadDeadline(time.Now().Add(c.timeout))
+		c.Conn.SetReadDeadline(time.Now().Add(writeWait))
 		// Respond with pong JSON
-		c.Conn.SetWriteDeadline(time.Now().Add(c.timeout))
-		err = c.Conn.WriteMessage(websocket.TextMessage, GetPongByteSlice())
+		select {
+		case c.send <- GetPongByteSlice():
+		default:
+			// optional: log or handle if send buffer is full
+			c.handleError(fmt.Errorf("failed to queue pong message"))
+			return
+		}
 		if err != nil {
 			c.handleError(fmt.Errorf("write pong error: %w", err))
 			return
@@ -132,33 +193,39 @@ func (c *Client) handleJsonPing(msg []byte) (isPing bool) {
 		}
 		isPing = true
 	}
+
 	return
 }
 
-func (c *Client) SendResponse(data []byte, timeout uint) error {
-	c.Conn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
-	return c.Conn.WriteMessage(websocket.TextMessage, data)
-}
-
-func (c *Client) Close(code int, reason string) {
-	c.handleClose(code, reason)
-}
-
-func (c *Client) handleClose(code int, text string) {
+func (c *Client) SendResponse(data []byte) {
 	if !c.Connected {
 		return
 	}
-	c.Connected = false
-	if c.OnClose != nil {
-		c.OnClose(code, text)
+	c.send <- data
+}
+
+func (c *Client) Close(code int, reason string) error {
+	closeMsg := websocket.FormatCloseMessage(code, reason)
+	select {
+	case c.unregister <- closeMsg: // Attempt to send
+	default: // If the channel is full, this runs
+		return fmt.Errorf("attempt close client socket failed")
 	}
-	c.Conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
-	c.Conn.Close()
+	if c.OnClose != nil {
+		c.OnClose(code, reason)
+	}
+	return nil
 }
 
 func (c *Client) handleError(err error) {
 	if c.OnError != nil {
 		c.OnError(err)
 	}
-	c.Close(websocket.CloseInternalServerErr, err.Error())
+	if err := c.Close(websocket.CloseInternalServerErr, err.Error()); err != nil {
+		if c.OnError != nil {
+			c.OnError(err)
+		} else {
+			fmt.Println("error: ", err)
+		}
+	}
 }