improve websocket connection according to gorrila example
This commit is contained in:
Adrian Zürcher
@@ -4,30 +4,51 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
json_dataModels "github.com/tecamino/tecamino-json_data/models"
|
||||
)
|
||||
|
||||
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 {
|
||||
conn *websocket.Conn
|
||||
writeMu sync.Mutex
|
||||
OnMessage func(data []byte)
|
||||
OnOpen func()
|
||||
OnClose func(code int, reason string)
|
||||
OnError func(err error)
|
||||
OnPing func()
|
||||
OnPong func()
|
||||
timeout time.Duration
|
||||
ip string
|
||||
port uint
|
||||
Connected bool
|
||||
conn *websocket.Conn
|
||||
OnOpen func()
|
||||
OnMessage func(data []byte)
|
||||
OnClose func(code int, reason string)
|
||||
OnError func(err error)
|
||||
OnPing func()
|
||||
OnPong func()
|
||||
sendPong chan string
|
||||
send chan []byte
|
||||
unregister chan []byte
|
||||
}
|
||||
|
||||
// Connect to websocket server
|
||||
// ip: ip address of server
|
||||
func NewClient(ip, id string, port uint) (*Client, error) {
|
||||
url := fmt.Sprintf("ws://%s:%d/ws?id=%s", ip, port, id)
|
||||
c := &Client{}
|
||||
c := &Client{
|
||||
ip: ip,
|
||||
port: port,
|
||||
Connected: true,
|
||||
sendPong: make(chan string),
|
||||
send: make(chan []byte),
|
||||
unregister: make(chan []byte),
|
||||
}
|
||||
|
||||
dialer := websocket.DefaultDialer
|
||||
conn, resp, err := dialer.Dial(url, nil)
|
||||
@@ -40,20 +61,25 @@ func NewClient(ip, id string, port uint) (*Client, error) {
|
||||
c.conn = conn
|
||||
|
||||
// Setup control handlers
|
||||
c.conn.SetPingHandler(func(appData string) error {
|
||||
conn.SetPingHandler(func(appData string) error {
|
||||
if c.OnPing != nil {
|
||||
c.OnPing()
|
||||
}
|
||||
return c.conn.WriteMessage(websocket.PongMessage, nil)
|
||||
conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
conn.SetReadDeadline(time.Now().Add(writeWait))
|
||||
c.sendPong <- appData
|
||||
return nil
|
||||
})
|
||||
c.conn.SetPongHandler(func(appData string) error {
|
||||
c.conn.SetReadDeadline(time.Now().Add(c.timeout))
|
||||
|
||||
conn.SetPongHandler(func(string) error {
|
||||
conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
if c.OnPong != nil {
|
||||
c.OnPong()
|
||||
}
|
||||
return nil
|
||||
})
|
||||
c.conn.SetCloseHandler(func(code int, text string) error {
|
||||
|
||||
conn.SetCloseHandler(func(code int, text string) error {
|
||||
if c.OnClose != nil {
|
||||
c.OnClose(code, text)
|
||||
}
|
||||
@@ -63,98 +89,103 @@ func NewClient(ip, id string, port uint) (*Client, error) {
|
||||
if c.OnOpen != nil {
|
||||
c.OnOpen()
|
||||
}
|
||||
|
||||
// Start reading messages from client
|
||||
go c.Read()
|
||||
go c.Write()
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func (c *Client) Connect(timeout uint) {
|
||||
if timeout > 0 {
|
||||
fmt.Println(1234, timeout)
|
||||
c.timeout = time.Duration(timeout) * time.Second
|
||||
func (c *Client) Read() {
|
||||
if c.OnOpen != nil {
|
||||
c.OnOpen()
|
||||
}
|
||||
|
||||
go c.pingLoop()
|
||||
|
||||
c.conn.SetReadDeadline(time.Now().Add(c.timeout))
|
||||
go func() {
|
||||
|
||||
for {
|
||||
msgType, msg, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||||
log.Println("WebSocket closed:", err)
|
||||
}
|
||||
if c.OnError != nil {
|
||||
c.OnError(fmt.Errorf("read error: %w", err))
|
||||
}
|
||||
return
|
||||
}
|
||||
switch msgType {
|
||||
case websocket.TextMessage, websocket.BinaryMessage:
|
||||
if c.OnMessage != nil {
|
||||
c.OnMessage(msg)
|
||||
}
|
||||
default:
|
||||
log.Printf("Unhandled message type: %d", msgType)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (c *Client) pingLoop() {
|
||||
interval := c.timeout / 2
|
||||
if interval <= 0 {
|
||||
interval = 5 * time.Second
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
if err := c.Write(websocket.PingMessage, nil); err != nil {
|
||||
if c.OnError != nil {
|
||||
c.OnError(fmt.Errorf("ping error: %w", err))
|
||||
}
|
||||
c.conn.SetReadDeadline(time.Now().Add(writeWait))
|
||||
for c.Connected {
|
||||
msgType, msg, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
c.handleError(fmt.Errorf("read error (id:%s:%d): %w", c.ip, c.port, err))
|
||||
return
|
||||
}
|
||||
if c.OnPing != nil {
|
||||
c.OnPing()
|
||||
switch msgType {
|
||||
case websocket.CloseMessage:
|
||||
c.Close(websocket.CloseNormalClosure, "Client closed")
|
||||
return
|
||||
case websocket.TextMessage:
|
||||
if c.OnMessage != nil {
|
||||
c.OnMessage(msg)
|
||||
} else {
|
||||
log.Printf("Received message but no handler set (id:%s:%d): %s", c.ip, c.port, string(msg))
|
||||
}
|
||||
default:
|
||||
log.Printf("Unhandled message type %d (id:%s:%d)", msgType, c.ip, c.port)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) Write(msgType int, data []byte) error {
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
|
||||
c.conn.SetWriteDeadline(time.Now().Add(c.timeout))
|
||||
|
||||
if err := c.conn.WriteMessage(msgType, data); err != nil {
|
||||
if c.OnError != nil {
|
||||
c.OnError(err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close connection to websocket server
|
||||
func (c *Client) Close(code int, reason string) {
|
||||
if c.conn != nil {
|
||||
if c.OnClose != nil {
|
||||
c.OnClose(code, reason)
|
||||
}
|
||||
c.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, reason))
|
||||
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:%d closed channel", c.ip, c.port))
|
||||
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, ok := <-c.sendPong:
|
||||
if ok {
|
||||
c.conn.WriteMessage(websocket.PongMessage, []byte(message))
|
||||
}
|
||||
case message := <-c.unregister:
|
||||
c.conn.WriteMessage(websocket.CloseMessage, message)
|
||||
c.Connected = false
|
||||
close(c.send)
|
||||
close(c.sendPong)
|
||||
close(c.unregister)
|
||||
return
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) SendData(data any) error {
|
||||
c.conn.SetWriteDeadline(time.Now().Add(c.timeout))
|
||||
if err := c.conn.WriteJSON(data); err != nil {
|
||||
if c.OnError != nil {
|
||||
c.OnError(err)
|
||||
}
|
||||
func (c *Client) SendRequest(req *json_dataModels.Request) error {
|
||||
if !c.Connected {
|
||||
return nil
|
||||
}
|
||||
|
||||
data, err := json.Marshal(*req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c.send <- data
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -163,3 +194,30 @@ func (c *Client) ReadJsonData(data []byte) (json_dataModels.Response, error) {
|
||||
err := json.Unmarshal(data, &resp)
|
||||
return resp, err
|
||||
}
|
||||
|
||||
// Close connection to websocket server
|
||||
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")
|
||||
}
|
||||
if c.OnClose != nil {
|
||||
c.OnClose(code, reason)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) handleError(err error) {
|
||||
if c.OnError != nil {
|
||||
c.OnError(err)
|
||||
}
|
||||
if err := c.Close(websocket.CloseInternalServerErr, err.Error()); err != nil {
|
||||
if c.OnError != nil {
|
||||
c.OnError(err)
|
||||
} else {
|
||||
fmt.Println("error: ", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user