90 lines
2.3 KiB
Go
90 lines
2.3 KiB
Go
|
// Package websocket 实现WebSocket连接管理和消息广播
|
|||
|
|
// 为在线协作功能预留，支持多客户端实时通信
|
|||
|
|
package websocket
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"log"
|
|||
|
|
"sync"
|
|||
|
|
|
|||
|
|
"github.com/gorilla/websocket"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// Hub WebSocket连接管理中心
|
|||
|
|
// 负责维护所有活跃的WebSocket连接，处理消息广播
|
|||
|
|
type Hub struct {
|
|||
|
|
clients map[*Client]bool // 所有已连接的客户端
|
|||
|
|
broadcast chan []byte // 待广播的消息通道
|
|||
|
|
register chan *Client // 客户端注册通道
|
|||
|
|
unregister chan *Client // 客户端注销通道
|
|||
|
|
mu sync.RWMutex // 读写锁，保护 clients map
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Client 单个WebSocket客户端连接
|
|||
|
|
type Client struct {
|
|||
|
|
hub *Hub // 所属的Hub
|
|||
|
|
conn *websocket.Conn // WebSocket连接实例
|
|||
|
|
send chan []byte // 待发送的消息队列
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NewHub 创建并返回一个新的Hub实例
|
|||
|
|
func NewHub() *Hub {
|
|||
|
|
return &Hub{
|
|||
|
|
clients: make(map[*Client]bool),
|
|||
|
|
broadcast: make(chan []byte, 256),
|
|||
|
|
register: make(chan *Client),
|
|||
|
|
unregister: make(chan *Client),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Run 启动Hub的主循环
|
|||
|
|
// 在独立的goroutine中运行，处理客户端注册、注销和消息广播
|
|||
|
|
func (h *Hub) Run() {
|
|||
|
|
for {
|
|||
|
|
select {
|
|||
|
|
case client := <-h.register:
|
|||
|
|
// 新客户端连接注册
|
|||
|
|
h.mu.Lock()
|
|||
|
|
h.clients[client] = true
|
|||
|
|
h.mu.Unlock()
|
|||
|
|
log.Printf("WebSocket客户端已连接，当前在线: %d", len(h.clients))
|
|||
|
|
|
|||
|
|
case client := <-h.unregister:
|
|||
|
|
// 客户端断开连接注销
|
|||
|
|
h.mu.Lock()
|
|||
|
|
if _, ok := h.clients[client]; ok {
|
|||
|
|
delete(h.clients, client)
|
|||
|
|
close(client.send)
|
|||
|
|
}
|
|||
|
|
h.mu.Unlock()
|
|||
|
|
log.Printf("WebSocket客户端已断开，当前在线: %d", 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()
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// BroadcastMessage 向所有连接的客户端广播消息
|
|||
|
|
// message: 要广播的字节数据
|
|||
|
|
func (h *Hub) BroadcastMessage(message []byte) {
|
|||
|
|
h.broadcast <- message
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ClientCount 返回当前在线客户端数量
|
|||
|
|
func (h *Hub) ClientCount() int {
|
|||
|
|
h.mu.RLock()
|
|||
|
|
defer h.mu.RUnlock()
|
|||
|
|
return len(h.clients)
|
|||
|
|
}
|