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}) }