Gorilla WebSocket Hub Room Broadcast

package hub

import (
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type Client struct{ conn *websocket.Conn; room string; send chan []byte }

type Hub struct {
	mu         sync.RWMutex
	rooms      map[string]map[*Client]struct{}
	register   chan *Client
	unregister chan *Client
	broadcast  chan message
}

type message struct{ room string; data []byte }

func newHub() *Hub {
	return &Hub{
		rooms:      map[string]map[*Client]struct{}{},
		register:   make(chan *Client),
		unregister: make(chan *Client),
		broadcast:  make(chan message, 64),
	}
}

func (h *Hub) run() {
	for {
		select {
		case c := <-h.register:
			h.mu.Lock()
			if h.rooms[c.room] == nil {
				h.rooms[c.room] = map[*Client]struct{}{}
			}
			h.rooms[c.room][c] = struct{}{}
			h.mu.Unlock()
		case c := <-h.unregister:
			h.mu.Lock()
			if room, ok := h.rooms[c.room]; ok {
				if _, alive := room[c]; alive {
					delete(room, c)
					close(c.send)
				}
				if len(room) == 0 {
					delete(h.rooms, c.room)
				}
			}
			h.mu.Unlock()
		case msg := <-h.broadcast:
			h.mu.RLock()
			for c := range h.rooms[msg.room] {
				select {
				case c.send <- msg.data:
				default: // 缓冲满：掉队者被移除
					go h.unregister <- c
				}
			}
			h.mu.RUnlock()
		}
	}
}

var upgrader = websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}

func (h *Hub) serveWS(w http.ResponseWriter, r *http.Request, room string) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("upgrade:", err)
		return
	}
	client := &Client{conn: conn, room: room, send: make(chan []byte, 16)}
	h.register <- client

	go func() {
		defer func() { h.unregister <- client }()
		for {
			conn.SetReadDeadline(time.Now().Add(90 * time.Second))
			if _, _, err := conn.ReadMessage(); err != nil {
				return
			}
		}
	}()

	go func() {
		for data := range client.send {
			if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
				return
			}
		}
	}()
}