neural-hunt/internal/ws/hub.go

package ws

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

	"github.com/gorilla/websocket"
)

type Event struct {
	Type   string `json:"type"`
	TaskID string `json:"task_id,omitempty"`
	Data   any    `json:"data,omitempty"`
}

type Client struct {
	Conn       *websocket.Conn
	TaskID     string
	ClientID   string
	All        bool
	Admin      bool
	send       chan []byte
	closed     chan struct{}
	onWrite    func(int)
	onDrop     func()
	closeOnce  sync.Once
	dropStreak atomic.Uint32
}

func NewClient(conn *websocket.Conn, taskID, clientID string, all bool) *Client {
	return &Client{Conn: conn, TaskID: taskID, ClientID: clientID, All: all, send: make(chan []byte, 32), closed: make(chan struct{})}
}

func NewAdminClient(conn *websocket.Conn) *Client {
	return &Client{Conn: conn, ClientID: "admin", All: true, Admin: true, send: make(chan []byte, 64), closed: make(chan struct{})}
}
func (c *Client) start() { go c.writeLoop() }
func (c *Client) writeLoop() {
	for {
		select {
		case <-c.closed:
			return
		case b := <-c.send:
			_ = c.Conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
			if err := c.Conn.WriteMessage(websocket.TextMessage, b); err != nil {
				c.Close()
				return
			}
			if c.onWrite != nil {
				c.onWrite(len(b))
			}
		}
	}
}
func (c *Client) Enqueue(v any) bool {
	b, err := json.Marshal(v)
	if err != nil {
		return false
	}
	return c.EnqueueBytes(b)
}
func (c *Client) EnqueueBytes(b []byte) bool {
	select {
	case <-c.closed:
		return false
	default:
	}
	select {
	case c.send <- b:
		c.dropStreak.Store(0)
		return true
	default:
		if c.onDrop != nil {
			c.onDrop()
		}
		// A persistently slow reader can otherwise hold memory/socket resources
		// forever while every broadcast is dropped. Close after a short burst of
		// consecutive queue overflows; healthy clients reset the streak on enqueue.
		if c.dropStreak.Add(1) >= 8 {
			c.Close()
		}
		return false
	}
}
func (c *Client) Close() { c.closeOnce.Do(func() { close(c.closed); _ = c.Conn.Close() }) }

// Ping sends a real WebSocket control ping. Browsers answer this automatically
// with a pong, which refreshes the server-side read deadline. Do not replace
// this with an application JSON message: that caused every otherwise-idle
// browser connection to be closed after the 90s read deadline.
func (c *Client) Ping() error {
	select {
	case <-c.closed:
		return context.Canceled
	default:
	}
	if err := c.Conn.WriteControl(websocket.PingMessage, []byte("nh"), time.Now().Add(5*time.Second)); err != nil {
		c.Close()
		return err
	}
	return nil
}

type pointUpdate struct {
	ClientID string
	Data     any
}
type rateBucket struct {
	unix                 int64
	frames, bytes, drops uint64
}

type Hub struct {
	mu        sync.RWMutex
	clients   map[*Client]struct{}
	pendingMu sync.Mutex
	pending   map[string]map[string]any
	frames    atomic.Uint64
	bytes     atomic.Uint64
	drops     atomic.Uint64
	rateMu    sync.Mutex
	rate      [60]rateBucket
}

func New() *Hub { return &Hub{clients: map[*Client]struct{}{}, pending: map[string]map[string]any{}} }
func (h *Hub) Add(c *Client) {
	c.onWrite = h.recordWrite
	c.onDrop = h.recordDrop
	// Same-session websocket reconnects are intentionally allowed by runtime.
	// Keep only the newest user socket for a ClientID so a half-dead TCP
	// connection cannot continue receiving frames beside its replacement.
	var stale []*Client
	h.mu.Lock()
	if !c.All && c.ClientID != "" {
		for old := range h.clients {
			if !old.All && old.ClientID == c.ClientID {
				delete(h.clients, old)
				stale = append(stale, old)
			}
		}
	}
	h.clients[c] = struct{}{}
	h.mu.Unlock()
	for _, old := range stale {
		old.Close()
	}
	c.start()
}
func (h *Hub) Remove(c *Client) { h.mu.Lock(); delete(h.clients, c); h.mu.Unlock(); c.Close() }
func (h *Hub) Connected() int   { h.mu.RLock(); defer h.mu.RUnlock(); return len(h.clients) }

func (h *Hub) Publish(ctx context.Context, e Event) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	b, err := json.Marshal(e)
	if err != nil {
		return err
	}
	h.broadcastBytes(e.TaskID, b)
	return nil
}

// PublishAdmin sends an ephemeral event only to authenticated admin websocket
// clients. It is intentionally separate from the public/task broadcast path so
// per-guess telemetry can never leak to normal players or the leaderboard.
func (h *Hub) PublishAdmin(ctx context.Context, e Event) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	h.mu.RLock()
	hasAdmin := false
	for c := range h.clients {
		if c.Admin {
			hasAdmin = true
			break
		}
	}
	h.mu.RUnlock()
	if !hasAdmin {
		return nil
	}
	b, err := json.Marshal(e)
	if err != nil {
		return err
	}
	h.broadcastAdminBytes(b)
	return nil
}

func (h *Hub) broadcastAdminBytes(b []byte) {
	h.mu.RLock()
	clients := make([]*Client, 0, 4)
	for c := range h.clients {
		if c.Admin {
			clients = append(clients, c)
		}
	}
	h.mu.RUnlock()
	for _, c := range clients {
		c.EnqueueBytes(b)
	}
}

// PublishPoint coalesces repeated improvements for the same client. Every 250ms
// all changed points for a task are emitted as one frame per websocket instead
// of one synchronous write per guess per client.
func (h *Hub) PublishPoint(taskID, clientID string, data any) {
	h.pendingMu.Lock()
	m := h.pending[taskID]
	if m == nil {
		m = map[string]any{}
		h.pending[taskID] = m
	}
	m[clientID] = data
	h.pendingMu.Unlock()
}
func (h *Hub) flushPoints() {
	h.pendingMu.Lock()
	pending := h.pending
	h.pending = map[string]map[string]any{}
	h.pendingMu.Unlock()
	for taskID, m := range pending {
		arr := make([]any, 0, len(m))
		for _, v := range m {
			arr = append(arr, v)
		}
		b, err := json.Marshal(Event{Type: "points", TaskID: taskID, Data: arr})
		if err == nil {
			h.broadcastBytes(taskID, b)
		}
	}
}
func (h *Hub) broadcastBytes(taskID string, b []byte) {
	h.mu.RLock()
	clients := make([]*Client, 0, len(h.clients))
	for c := range h.clients {
		if c.All || taskID == "" || c.TaskID == taskID {
			clients = append(clients, c)
		}
	}
	h.mu.RUnlock()
	for _, c := range clients {
		c.EnqueueBytes(b)
	}
}
func (h *Hub) recordWrite(n int) {
	h.frames.Add(1)
	h.bytes.Add(uint64(n))
	now := time.Now().Unix()
	h.rateMu.Lock()
	i := now % 60
	if h.rate[i].unix != now {
		h.rate[i] = rateBucket{unix: now}
	}
	h.rate[i].frames++
	h.rate[i].bytes += uint64(n)
	h.rateMu.Unlock()
}
func (h *Hub) recordDrop() {
	h.drops.Add(1)
	now := time.Now().Unix()
	h.rateMu.Lock()
	i := now % 60
	if h.rate[i].unix != now {
		h.rate[i] = rateBucket{unix: now}
	}
	h.rate[i].drops++
	h.rateMu.Unlock()
}

type Metrics struct {
	Connected     int     `json:"connected"`
	FramesPerSec  float64 `json:"frames_per_sec"`
	BytesPerSec   float64 `json:"bytes_per_sec"`
	DroppedPerSec float64 `json:"dropped_per_sec"`
	FramesTotal   uint64  `json:"frames_total"`
	BytesTotal    uint64  `json:"bytes_total"`
	DroppedTotal  uint64  `json:"dropped_total"`
}

func (h *Hub) Metrics() Metrics {
	now := time.Now().Unix()
	var f, b, d uint64
	h.rateMu.Lock()
	for _, x := range h.rate {
		if x.unix > now-5 {
			f += x.frames
			b += x.bytes
			d += x.drops
		}
	}
	h.rateMu.Unlock()
	return Metrics{Connected: h.Connected(), FramesPerSec: float64(f) / 5, BytesPerSec: float64(b) / 5, DroppedPerSec: float64(d) / 5, FramesTotal: h.frames.Load(), BytesTotal: h.bytes.Load(), DroppedTotal: h.drops.Load()}
}
func (h *Hub) Run(ctx context.Context) {
	t := time.NewTicker(250 * time.Millisecond)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			h.flushPoints()
		}
	}
}