wings/api/websockets/socket.go

82 lines
1.6 KiB
Go
Raw Normal View History

2018-03-14 09:34:06 +00:00
package websockets
import "github.com/gorilla/websocket"
import (
"time"
log "github.com/sirupsen/logrus"
)
type Socket struct {
collection *Collection
conn *websocket.Conn
send chan []byte
}
func (s *Socket) readPump() {
defer func() {
s.collection.unregister <- s
s.conn.Close()
}()
s.conn.SetReadLimit(maxMessageSize)
s.conn.SetReadDeadline(time.Now().Add(pongWait))
s.conn.SetPongHandler(func(string) error {
s.conn.SetReadDeadline(time.Now().Add(pongWait))
return nil
})
for {
t, m, err := s.conn.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
log.WithError(err).Debug("Websocket closed unexpectedly.")
}
return
}
// Handle websocket responses somehow
if t == websocket.TextMessage {
log.Debug("Received websocket message: " + string(m))
}
}
}
func (s *Socket) writePump() {
ticker := time.NewTicker(pingPeriod)
defer func() {
ticker.Stop()
s.conn.Close()
}()
for {
select {
case m, ok := <-s.send:
s.conn.SetWriteDeadline(time.Now().Add(writeWait))
if !ok {
// The collection closed the channel
s.conn.WriteMessage(websocket.CloseMessage, []byte{})
return
}
w, err := s.conn.NextWriter(websocket.TextMessage)
if err != nil {
return
}
w.Write([]byte{'['})
w.Write(m)
n := len(s.send) - 1
for i := 0; i < n; i++ {
w.Write([]byte{','})
w.Write(<-s.send)
}
w.Write([]byte{']'})
if err := w.Close(); err != nil {
return
}
case <-ticker.C:
s.conn.SetWriteDeadline(time.Now().Add(writeWait))
if err := s.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
return
}
}
}
}