Edu/services/push-gateway/internal/hub/hub.go

// Package hub manages the in-memory WebSocket connection pool for push-gateway.
//
// Hub is the single process-wide registry of all live WebSocket connections,
// indexed by userID -> connID -> *Connection. Each Connection owns a send
// channel consumed by a dedicated writer goroutine; Hub exposes Register /
// Unregister / SendToUser / Broadcast operations.
//
// The Hub also tracks per-user connection counts to enforce the
// MaxConnectionsPerUser limit (default 5) and provides CloseAll for graceful
// shutdown. Cross-instance presence and message fanout are delegated to a
// RedisClient (see internal/redis), keeping Hub free of any persistence layer.
package hub

import (
	"errors"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"github.com/google/uuid"
)

// ErrTooManyConnections is returned by Register when the per-user connection
// cap has been reached. Callers should translate this into a 429 /
// PUSH_TOO_MANY_CONNECTIONS response or a close frame with code 1008.
var ErrTooManyConnections = errors.New("too many connections per user")

// ErrHubClosing is returned by Register when the Hub is shutting down and no
// new connections are accepted.
var ErrHubClosing = errors.New("hub is closing")

// sendBufferSize is the per-connection send channel capacity. A full channel
// triggers message drop with metric increment (see Send).
const sendBufferSize = 64

// Connection wraps a single WebSocket connection together with its asynchronous
// send channel and lifecycle metadata.
type Connection struct {
	UserID    string
	ConnID    string
	conn      *websocket.Conn
	send      chan []byte
	closeOnce sync.Once
	closed    bool
	mu        sync.Mutex
	lastPing  time.Time
}

// Send delivers message to the connection's writer goroutine. When the send
// channel is full the message is dropped (non-blocking) and the caller is
// notified via the returned false value so it can increment the
// push_gateway_messages_dropped_total{reason="channel_full"} metric.
func (c *Connection) Send(message []byte) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return false
	}
	select {
	case c.send <- message:
		return true
	default:
		return false
	}
}

// Outgoing exposes the send channel for the writer goroutine. The channel is
// closed by Close, terminating the consumer loop.
func (c *Connection) Outgoing() <-chan []byte {
	return c.send
}

// Close shuts down the connection's send channel exactly once. Subsequent
// Send calls become no-ops.
func (c *Connection) Close() {
	c.closeOnce.Do(func() {
		c.mu.Lock()
		c.closed = true
		c.mu.Unlock()
		close(c.send)
	})
}

// LastPing returns the most recent heartbeat timestamp.
func (c *Connection) LastPing() time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.lastPing
}

// TouchPing updates the last heartbeat timestamp to now.
func (c *Connection) TouchPing() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.lastPing = time.Now()
}

// Conn returns the underlying gorilla websocket.Conn for control-frame
// operations (Ping/Pong/Close). Used by the reader goroutine.
func (c *Connection) Conn() *websocket.Conn {
	return c.conn
}

// Hub is the process-wide registry of live WebSocket connections.
//
// clients maps userID -> connID -> *Connection so a single user may hold
// multiple simultaneous connections (multi-device). counters caches the
// per-user connection count to make the MaxConnectionsPerUser check O(1).
type Hub struct {
	mu       sync.RWMutex
	clients  map[string]map[string]*Connection
	counters map[string]int
	closing  bool

	// maxConnsPerUser caps the number of simultaneous connections per user.
	maxConnsPerUser int

	// onRegister / onUnregister are optional hooks used by the Redis presence
	// layer to maintain the cross-instance online SET. They are invoked outside
	// the Hub lock to avoid blocking the registry on Redis round-trips.
	onRegister   func(userID, connID string)
	onUnregister func(userID, connID string)
}

// NewHub returns a Hub with the given per-user connection cap.
func NewHub(maxConnsPerUser int) *Hub {
	if maxConnsPerUser <= 0 {
		maxConnsPerUser = 5
	}
	return &Hub{
		clients:         make(map[string]map[string]*Connection),
		counters:        make(map[string]int),
		maxConnsPerUser: maxConnsPerUser,
	}
}

// SetPresenceHooks registers callbacks invoked after Register / Unregister.
// Used by the Redis presence layer to keep edu:push:online:<userID> in sync.
func (h *Hub) SetPresenceHooks(onRegister, onUnregister func(userID, connID string)) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.onRegister = onRegister
	h.onUnregister = onUnregister
}

// ActiveConnections returns the total number of live connections across all
// users. Exposed via /healthz and the push_gateway_active_connections gauge.
func (h *Hub) ActiveConnections() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	count := 0
	for _, conns := range h.clients {
		count += len(conns)
	}
	return count
}

// UserCount returns the number of distinct online users (used by /healthz).
func (h *Hub) UserCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients)
}

// HasUser reports whether at least one local connection exists for userID.
// Used by the /internal/push path to decide between local delivery and Redis
// Pub/Sub fanout.
func (h *Hub) HasUser(userID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	conns, ok := h.clients[userID]
	return ok && len(conns) > 0
}

// IsClosing reports whether CloseAll has been invoked. The /readyz probe uses
// this to return 503 during graceful shutdown so the load balancer stops
// sending new WebSocket upgrades while existing connections drain.
func (h *Hub) IsClosing() bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.closing
}

// Register adds a new connection for userID. Returns ErrTooManyConnections
// when the per-user cap is exceeded and ErrHubClosing when the Hub is shutting
// down. The returned Connection must be Unregistered by the caller on close.
func (h *Hub) Register(userID string, conn *websocket.Conn) (*Connection, error) {
	h.mu.Lock()
	if h.closing {
		h.mu.Unlock()
		return nil, ErrHubClosing
	}
	if h.counters[userID] >= h.maxConnsPerUser {
		h.mu.Unlock()
		return nil, ErrTooManyConnections
	}
	connID := uuid.NewString()
	c := &Connection{
		UserID:   userID,
		ConnID:   connID,
		conn:     conn,
		send:     make(chan []byte, sendBufferSize),
		lastPing: time.Now(),
	}
	if h.clients[userID] == nil {
		h.clients[userID] = make(map[string]*Connection)
	}
	h.clients[userID][connID] = c
	h.counters[userID]++
	onRegister := h.onRegister
	h.mu.Unlock()

	if onRegister != nil {
		onRegister(userID, connID)
	}
	return c, nil
}

// Unregister removes a connection and closes its send channel. Safe to call
// multiple times for the same Connection.
func (h *Hub) Unregister(c *Connection) {
	h.mu.Lock()
	conns, ok := h.clients[c.UserID]
	if !ok {
		h.mu.Unlock()
		return
	}
	if _, exists := conns[c.ConnID]; exists {
		delete(conns, c.ConnID)
		h.counters[c.UserID]--
		if h.counters[c.UserID] <= 0 {
			delete(h.counters, c.UserID)
		}
		if len(conns) == 0 {
			delete(h.clients, c.UserID)
		}
	}
	onUnregister := h.onUnregister
	h.mu.Unlock()

	c.Close()
	if onUnregister != nil {
		onUnregister(c.UserID, c.ConnID)
	}
}

// SendToUser delivers message to every local connection of userID. Returns
// the number of connections the message was successfully queued to. A zero
// return value signals "user not online locally" so the caller can fall back
// to Redis Pub/Sub fanout.
func (h *Hub) SendToUser(userID string, message []byte) int {
	h.mu.RLock()
	conns := h.clients[userID]
	// Snapshot the slice under the read lock to avoid holding it during Send.
	snapshot := make([]*Connection, 0, len(conns))
	for _, c := range conns {
		snapshot = append(snapshot, c)
	}
	h.mu.RUnlock()

	delivered := 0
	for _, c := range snapshot {
		if c.Send(message) {
			delivered++
		}
	}
	return delivered
}

// Broadcast delivers message to every live connection in this instance.
// Cross-instance broadcast is handled separately via Redis Pub/Sub.
func (h *Hub) Broadcast(message []byte) int {
	h.mu.RLock()
	// Count inline: calling ActiveConnections() here would re-enter RLock
	// while we already hold it, which can deadlock if a writer is waiting.
	total := 0
	for _, conns := range h.clients {
		total += len(conns)
	}
	snapshot := make([]*Connection, 0, total)
	for _, conns := range h.clients {
		for _, c := range conns {
			snapshot = append(snapshot, c)
		}
	}
	h.mu.RUnlock()

	delivered := 0
	for _, c := range snapshot {
		if c.Send(message) {
			delivered++
		}
	}
	return delivered
}

// CloseAll marks the Hub as closing (rejecting further Register calls) and
// sends a close frame to every live connection. Used during graceful shutdown.
// The caller should wait for connections to drain after CloseAll returns.
func (h *Hub) CloseAll() {
	h.mu.Lock()
	h.closing = true
	// Count inline instead of calling ActiveConnections() to avoid reentrant
	// RLock while holding the write lock (Go's sync.RWMutex is non-reentrant).
	total := 0
	for _, conns := range h.clients {
		total += len(conns)
	}
	all := make([]*Connection, 0, total)
	for _, conns := range h.clients {
		for _, c := range conns {
			all = append(all, c)
		}
	}
	h.mu.Unlock()

	for _, c := range all {
		// Send a graceful close frame (1001 going away) and close the send
		// channel so the writer goroutine exits.
		_ = c.conn.WriteControl(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseGoingAway, "server shutting down"),
			time.Now().Add(2*time.Second),
		)
		c.Close()
	}
}

// ForEachUser iterates over all locally online userIDs. Used by the Redis
// presence layer to rebuild the online SET on startup (ISSUE-058).
func (h *Hub) ForEachUser(fn func(userID string)) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for userID := range h.clients {
		fn(userID)
	}
}