64 lines
1.1 KiB
Go
64 lines
1.1 KiB
Go
|
package websocket
|
||
|
|
|
||
|
|
import "sync"
|
||
|
|
|
||
|
|
type Hub struct {
|
||
|
|
Clients map[string]*Client
|
||
|
|
Register chan *Client
|
||
|
|
Unregister chan *Client
|
||
|
|
mu sync.RWMutex
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewHub() *Hub {
|
||
|
|
return &Hub{
|
||
|
|
Clients: make(map[string]*Client),
|
||
|
|
Register: make(chan *Client),
|
||
|
|
Unregister: make(chan *Client),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *Hub) Run() {
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case client := <-h.Register:
|
||
|
|
h.mu.Lock()
|
||
|
|
h.Clients[client.ID] = client
|
||
|
|
h.mu.Unlock()
|
||
|
|
case client := <-h.Unregister:
|
||
|
|
h.mu.Lock()
|
||
|
|
if _, ok := h.Clients[client.ID]; ok {
|
||
|
|
delete(h.Clients, client.ID)
|
||
|
|
close(client.Send)
|
||
|
|
}
|
||
|
|
h.mu.Unlock()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *Hub) SendToClient(id string, msg []byte) {
|
||
|
|
h.mu.RLock()
|
||
|
|
defer h.mu.RUnlock()
|
||
|
|
if c, ok := h.Clients[id]; ok {
|
||
|
|
c.Send <- msg
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *Hub) SendToAll(ids []string, msg []byte) {
|
||
|
|
h.mu.RLock()
|
||
|
|
defer h.mu.RUnlock()
|
||
|
|
|
||
|
|
for _, id := range ids {
|
||
|
|
if c, ok := h.Clients[id]; ok {
|
||
|
|
c.Send <- msg
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (h *Hub) Broadcast(msg []byte) {
|
||
|
|
h.mu.RLock()
|
||
|
|
defer h.mu.RUnlock()
|
||
|
|
for _, client := range h.Clients {
|
||
|
|
client.Send <- msg
|
||
|
|
}
|
||
|
|
}
|