llink/go/internal/pusher/hub.go

package pusher

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

type subscribeRequest struct {
	conn      *Conn
	channelID string
}

type unsubscribeRequest struct {
	conn      *Conn
	channelID string
}

type broadcastRequest struct {
	conn      *Conn
	channelID string
	payload   json.RawMessage
}

type remoteEvent struct {
	channelID string
	event     redisEvent
}

// Hub manages all local WebSocket connections and channels on this pod.
// All state mutations happen in a single goroutine via Go channels — no locks.
type Hub struct {
	channels     map[string]*Channel
	connChannels map[*Conn]map[string]bool // reverse index: conn → set of channel IDs

	bridge     *RedisBridge
	authorizer *Authorizer

	subscribeCh   chan *subscribeRequest
	unsubscribeCh chan *unsubscribeRequest
	broadcastCh   chan *broadcastRequest
	disconnectCh  chan *Conn
	remoteEventCh chan *remoteEvent
}

func NewHub(bridge *RedisBridge, authorizer *Authorizer) *Hub {
	return &Hub{
		channels:      make(map[string]*Channel),
		connChannels:  make(map[*Conn]map[string]bool),
		bridge:        bridge,
		authorizer:    authorizer,
		subscribeCh:   make(chan *subscribeRequest, 256),
		unsubscribeCh: make(chan *unsubscribeRequest, 256),
		broadcastCh:   make(chan *broadcastRequest, 256),
		disconnectCh:  make(chan *Conn, 256),
		remoteEventCh: make(chan *remoteEvent, 256),
	}
}

// Run starts the hub event loop. Blocks until the context is cancelled.
func (h *Hub) Run(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return

		case req := <-h.subscribeCh:
			h.handleSubscribe(ctx, req)

		case req := <-h.unsubscribeCh:
			h.handleUnsubscribe(ctx, req)

		case req := <-h.broadcastCh:
			h.handleBroadcast(ctx, req)

		case conn := <-h.disconnectCh:
			h.handleDisconnect(ctx, conn)

		case evt := <-h.remoteEventCh:
			h.handleRemoteEvent(evt)
		}
	}
}

func (h *Hub) handleSubscribe(ctx context.Context, req *subscribeRequest) {
	if err := h.authorizer.Authorize(ctx, req.channelID, req.conn.humanID); err != nil {
		req.conn.Send(ServerMessage{
			Type:    TypeError,
			Channel: req.channelID,
			Message: "unauthorized",
		})
		return
	}

	ch, ok := h.channels[req.channelID]
	if !ok {
		ch = newChannel(req.channelID)
		h.channels[req.channelID] = ch
	}

	// Capture before addMember so multi-tab joins don't emit a spurious join.
	wasPresentLocally := ch.hasHumanID(req.conn.humanID)

	ch.addMember(req.conn, req.conn.humanID)

	if h.connChannels[req.conn] == nil {
		h.connChannels[req.conn] = make(map[string]bool)
	}
	h.connChannels[req.conn][req.channelID] = true

	presence, err := h.bridge.Subscribe(ctx, req.channelID, req.conn.id, req.conn.humanID)
	if err != nil {
		slog.Error("redis subscribe failed", "channelId", req.channelID, "error", err)
		// Fall back to local-only presence.
		presence = ch.localHumanIDs()
	}

	req.conn.Send(ServerMessage{
		Type:     TypeSubscribed,
		Channel:  req.channelID,
		Presence: presence,
	})

	// Redis self-filter drops our own echo, so same-pod peers need a direct nudge.
	if !wasPresentLocally {
		ch.broadcast(ServerMessage{
			Type:    TypeJoin,
			Channel: req.channelID,
			HumanID: req.conn.humanID,
		}, req.conn)
	}
}

func (h *Hub) handleUnsubscribe(ctx context.Context, req *unsubscribeRequest) {
	ch, ok := h.channels[req.channelID]
	if !ok {
		return
	}

	ch.removeMember(req.conn)

	if chans, ok := h.connChannels[req.conn]; ok {
		delete(chans, req.channelID)
	}

	if err := h.bridge.Unsubscribe(ctx, req.channelID, req.conn.id, req.conn.humanID); err != nil {
		slog.Error("redis unsubscribe failed", "channelId", req.channelID, "error", err)
	}

	// Only emit leave once the humanID has no remaining tabs on this pod.
	if !ch.hasHumanID(req.conn.humanID) {
		ch.broadcast(ServerMessage{
			Type:    TypeLeave,
			Channel: req.channelID,
			HumanID: req.conn.humanID,
		}, req.conn)
	}

	if ch.isEmpty() {
		delete(h.channels, req.channelID)
	}
}

func (h *Hub) handleBroadcast(ctx context.Context, req *broadcastRequest) {
	ch, ok := h.channels[req.channelID]
	if !ok {
		return
	}

	if _, isMember := ch.members[req.conn]; !isMember {
		req.conn.sendError("not subscribed to channel: " + req.channelID)
		return
	}

	// Local fanout (excluding sender), then publish for other pods.
	ch.broadcast(ServerMessage{
		Type:    TypeMessage,
		Channel: req.channelID,
		HumanID: req.conn.humanID,
		Payload: req.payload,
	}, req.conn)

	h.bridge.Broadcast(ctx, req.channelID, req.conn.humanID, req.payload)
}

func (h *Hub) handleDisconnect(ctx context.Context, conn *Conn) {
	chans, ok := h.connChannels[conn]
	if !ok {
		return
	}

	for channelID := range chans {
		ch, ok := h.channels[channelID]
		if !ok {
			continue
		}

		ch.removeMember(conn)

		if err := h.bridge.Unsubscribe(ctx, channelID, conn.id, conn.humanID); err != nil {
			slog.Error("redis unsubscribe on disconnect failed", "channelId", channelID, "error", err)
		}

		if !ch.hasHumanID(conn.humanID) {
			ch.broadcast(ServerMessage{
				Type:    TypeLeave,
				Channel: channelID,
				HumanID: conn.humanID,
			}, conn)
		}

		if ch.isEmpty() {
			delete(h.channels, channelID)
		}
	}

	delete(h.connChannels, conn)
}

func (h *Hub) handleRemoteEvent(evt *remoteEvent) {
	ch, ok := h.channels[evt.channelID]
	if !ok {
		// No local subscribers — drop the event.
		return
	}

	switch evt.event.Type {
	case TypeJoin:
		ch.broadcast(ServerMessage{
			Type:    TypeJoin,
			Channel: evt.channelID,
			HumanID: evt.event.HumanID,
		}, nil)

	case TypeLeave:
		ch.broadcast(ServerMessage{
			Type:    TypeLeave,
			Channel: evt.channelID,
			HumanID: evt.event.HumanID,
		}, nil)

	case TypeMessage:
		ch.broadcast(ServerMessage{
			Type:    TypeMessage,
			Channel: evt.channelID,
			HumanID: evt.event.HumanID,
			Payload: evt.event.Payload,
		}, nil)
	}
}

func (h *Hub) Subscribe(conn *Conn, channelID string) {
	h.subscribeCh <- &subscribeRequest{conn: conn, channelID: channelID}
}

func (h *Hub) disconnect(conn *Conn) {
	h.disconnectCh <- conn
}