llink/go/internal/pusher/conn.go

package pusher

import (
	"context"
	"encoding/json"
	"log/slog"
	"sync"

	"nhooyr.io/websocket"
)

const sendBufferSize = 256

// Conn wraps a WebSocket connection with identity and a send buffer.
type Conn struct {
	id      string
	humanID string
	ws      *websocket.Conn
	send    chan []byte
	once    sync.Once // ensures close logic runs once
}

func newConn(id, humanID string, ws *websocket.Conn) *Conn {
	return &Conn{
		id:      id,
		humanID: humanID,
		ws:      ws,
		send:    make(chan []byte, sendBufferSize),
	}
}

// ReadPump reads messages from the WebSocket and forwards them to the hub.
// It blocks until the connection is closed or the context 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 {
				return
			}
			slog.Debug("websocket read error", "connId", c.id, "error", err)
			return
		}

		// Respond to keep-alive pings
		if string(data) == "ping" {
			if err := c.ws.Write(ctx, websocket.MessageText, []byte("pong")); err != nil {
				slog.Debug("websocket pong write error", "connId", c.id, "error", err)
				return
			}
			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)
		}
	}
}

// WritePump drains the send buffer and writes messages to the WebSocket.
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 {
				slog.Debug("websocket write error", "connId", c.id, "error", err)
				return
			}
		}
	}
}

// Send enqueues a ServerMessage to be written to the WebSocket.
// If the send buffer is full, the connection is closed (slow client).
func (c *Conn) Send(msg ServerMessage) {
	data, err := json.Marshal(msg)
	if err != nil {
		slog.Error("failed to marshal server message", "error", err)
		return
	}

	select {
	case c.send <- data:
	default:
		slog.Warn("slow client, closing connection", "connId", c.id, "humanId", c.humanID)
		c.Close()
	}
}

// Close closes the WebSocket connection and the send channel.
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})
}