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) } } }