67 lines
1.3 KiB
Go
67 lines
1.3 KiB
Go
package ws
|
|
|
|
import "sync"
|
|
|
|
// Hub 管理房间及连接。它只负责单进程广播；跨实例广播可在此接入 Redis Pub/Sub。
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
rooms map[string]map[*Client]struct{}
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{rooms: make(map[string]map[*Client]struct{})}
|
|
}
|
|
|
|
func (h *Hub) Register(client *Client) int {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if h.rooms[client.room] == nil {
|
|
h.rooms[client.room] = make(map[*Client]struct{})
|
|
}
|
|
h.rooms[client.room][client] = struct{}{}
|
|
return len(h.rooms[client.room])
|
|
}
|
|
|
|
func (h *Hub) Unregister(client *Client) int {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
clients, ok := h.rooms[client.room]
|
|
if !ok {
|
|
return 0
|
|
}
|
|
if _, ok := clients[client]; ok {
|
|
delete(clients, client)
|
|
close(client.send)
|
|
}
|
|
remaining := len(clients)
|
|
if remaining == 0 {
|
|
delete(h.rooms, client.room)
|
|
}
|
|
return remaining
|
|
}
|
|
|
|
func (h *Hub) Broadcast(room string, message OutgoingMessage) {
|
|
h.mu.RLock()
|
|
clients := h.rooms[room]
|
|
stale := make([]*Client, 0)
|
|
for client := range clients {
|
|
select {
|
|
case client.send <- message:
|
|
default:
|
|
stale = append(stale, client)
|
|
}
|
|
}
|
|
h.mu.RUnlock()
|
|
|
|
for _, client := range stale {
|
|
h.Unregister(client)
|
|
_ = client.conn.Close()
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Count(room string) int {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
return len(h.rooms[room])
|
|
}
|