Messenger/server/internal/realtime/hub.go

package realtime

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

	"github.com/coder/websocket"
)

// MembershipProvider resolves conversation membership and read state.
type MembershipProvider interface {
	ConversationIDsForUser(ctx context.Context, userID int64) ([]int64, error)
	IsMember(ctx context.Context, conversationID, userID int64) (bool, error)
	UpdateLastRead(ctx context.Context, conversationID, userID int64) error
}

// Hub tracks WebSocket connections and fans out events to the
// users subscribed to a conversation.
type Hub struct {
	mu      sync.Mutex
	conns   map[int64]map[*client]struct{}
	subs    map[int64]map[int64]struct{}
	members MembershipProvider
}

// NewHub creates a hub that resolves membership through members.
func NewHub(members MembershipProvider) *Hub {
	return &Hub{
		conns:   make(map[int64]map[*client]struct{}),
		subs:    make(map[int64]map[int64]struct{}),
		members: members,
	}
}

// HandleWS upgrades the HTTP request to a WebSocket connection and
// runs the client loops until the connection closes.
func (h *Hub) HandleWS(w http.ResponseWriter, r *http.Request, userID int64) {
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}})
	if err != nil {
		slog.Error("ws upgrade failed", "error", err)
		return
	}
	conn.SetReadLimit(32 << 10)

	ctx := r.Context()
	conversationIDs, err := h.members.ConversationIDsForUser(ctx, userID)
	if err != nil {
		slog.Error("ws: load memberships", "user_id", userID, "error", err)
		conn.Close(websocket.StatusInternalError, "membership lookup failed")
		return
	}

	c := newClient(h, conn, userID)
	h.register(c, conversationIDs)
	defer h.unregister(c, conversationIDs)

	h.broadcastPresence(userID, conversationIDs, StatusOnline)
	defer h.broadcastPresence(userID, conversationIDs, StatusOffline)

	go c.writeLoop()
	if err := c.readLoop(); err != nil {
		slog.Debug("ws connection closed", "user_id", userID, "error", err)
	}
}

// PublishMessageCreated delivers a stored message to everyone subscribed
// to the conversation, including the sender.
func (h *Hub) PublishMessageCreated(conversationID int64, evt MessageCreated) {
	h.sendToConversation(conversationID, envelope{Type: EventMessageCreated, Data: mustMarshal(evt)}, 0)
}

// PublishTyping delivers a typing state change to everyone except the sender.
func (h *Hub) PublishTyping(conversationID, userID int64, isTyping bool) {
	evt := Typing{ConversationID: conversationID, UserID: userID, IsTyping: isTyping}
	h.sendToConversation(conversationID, envelope{Type: EventTyping, Data: mustMarshal(evt)}, userID)
}

// PublishRead delivers a read receipt to everyone except the reader.
func (h *Hub) PublishRead(conversationID, userID, messageID int64) {
	evt := Read{ConversationID: conversationID, UserID: userID, MessageID: messageID}
	h.sendToConversation(conversationID, envelope{Type: EventRead, Data: mustMarshal(evt)}, userID)
}

// Subscribe adds a user to a conversation's subscriber set so they receive its
// events without reconnecting (used when a member is added to a group).
func (h *Hub) Subscribe(userID, conversationID int64) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.subs[conversationID] == nil {
		h.subs[conversationID] = make(map[int64]struct{})
	}
	h.subs[conversationID][userID] = struct{}{}
}

// Unsubscribe removes a user from a conversation's subscriber set.
func (h *Hub) Unsubscribe(userID, conversationID int64) {
	h.mu.Lock()
	defer h.mu.Unlock()
	users := h.subs[conversationID]
	delete(users, userID)
	if len(users) == 0 {
		delete(h.subs, conversationID)
	}
}

// SendToUser delivers an event to every connection of a single user.
func (h *Hub) SendToUser(userID int64, evtType string, data any) {
	payload, err := json.Marshal(envelope{Type: evtType, Data: mustMarshal(data)})
	if err != nil {
		slog.Error("ws: marshal direct event", "error", err)
		return
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	for c := range h.conns[userID] {
		c.enqueueRaw(payload)
	}
}

// PublishMemberChanged tells a conversation's members that someone joined or left.
func (h *Hub) PublishMemberChanged(conversationID, userID int64, added bool) {
	evtType := EventMemberAdded
	if !added {
		evtType = EventMemberRemoved
	}
	evt := MemberChanged{ConversationID: conversationID, UserID: userID}
	h.sendToConversation(conversationID, envelope{Type: evtType, Data: mustMarshal(evt)}, 0)
}

// PublishReactionChanged tells a conversation's members that a reaction changed.
func (h *Hub) PublishReactionChanged(conversationID, messageID, userID int64, emoji string, added bool) {
	evtType := EventReactionAdded
	if !added {
		evtType = EventReactionRemoved
	}
	evt := ReactionChanged{
		ConversationID: conversationID,
		MessageID:      messageID,
		UserID:         userID,
		Emoji:          emoji,
	}
	h.sendToConversation(conversationID, envelope{Type: evtType, Data: mustMarshal(evt)}, 0)
}

func (h *Hub) register(c *client, conversationIDs []int64) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if h.conns[c.userID] == nil {
		h.conns[c.userID] = make(map[*client]struct{})
	}
	h.conns[c.userID][c] = struct{}{}
	for _, cid := range conversationIDs {
		if h.subs[cid] == nil {
			h.subs[cid] = make(map[int64]struct{})
		}
		h.subs[cid][c.userID] = struct{}{}
	}
}

func (h *Hub) unregister(c *client, conversationIDs []int64) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if conns := h.conns[c.userID]; conns != nil {
		delete(conns, c)
		if len(conns) == 0 {
			delete(h.conns, c.userID)
		}
	}
	for _, cid := range conversationIDs {
		if users := h.subs[cid]; users != nil {
			delete(users, c.userID)
			if len(users) == 0 {
				delete(h.subs, cid)
			}
		}
	}
	close(c.done)
}

func (h *Hub) broadcastPresence(userID int64, conversationIDs []int64, status string) {
	evt := Presence{UserID: userID, Status: status}
	env := envelope{Type: EventPresence, Data: mustMarshal(evt)}
	h.mu.Lock()
	defer h.mu.Unlock()
	for _, cid := range conversationIDs {
		for uid := range h.subs[cid] {
			if uid == userID {
				continue
			}
			for c := range h.conns[uid] {
				c.enqueue(env)
			}
		}
	}
}

func (h *Hub) sendToConversation(conversationID int64, env envelope, exceptUserID int64) {
	payload, err := json.Marshal(env)
	if err != nil {
		slog.Error("ws: marshal event", "error", err)
		return
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	for uid := range h.subs[conversationID] {
		if uid == exceptUserID {
			continue
		}
		for c := range h.conns[uid] {
			c.enqueueRaw(payload)
		}
	}
}

func mustMarshal(v any) json.RawMessage {
	data, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	return data
}