133 lines
2.8 KiB
Go
133 lines
2.8 KiB
Go
package ws
|
|
|
|
import (
|
|
"log"
|
|
"sync"
|
|
)
|
|
|
|
// Hub WebSocket连接池
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
clients map[*Client]bool
|
|
broadcast chan []byte
|
|
register chan *Client
|
|
unregister chan *Client
|
|
|
|
// 按用户ID索引的客户端映射
|
|
userClients map[string]map[*Client]bool
|
|
}
|
|
|
|
// NewHub 创建WebSocket Hub
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
clients: make(map[*Client]bool),
|
|
broadcast: make(chan []byte, 256),
|
|
register: make(chan *Client),
|
|
unregister: make(chan *Client),
|
|
userClients: make(map[string]map[*Client]bool),
|
|
}
|
|
}
|
|
|
|
// Run 启动Hub主循环
|
|
func (h *Hub) Run() {
|
|
for {
|
|
select {
|
|
case client := <-h.register:
|
|
h.mu.Lock()
|
|
h.clients[client] = true
|
|
|
|
// 用户索引
|
|
if h.userClients[client.UserID] == nil {
|
|
h.userClients[client.UserID] = make(map[*Client]bool)
|
|
}
|
|
h.userClients[client.UserID][client] = true
|
|
h.mu.Unlock()
|
|
|
|
log.Printf("[WS] 客户端连接: user=%s session=%s (当前连接数: %d)",
|
|
client.UserID, client.SessionID, len(h.clients))
|
|
|
|
case client := <-h.unregister:
|
|
h.mu.Lock()
|
|
if _, ok := h.clients[client]; ok {
|
|
delete(h.clients, client)
|
|
close(client.Send)
|
|
|
|
// 清理用户索引
|
|
if h.userClients[client.UserID] != nil {
|
|
delete(h.userClients[client.UserID], client)
|
|
if len(h.userClients[client.UserID]) == 0 {
|
|
delete(h.userClients, client.UserID)
|
|
}
|
|
}
|
|
}
|
|
h.mu.Unlock()
|
|
|
|
log.Printf("[WS] 客户端断开: user=%s session=%s (当前连接数: %d)",
|
|
client.UserID, client.SessionID, len(h.clients))
|
|
|
|
case message := <-h.broadcast:
|
|
h.mu.RLock()
|
|
for client := range h.clients {
|
|
select {
|
|
case client.Send <- message:
|
|
default:
|
|
// 客户端发送通道已满，跳过
|
|
close(client.Send)
|
|
delete(h.clients, client)
|
|
}
|
|
}
|
|
h.mu.RUnlock()
|
|
}
|
|
}
|
|
}
|
|
|
|
// SendToUser 向指定用户的所有连接发送消息
|
|
func (h *Hub) SendToUser(userID string, message []byte) {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
|
|
if clients, ok := h.userClients[userID]; ok {
|
|
for client := range clients {
|
|
select {
|
|
case client.Send <- message:
|
|
default:
|
|
// 跳过阻塞的客户端
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// SendToSession 向指定会话的连接发送消息
|
|
func (h *Hub) SendToSession(userID, sessionID string, message []byte) {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
|
|
if clients, ok := h.userClients[userID]; ok {
|
|
for client := range clients {
|
|
if client.SessionID == sessionID {
|
|
select {
|
|
case client.Send <- message:
|
|
default:
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ClientCount 获取当前连接数
|
|
func (h *Hub) ClientCount() int {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
return len(h.clients)
|
|
}
|
|
|
|
// UserClientCount 获取指定用户的连接数
|
|
func (h *Hub) UserClientCount(userID string) int {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
if clients, ok := h.userClients[userID]; ok {
|
|
return len(clients)
|
|
}
|
|
return 0
|
|
}
|