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)
|
||
}
|