67 lines
1.2 KiB
Go
67 lines
1.2 KiB
Go
package ws
|
|
|
|
import "sync"
|
|
|
|
// Client represents a single WebSocket connection.
|
|
// The handler owns the actual *websocket.Conn; the hub only needs the send channel.
|
|
type Client struct {
|
|
Send chan []byte
|
|
}
|
|
|
|
// Hub manages all active clients and routes broadcast messages to them.
|
|
type Hub struct {
|
|
mu sync.Mutex
|
|
clients map[*Client]struct{}
|
|
Register chan *Client
|
|
Unregister chan *Client
|
|
broadcast chan []byte
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
clients: make(map[*Client]struct{}),
|
|
Register: make(chan *Client, 8),
|
|
Unregister: make(chan *Client, 8),
|
|
broadcast: make(chan []byte, 512),
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Run() {
|
|
for {
|
|
select {
|
|
case c := <-h.Register:
|
|
h.mu.Lock()
|
|
h.clients[c] = struct{}{}
|
|
h.mu.Unlock()
|
|
|
|
case c := <-h.Unregister:
|
|
h.mu.Lock()
|
|
if _, ok := h.clients[c]; ok {
|
|
delete(h.clients, c)
|
|
close(c.Send)
|
|
}
|
|
h.mu.Unlock()
|
|
|
|
case msg := <-h.broadcast:
|
|
h.mu.Lock()
|
|
for c := range h.clients {
|
|
select {
|
|
case c.Send <- msg:
|
|
default:
|
|
delete(h.clients, c)
|
|
close(c.Send)
|
|
}
|
|
}
|
|
h.mu.Unlock()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Broadcast sends msg to every connected client (non-blocking).
|
|
func (h *Hub) Broadcast(msg []byte) {
|
|
select {
|
|
case h.broadcast <- msg:
|
|
default:
|
|
}
|
|
}
|