49 lines
715 B
Go
49 lines
715 B
Go
package ws
|
|
|
|
import "sync"
|
|
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
clients map[*Client]struct{}
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
clients: make(map[*Client]struct{}),
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Register(c *Client) {
|
|
h.mu.Lock()
|
|
h.clients[c] = struct{}{}
|
|
h.mu.Unlock()
|
|
}
|
|
|
|
func (h *Hub) Unregister(c *Client) {
|
|
h.mu.Lock()
|
|
if _, ok := h.clients[c]; ok {
|
|
delete(h.clients, c)
|
|
close(c.send)
|
|
}
|
|
h.mu.Unlock()
|
|
}
|
|
|
|
func (h *Hub) Broadcast(msg []byte) {
|
|
if msg == nil {
|
|
return
|
|
}
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
for c := range h.clients {
|
|
select {
|
|
case c.send <- msg:
|
|
default:
|
|
// Client too slow, schedule disconnect
|
|
go func(c *Client) {
|
|
h.Unregister(c)
|
|
c.conn.Close()
|
|
}(c)
|
|
}
|
|
}
|
|
}
|