improvement websocket according to gorilla example

dbm/set.go

@@ -20,6 +20,7 @@ func (d *DBMHandler) Set(req *json_dataModels.Request, id string) {
 
 	for _, set := range req.Set {
 		dps := d.DBM.QueryDatapoints(1, set.Uuid, set.Path)
+
 		if len(dps) == 0 {
 			resp.SetError()
 			if resp.Message == "" {
@@ -27,8 +28,10 @@ func (d *DBMHandler) Set(req *json_dataModels.Request, id string) {
 			}
 			continue
 		}
+
 		for _, dp := range dps {
 			dp.UpdateValue(d.Conns, set.Value)
+
 			resp.AddSet(json_dataModels.Set{
 				Uuid:  dp.Uuid,
 				Path:  dp.Path,
@@ -36,6 +39,7 @@ func (d *DBMHandler) Set(req *json_dataModels.Request, id string) {
 			})
 		}
 	}
+
 	if err := d.Conns.SendResponse(id, resp); err != nil {
 		d.Log.Error("get.Set", err.Error())
 	}

dbm/subscribe.go

@@ -54,7 +54,6 @@ func (d *DBMHandler) Subscribe(req *json_dataModels.Request, id string) {
 			})
 		}
 	}
-
 	if err := d.Conns.SendResponse(id, resp); err != nil {
 		d.Log.Error("subscribe.Subscribe", err.Error())
 	}

dbm/webSocket.go

@@ -33,12 +33,15 @@ func (d *DBMHandler) WebSocket(c *gin.Context) {
 		if err != nil {
 			d.Log.Error("dbmHandler.webSocket.Websocket", "read json: "+err.Error())
 		}
+
 		// Sets
 		d.Get(request, id)
 		// Sets
 		d.Set(request, id)
+
 		// Subscribe
 		d.Subscribe(request, id)
+
 		// Unsubscribe
 		d.Unsubscribe(request, id)
 	}

models/datapoint.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"regexp"
 	"strings"
-	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -34,7 +33,6 @@ type Datapoint struct {
 	ReadWrite      json_dataModels.Rights  `json:"readWrite"`
 	Drivers        json_dataModels.Drivers `json:"drivers,omitempty"`
 	Subscriptions  Subscriptions           `json:"-"`
-	sync.RWMutex   `json:"-"`
 }
 
 func (d *Datapoint) Set(path string, set json_dataModels.Set) (bool, error) {
@@ -381,9 +379,6 @@ func (d *Datapoint) RemoveSubscribtion(client *wsModels.Client) {
 }
 
 func (d *Datapoint) Publish(eventType string) error {
-	d.RLock()
-	defer d.RUnlock()
-
 	for client := range d.Subscriptions {
 		r := json_data.NewResponse()
 		r.AddPublish(json_dataModels.Publish{
@@ -392,14 +387,13 @@ func (d *Datapoint) Publish(eventType string) error {
 			Path:  d.Path,
 			Value: d.Value,
 		})
+
 		b, err := json.Marshal(r)
 		if err != nil {
 			return err
 		}
-		if err := client.SendResponse(b, 5); err != nil {
-			return err
-		}
 
+		client.SendResponse(b)
 	}
 	return nil
 }

websocket/clientHandler.go

@@ -60,14 +60,12 @@ func (c *ClientHandler) SendResponse(id string, r *json_dataModels.Response) err
 		return fmt.Errorf("client not found for id %s", id)
 
 	}
-	b, err := json.Marshal(r)
+
+	b, err := json.Marshal(*r)
 	if err != nil {
 		return err
 	}
-
-	if err := client.SendResponse(b, 5); err != nil {
-		return err
-	}
+	client.SendResponse(b)
 
 	return nil
 }

websocket/models/client.go

@@ -13,6 +13,17 @@ 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"`
@@ -23,7 +34,8 @@ type Client struct {
 	OnError    func(err error)
 	OnPing     func()
 	OnPong     func()
-	timeout   time.Duration
+	send       chan []byte
+	unregister chan []byte
 }
 
 var upgrader = websocket.Upgrader{
@@ -55,41 +67,39 @@ func ConnectNewClient(id string, c *gin.Context) (*Client, error) {
 		Id:         id,
 		Connected:  true,
 		Conn:       conn,
-		timeout:   5,
+		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,31 +108,82 @@ 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 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)
 		}
 	}
 }
 
+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)
+		}
+	}
 }