63 lines
1.1 KiB
Go
63 lines
1.1 KiB
Go
package svc
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type Client struct {
|
|
SessionID string
|
|
Conn *websocket.Conn
|
|
sendMu sync.Mutex
|
|
}
|
|
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
clients map[string]*Client
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{clients: make(map[string]*Client)}
|
|
}
|
|
|
|
func (h *Hub) Add(c *Client) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
if old, ok := h.clients[c.SessionID]; ok && old.Conn != c.Conn {
|
|
_ = old.Conn.Close()
|
|
}
|
|
h.clients[c.SessionID] = c
|
|
}
|
|
|
|
func (h *Hub) Remove(sessionID string, conn *websocket.Conn) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
cur, ok := h.clients[sessionID]
|
|
if !ok {
|
|
return
|
|
}
|
|
if conn != nil && cur.Conn != conn {
|
|
return
|
|
}
|
|
delete(h.clients, sessionID)
|
|
}
|
|
|
|
func (h *Hub) Send(sessionID string, payload []byte) bool {
|
|
h.mu.RLock()
|
|
c, ok := h.clients[sessionID]
|
|
h.mu.RUnlock()
|
|
if !ok {
|
|
return false
|
|
}
|
|
|
|
c.sendMu.Lock()
|
|
defer c.sendMu.Unlock()
|
|
if err := c.Conn.WriteMessage(websocket.TextMessage, payload); err != nil {
|
|
_ = c.Conn.Close()
|
|
h.Remove(sessionID, c.Conn)
|
|
return false
|
|
}
|
|
return true
|
|
}
|