121 lines
2.1 KiB
Go
121 lines
2.1 KiB
Go
package websockets
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
log "github.com/sirupsen/logrus"
|
|
)
|
|
|
|
const (
|
|
writeWait = 10 * time.Second
|
|
pongWait = 60 * time.Second
|
|
|
|
pingPeriod = pongWait * 9 / 10
|
|
|
|
maxMessageSize = 512
|
|
)
|
|
|
|
var wsupgrader = websocket.Upgrader{
|
|
ReadBufferSize: 1024,
|
|
WriteBufferSize: 1024,
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
return true
|
|
},
|
|
}
|
|
|
|
type websocketMap map[*Socket]bool
|
|
|
|
type Collection struct {
|
|
sockets websocketMap
|
|
|
|
Broadcast chan Message
|
|
|
|
register chan *Socket
|
|
unregister chan *Socket
|
|
close chan bool
|
|
}
|
|
|
|
//var _ io.Writer = &Collection{}
|
|
|
|
func NewCollection() *Collection {
|
|
return &Collection{
|
|
Broadcast: make(chan Message),
|
|
register: make(chan *Socket),
|
|
unregister: make(chan *Socket),
|
|
close: make(chan bool),
|
|
sockets: make(websocketMap),
|
|
}
|
|
}
|
|
|
|
func (c *Collection) Upgrade(w http.ResponseWriter, r *http.Request) {
|
|
socket, err := wsupgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
log.WithError(err).Error("Failed to upgrade to websocket")
|
|
return
|
|
}
|
|
s := &Socket{
|
|
collection: c,
|
|
conn: socket,
|
|
send: make(chan []byte, 256),
|
|
}
|
|
c.register <- s
|
|
|
|
go s.readPump()
|
|
go s.writePump()
|
|
}
|
|
|
|
func (c *Collection) Add(s *Socket) {
|
|
c.register <- s
|
|
}
|
|
|
|
func (c *Collection) Remove(s *Socket) {
|
|
c.unregister <- s
|
|
}
|
|
|
|
func (c *Collection) Run() {
|
|
defer func() {
|
|
for s := range c.sockets {
|
|
close(s.send)
|
|
delete(c.sockets, s)
|
|
}
|
|
close(c.register)
|
|
close(c.unregister)
|
|
close(c.Broadcast)
|
|
close(c.close)
|
|
}()
|
|
for {
|
|
select {
|
|
case s := <-c.register:
|
|
c.sockets[s] = true
|
|
case s := <-c.unregister:
|
|
if _, ok := c.sockets[s]; ok {
|
|
delete(c.sockets, s)
|
|
close(s.send)
|
|
}
|
|
case m := <-c.Broadcast:
|
|
b, err := json.Marshal(m)
|
|
if err != nil {
|
|
log.WithError(err).Error("Failed to encode websocket message.")
|
|
continue
|
|
}
|
|
for s := range c.sockets {
|
|
select {
|
|
case s.send <- b:
|
|
default:
|
|
close(s.send)
|
|
delete(c.sockets, s)
|
|
}
|
|
}
|
|
case <-c.close:
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Collection) Close() {
|
|
c.close <- true
|
|
}
|