Arjun Patel
129 lines
2.8 KiB
Go
129 lines
2.8 KiB
Go
package pusher
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"sync"
|
|
|
|
"github.com/flowy-live/llink/internal/utils/flog"
|
|
|
|
"nhooyr.io/websocket"
|
|
)
|
|
|
|
const sendBufferSize = 256
|
|
|
|
type Conn struct {
|
|
id string
|
|
humanID string
|
|
ws *websocket.Conn
|
|
send chan []byte
|
|
once sync.Once // guards Close
|
|
}
|
|
|
|
func newConn(id, humanID string, ws *websocket.Conn) *Conn {
|
|
return &Conn{
|
|
id: id,
|
|
humanID: humanID,
|
|
ws: ws,
|
|
send: make(chan []byte, sendBufferSize),
|
|
}
|
|
}
|
|
|
|
// ReadPump forwards inbound frames to the hub; blocks until the connection
|
|
// closes or ctx is cancelled.
|
|
func (c *Conn) ReadPump(ctx context.Context, hub *Hub) {
|
|
defer hub.disconnect(c)
|
|
|
|
for {
|
|
_, data, err := c.ws.Read(ctx)
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
flog.Info("websocket context cancelled", "connId", c.id, "humanId", c.humanID, "error", ctx.Err())
|
|
return
|
|
}
|
|
flog.Warn("websocket read error", "connId", c.id, "humanId", c.humanID, "error", err)
|
|
return
|
|
}
|
|
|
|
if string(data) == "ping" {
|
|
if err := c.ws.Write(ctx, websocket.MessageText, []byte("pong")); err != nil {
|
|
flog.Warn("websocket pong write error", "connId", c.id, "error", err)
|
|
}
|
|
continue
|
|
}
|
|
|
|
var msg ClientMessage
|
|
if err := json.Unmarshal(data, &msg); err != nil {
|
|
c.sendError("invalid message format")
|
|
continue
|
|
}
|
|
|
|
switch msg.Type {
|
|
case TypeSubscribe:
|
|
if msg.Channel == "" {
|
|
c.sendError("channel is required")
|
|
continue
|
|
}
|
|
hub.subscribeCh <- &subscribeRequest{conn: c, channelID: msg.Channel}
|
|
case TypeUnsubscribe:
|
|
if msg.Channel == "" {
|
|
c.sendError("channel is required")
|
|
continue
|
|
}
|
|
hub.unsubscribeCh <- &unsubscribeRequest{conn: c, channelID: msg.Channel}
|
|
case TypeMessage:
|
|
if msg.Channel == "" {
|
|
c.sendError("channel is required")
|
|
continue
|
|
}
|
|
hub.broadcastCh <- &broadcastRequest{conn: c, channelID: msg.Channel, payload: msg.Payload}
|
|
default:
|
|
c.sendError("unknown message type: " + msg.Type)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Conn) WritePump(ctx context.Context) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case data, ok := <-c.send:
|
|
if !ok {
|
|
return
|
|
}
|
|
if err := c.ws.Write(ctx, websocket.MessageText, data); err != nil {
|
|
flog.Debug("websocket write error", "connId", c.id, "error", err)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// A full send buffer closes the connection (slow client policy).
|
|
func (c *Conn) Send(msg ServerMessage) {
|
|
data, err := json.Marshal(msg)
|
|
if err != nil {
|
|
flog.Error("failed to marshal server message", "error", err)
|
|
return
|
|
}
|
|
|
|
select {
|
|
case c.send <- data:
|
|
default:
|
|
flog.Warn("slow client, closing connection", "connId", c.id, "humanId", c.humanID)
|
|
c.Close()
|
|
}
|
|
}
|
|
|
|
func (c *Conn) Close() {
|
|
c.once.Do(func() {
|
|
c.ws.Close(websocket.StatusNormalClosure, "closing")
|
|
close(c.send)
|
|
})
|
|
}
|
|
|
|
func (c *Conn) sendError(msg string) {
|
|
c.Send(ServerMessage{Type: TypeError, Message: msg})
|
|
}
|