qkeeper/hub.go

package main

import (
	"encoding/json"
	"net/http"
	"sync/atomic"

	"github.com/gorilla/websocket"
)

// ---- WebSocket hub ----

// Hub serialises all queue mutations through a single goroutine (run).
// This eliminates data races without requiring a queue-level mutex.
type Hub struct {
	clients     map[*Client]bool
	registerC   chan *Client
	unregisterC chan *Client
	broadcastC  chan []byte   // pre-serialised JSON pushed to all clients
	dispatchC   chan dispatch // inbound client messages
	timerC      chan string   // entry IDs for 60s-timer-triggered removal
	stopC       chan struct{}
	activeCount atomic.Int32 // len(q.Entries); read by Registry.List without locking
}

type dispatch struct {
	c       *Client // nil for timer-triggered removals
	payload []byte
}

func newHub() *Hub {
	return &Hub{
		clients:     make(map[*Client]bool),
		registerC:   make(chan *Client, 16),
		unregisterC: make(chan *Client, 16),
		broadcastC:  make(chan []byte, 256),
		dispatchC:   make(chan dispatch, 256),
		timerC:      make(chan string, 64),
		stopC:       make(chan struct{}),
	}
}

// shutdown signals the hub goroutine to exit.
func (h *Hub) shutdown() {
	close(h.stopC)
}

// run is the single goroutine that owns all hub and queue state.
func (h *Hub) run(q *QueueState, reg *Registry, queueID string) {
	for {
		select {
		case <-h.stopC:
			return

		case c := <-h.registerC:
			h.clients[c] = true

		case c := <-h.unregisterC:
			if _, ok := h.clients[c]; ok {
				delete(h.clients, c)
				close(c.send)
			}

		case msg := <-h.broadcastC:
			for c := range h.clients {
				select {
				case c.send <- msg:
				default:
					// Slow client: drop and disconnect.
					delete(h.clients, c)
					close(c.send)
				}
			}

		case d := <-h.dispatchC:
			h.handleMessage(d.c, d.payload, q, reg, queueID)

		case entryID := <-h.timerC:
			if q.timerRemove(entryID) {
				h.pushState(q)
			}
		}
	}
}

// pushState broadcasts the full current queue state to all clients.
func (h *Hub) pushState(q *QueueState) {
	h.activeCount.Store(int32(len(q.Entries)))
	b := q.stateJSON()
	select {
	case h.broadcastC <- b:
	default:
	}
}

// sendTo sends a JSON-serialisable value to one client.
func (h *Hub) sendTo(c *Client, v interface{}) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	select {
	case c.send <- b:
	default:
	}
}

// ---- WebSocket client ----

type Client struct {
	hub   *Hub
	conn  *websocket.Conn
	send  chan []byte
	token string // set after init/join; used for all subsequent actions
}

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func (h *Hub) wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	c := &Client{
		hub:  h,
		conn: conn,
		send: make(chan []byte, 64),
	}
	h.registerC <- c
	go c.writePump()
	c.readPump()
}

func (c *Client) readPump() {
	defer func() {
		c.hub.unregisterC <- c
		c.conn.Close()
	}()
	c.conn.SetReadLimit(4096)
	for {
		_, raw, err := c.conn.ReadMessage()
		if err != nil {
			return
		}
		c.hub.dispatchC <- dispatch{c: c, payload: raw}
	}
}

func (c *Client) writePump() {
	defer c.conn.Close()
	for msg := range c.send {
		if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
			return
		}
	}
}

// ---- Message handling ----

type inMsg struct {
	Type    string `json:"type"`
	Token   string `json:"token,omitempty"`
	Name    string `json:"name,omitempty"`
	EntryID string `json:"entryId,omitempty"`
}

func (h *Hub) handleMessage(c *Client, raw []byte, q *QueueState, reg *Registry, queueID string) {
	var msg inMsg
	if err := json.Unmarshal(raw, &msg); err != nil {
		return
	}

	switch msg.Type {

	case "init":
		// Client sends their stored token (empty string if none).
		c.token = msg.Token
		var entryID interface{} // nil → JSON null
		if msg.Token != "" {
			if e := q.findByToken(msg.Token); e != nil {
				entryID = e.ID
			} else {
				c.token = "" // stale token
			}
		}
		h.sendTo(c, map[string]interface{}{"type": "hello", "yourEntryId": entryID})
		// Send current state so client renders queue immediately.
		h.sendTo(c, json.RawMessage(q.stateJSON()))

	case "join":
		// Anti-spam: reject if this client already has an active entry.
		if c.token != "" {
			if e := q.findByToken(c.token); e != nil {
				h.sendTo(c, map[string]interface{}{"type": "hello", "yourEntryId": e.ID})
				return
			}
		}
		e, token := q.join(msg.Name)
		c.token = token
		h.sendTo(c, map[string]interface{}{
			"type":        "joined",
			"yourToken":   token,
			"yourEntryId": e.ID,
		})
		h.pushState(q)
		reg.Touch(queueID)

	case "done":
		if c.token == "" {
			return
		}
		if q.selfRemove(c.token) {
			c.token = ""
			h.pushState(q)
		}

	case "mark_done":
		if c.token == "" || msg.EntryID == "" {
			return
		}
		actor := q.findByToken(c.token)
		if actor == nil {
			return
		}
		if q.markDone(actor.Number, msg.EntryID, h.timerC) {
			h.pushState(q)
		}

	case "vote_remove":
		if c.token == "" || msg.EntryID == "" {
			return
		}
		actor := q.findByToken(c.token)
		if actor == nil {
			return
		}
		switch q.voteRemove(actor.Number, msg.EntryID, h.timerC) {
		case "ok":
			h.pushState(q)
		case "already_voting":
			onNum := q.findVoteTarget(actor.Number)
			h.sendTo(c, map[string]interface{}{
				"type": "error", "reason": "already_voting", "onNumber": onNum,
			})
		}

	case "restore":
		if c.token == "" {
			return
		}
		if q.restore(c.token) {
			h.pushState(q)
		}
	}
}