x0gp/server/internal/realtime/hub.go

// Package realtime manages WebSocket connections: registration, fan-out of
// race snapshots, and per-client send queues with backpressure handling.
package realtime

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"github.com/x0gp/server/internal/transport"
)

// Client is a single connected WebSocket client.
type Client struct {
	ID        string
	Send      chan []byte
	SessionID string
	Done      chan struct{}
	once      sync.Once
}

// Close signals the read/write pumps to exit and frees the send channel.
func (c *Client) Close() {
	c.once.Do(func() { close(c.Done) })
}

// Hub multiplexes snapshot broadcast to many clients.
type Hub struct {
	mu         sync.RWMutex
	clients    map[*Client]struct{}
	register   chan *Client
	unregister chan *Client
	broadcast  chan []byte
	stopCh     chan struct{}
	doneCh     chan struct{}

	// Stats.
	connTotal atomic.Uint64
	dropTotal atomic.Uint64
}

// NewHub creates a hub with sane defaults (send-buffer 64 frames).
func NewHub() *Hub {
	return &Hub{
		clients:    make(map[*Client]struct{}),
		register:   make(chan *Client, 16),
		unregister: make(chan *Client, 16),
		broadcast:  make(chan []byte, 64),
		stopCh:     make(chan struct{}),
		doneCh:     make(chan struct{}),
	}
}

// Stats returns current hub statistics.
type Stats struct {
	Connections uint64
	DropsTotal  uint64
}

func (h *Hub) Stats() Stats {
	h.mu.RLock()
	n := uint64(len(h.clients))
	h.mu.RUnlock()
	return Stats{Connections: n, DropsTotal: h.dropTotal.Load()}
}

// Register adds a client to the hub.
func (h *Hub) Register(c *Client) { h.register <- c }

// Unregister removes a client from the hub.
func (h *Hub) Unregister(c *Client) {
	select {
	case h.unregister <- c:
	default:
		// Channel full; force-close anyway.
		c.Close()
	}
}

// Run drives the hub until ctx is cancelled or Stop() is called.
func (h *Hub) Run(ctx context.Context) {
	defer close(h.doneCh)

	// Snapshot fan-out loop.
	go h.snapshotFanout(ctx)

	for {
		select {
		case <-ctx.Done():
			return
		case <-h.stopCh:
			return
		case c := <-h.register:
			h.mu.Lock()
			h.clients[c] = struct{}{}
			h.connTotal.Add(1)
			h.mu.Unlock()
		case c := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[c]; ok {
				delete(h.clients, c)
				close(c.Send)
			}
			h.mu.Unlock()
		}
	}
}

// Stop terminates the hub.
func (h *Hub) Stop() {
	close(h.stopCh)
	<-h.doneCh
}

// snapshotFanout consumes from the broadcast channel and sends to clients,
// dropping frames to slow consumers.
func (h *Hub) snapshotFanout(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-h.stopCh:
			return
		case frame, ok := <-h.broadcast:
			if !ok {
				return
			}
			h.fanout(frame)
		}
	}
}

// fanout delivers a frame to every client (non-blocking per-client).
func (h *Hub) fanout(frame []byte) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for c := range h.clients {
		select {
		case c.Send <- frame:
		default:
			// Slow consumer; drop and count.
			h.dropTotal.Add(1)
		}
	}
}

// Publish enqueues an envelope to be broadcast to all clients.
func (h *Hub) Publish(env *transport.Envelope) {
	data, err := transport.Encode(env)
	if err != nil {
		return
	}
	select {
	case h.broadcast <- data:
	default:
		h.dropTotal.Add(1)
	}
}

// helper to construct a server -> client envelope.
func MustEnvelope(t transport.MessageType, payload any) *transport.Envelope {
	return &transport.Envelope{
		Type:    t,
		TSMs:    time.Now().UnixMilli(),
		Payload: payload,
	}
}