112 lines
2.7 KiB
Go
112 lines
2.7 KiB
Go
package main
|
|
|
|
import (
|
|
"log"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type WSConnection struct {
|
|
id int32
|
|
conn *websocket.Conn
|
|
writeLock sync.Mutex
|
|
alive bool
|
|
ReadChan chan string
|
|
WriteChan chan string
|
|
ErrorChan chan error
|
|
IdleTimeout time.Duration
|
|
PingInterval time.Duration
|
|
}
|
|
|
|
func NewConn(conn *websocket.Conn, server *WSServer) *WSConnection {
|
|
wsconn := &WSConnection{
|
|
id: server.clientId,
|
|
conn: conn,
|
|
alive: true,
|
|
IdleTimeout: server.IdleTimeout,
|
|
PingInterval: server.PingInterval,
|
|
ReadChan: make(chan string, 1024),
|
|
WriteChan: make(chan string, 1024),
|
|
ErrorChan: make(chan error, 1),
|
|
}
|
|
atomic.AddInt32(&server.clientId, 1)
|
|
|
|
return wsconn
|
|
}
|
|
|
|
func (ws *WSConnection) Open() {
|
|
go ws.messageReader()
|
|
go ws.messageSender()
|
|
go ws.pinger()
|
|
ws.conn.SetPongHandler(func(string) error {
|
|
// log.Printf("Client %d: Pong OK", ws.id)
|
|
ws.conn.SetReadDeadline(time.Now().Add(ws.IdleTimeout))
|
|
return nil
|
|
})
|
|
}
|
|
|
|
func (ws *WSConnection) messageReader() {
|
|
log.Printf("Client %d: Reading messages", ws.id)
|
|
for {
|
|
_, message, err := ws.conn.ReadMessage()
|
|
if !ws.alive {
|
|
return
|
|
}
|
|
ws.conn.SetReadDeadline(time.Now().Add(ws.IdleTimeout))
|
|
if err != nil {
|
|
ws.ErrorChan <- err
|
|
return
|
|
}
|
|
log.Printf("Client %d: Received: %s, %d in output channel", ws.id, message, len(ws.ReadChan))
|
|
ws.ReadChan <- string(message)
|
|
}
|
|
}
|
|
|
|
func (ws *WSConnection) messageSender() {
|
|
log.Printf("Client %d: Sending messages", ws.id)
|
|
for {
|
|
func() {
|
|
msg, ok := <-ws.WriteChan
|
|
if !ok || !ws.alive {
|
|
return
|
|
}
|
|
ws.writeLock.Lock()
|
|
defer ws.writeLock.Unlock()
|
|
|
|
ws.conn.SetWriteDeadline(time.Now().Add(ws.IdleTimeout))
|
|
log.Printf("Client %d: Sending: %s, %d in input channel", ws.id, msg, len(ws.WriteChan))
|
|
err := ws.conn.WriteMessage(websocket.TextMessage, []byte(msg))
|
|
if err != nil {
|
|
log.Printf("Client %d: Error during message writing: %v", ws.id, err)
|
|
ws.ErrorChan <- err
|
|
return
|
|
}
|
|
}()
|
|
}
|
|
}
|
|
|
|
func (ws *WSConnection) pinger() {
|
|
log.Printf("Client %d: Starting pinger, sleeping for %v", ws.id, ws.PingInterval)
|
|
for {
|
|
time.Sleep(ws.PingInterval)
|
|
if !ws.alive {
|
|
return
|
|
}
|
|
|
|
// log.Printf("Client %d: Ping", ws.id)
|
|
ws.writeLock.Lock()
|
|
err := ws.conn.WriteMessage(websocket.PingMessage, nil)
|
|
if err != nil {
|
|
log.Printf("Client %d: Error during ping: %+v", ws.id, err)
|
|
ws.ErrorChan <- err
|
|
return
|
|
}
|
|
ws.conn.SetWriteDeadline(time.Now().Add(ws.IdleTimeout))
|
|
ws.writeLock.Unlock()
|
|
// log.Printf("Client %d: Ping OK", ws.id)
|
|
}
|
|
}
|