318 lines
7.0 KiB
Go
318 lines
7.0 KiB
Go
|
package ws
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"encoding/json"
|
|||
|
|
"sync"
|
|||
|
|
"sync/atomic"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"go.uber.org/zap"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
defaultUserBufferSize = 128
|
|||
|
|
maxReplayEvents = 200
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// Event 表示一个WebSocket事件
|
|||
|
|
type Event struct {
|
|||
|
|
ID uint64 `json:"event_id"`
|
|||
|
|
Type string `json:"type"`
|
|||
|
|
Event string `json:"event"` // 事件名称（兼容多种协议）
|
|||
|
|
TS int64 `json:"ts"`
|
|||
|
|
Payload interface{} `json:"payload"`
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Client 表示一个WebSocket客户端连接
|
|||
|
|
type Client struct {
|
|||
|
|
ID uint64
|
|||
|
|
UserID string
|
|||
|
|
Send chan []byte
|
|||
|
|
Quit chan struct{}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Message 表示WebSocket消息
|
|||
|
|
type Message struct {
|
|||
|
|
Type string `json:"type"`
|
|||
|
|
Payload json.RawMessage `json:"payload"`
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ResponseMessage 表示发送给客户端的消息
|
|||
|
|
type ResponseMessage struct {
|
|||
|
|
EventID uint64 `json:"event_id"`
|
|||
|
|
Type string `json:"type"`
|
|||
|
|
TS int64 `json:"ts"`
|
|||
|
|
Payload interface{} `json:"payload,omitempty"`
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// ErrorMessage 表示错误消息
|
|||
|
|
type ErrorMessage struct {
|
|||
|
|
Type string `json:"type"`
|
|||
|
|
Payload struct {
|
|||
|
|
Code string `json:"code"`
|
|||
|
|
Message string `json:"message"`
|
|||
|
|
} `json:"payload"`
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// subscriber 表示一个事件订阅者
|
|||
|
|
type subscriber struct {
|
|||
|
|
id uint64
|
|||
|
|
ch chan Event
|
|||
|
|
quit chan struct{}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Hub WebSocket连接管理器
|
|||
|
|
type Hub struct {
|
|||
|
|
seq uint64
|
|||
|
|
|
|||
|
|
mu sync.RWMutex
|
|||
|
|
clients map[string]map[uint64]*Client // userID -> clientID -> Client
|
|||
|
|
history map[string][]Event // userID -> []Event (用于断线重连历史回放)
|
|||
|
|
subscribers map[string]map[uint64]*subscriber // userID -> subscriberID -> subscriber
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NewHub 创建WebSocket Hub
|
|||
|
|
func NewHub() *Hub {
|
|||
|
|
return &Hub{
|
|||
|
|
clients: make(map[string]map[uint64]*Client),
|
|||
|
|
history: make(map[string][]Event),
|
|||
|
|
subscribers: make(map[string]map[uint64]*subscriber),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NextID 获取下一个事件ID
|
|||
|
|
func (h *Hub) NextID() uint64 {
|
|||
|
|
return atomic.AddUint64(&h.seq, 1)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Register 注册客户端连接
|
|||
|
|
func (h *Hub) Register(client *Client) []Event {
|
|||
|
|
h.mu.Lock()
|
|||
|
|
defer h.mu.Unlock()
|
|||
|
|
|
|||
|
|
if _, ok := h.clients[client.UserID]; !ok {
|
|||
|
|
h.clients[client.UserID] = make(map[uint64]*Client)
|
|||
|
|
}
|
|||
|
|
h.clients[client.UserID][client.ID] = client
|
|||
|
|
|
|||
|
|
// 获取历史事件用于回放
|
|||
|
|
replay := make([]Event, 0)
|
|||
|
|
for _, e := range h.history[client.UserID] {
|
|||
|
|
replay = append(replay, e)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
zap.L().Debug("WebSocket client registered",
|
|||
|
|
zap.String("user_id", client.UserID),
|
|||
|
|
zap.Uint64("client_id", client.ID),
|
|||
|
|
zap.Int("total_clients", len(h.clients[client.UserID])),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return replay
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Unregister 注销客户端连接
|
|||
|
|
func (h *Hub) Unregister(client *Client) {
|
|||
|
|
h.mu.Lock()
|
|||
|
|
defer h.mu.Unlock()
|
|||
|
|
|
|||
|
|
if userClients, ok := h.clients[client.UserID]; ok {
|
|||
|
|
if _, exists := userClients[client.ID]; exists {
|
|||
|
|
delete(userClients, client.ID)
|
|||
|
|
close(client.Quit)
|
|||
|
|
|
|||
|
|
if len(userClients) == 0 {
|
|||
|
|
delete(h.clients, client.UserID)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
zap.L().Debug("WebSocket client unregistered",
|
|||
|
|
zap.String("user_id", client.UserID),
|
|||
|
|
zap.Uint64("client_id", client.ID),
|
|||
|
|
)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// HasClients 检查用户是否有活跃连接
|
|||
|
|
func (h *Hub) HasClients(userID string) bool {
|
|||
|
|
h.mu.RLock()
|
|||
|
|
defer h.mu.RUnlock()
|
|||
|
|
return len(h.clients[userID]) > 0
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// GetClientCount 获取用户的连接数
|
|||
|
|
func (h *Hub) GetClientCount(userID string) int {
|
|||
|
|
h.mu.RLock()
|
|||
|
|
defer h.mu.RUnlock()
|
|||
|
|
return len(h.clients[userID])
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// PublishToUser 向指定用户发送事件
|
|||
|
|
func (h *Hub) PublishToUser(userID string, eventType string, payload interface{}) Event {
|
|||
|
|
ev := Event{
|
|||
|
|
ID: h.NextID(),
|
|||
|
|
Type: eventType,
|
|||
|
|
TS: time.Now().UnixMilli(),
|
|||
|
|
Payload: payload,
|
|||
|
|
}
|
|||
|
|
h.publish(userID, ev)
|
|||
|
|
return ev
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// PublishToUsers 向多个用户发送事件
|
|||
|
|
func (h *Hub) PublishToUsers(userIDs []string, eventType string, payload interface{}) {
|
|||
|
|
for _, uid := range userIDs {
|
|||
|
|
h.PublishToUser(uid, eventType, payload)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// publish 内部发布方法
|
|||
|
|
func (h *Hub) publish(userID string, ev Event) {
|
|||
|
|
h.mu.Lock()
|
|||
|
|
// 存储历史事件（最多保留200条）
|
|||
|
|
history := append(h.history[userID], ev)
|
|||
|
|
if len(history) > maxReplayEvents {
|
|||
|
|
history = history[len(history)-maxReplayEvents:]
|
|||
|
|
}
|
|||
|
|
h.history[userID] = history
|
|||
|
|
|
|||
|
|
// 复制客户端列表避免锁竞争
|
|||
|
|
targets := make([]*Client, 0, len(h.clients[userID]))
|
|||
|
|
for _, c := range h.clients[userID] {
|
|||
|
|
targets = append(targets, c)
|
|||
|
|
}
|
|||
|
|
h.mu.Unlock()
|
|||
|
|
|
|||
|
|
// 构建消息
|
|||
|
|
msg := ResponseMessage{
|
|||
|
|
EventID: ev.ID,
|
|||
|
|
Type: ev.Type,
|
|||
|
|
TS: ev.TS,
|
|||
|
|
Payload: ev.Payload,
|
|||
|
|
}
|
|||
|
|
data, err := json.Marshal(msg)
|
|||
|
|
if err != nil {
|
|||
|
|
zap.L().Error("Failed to marshal WebSocket message",
|
|||
|
|
zap.String("user_id", userID),
|
|||
|
|
zap.String("event_type", ev.Type),
|
|||
|
|
zap.Error(err),
|
|||
|
|
)
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 非阻塞发送，慢消费者丢弃消息
|
|||
|
|
for _, c := range targets {
|
|||
|
|
select {
|
|||
|
|
case <-c.Quit:
|
|||
|
|
// 客户端已断开
|
|||
|
|
case c.Send <- data:
|
|||
|
|
// 消息已发送
|
|||
|
|
default:
|
|||
|
|
// 发送缓冲区满，丢弃消息
|
|||
|
|
zap.L().Warn("WebSocket client buffer full, dropping message",
|
|||
|
|
zap.String("user_id", userID),
|
|||
|
|
zap.Uint64("client_id", c.ID),
|
|||
|
|
)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// SendError 向客户端发送错误消息
|
|||
|
|
func (h *Hub) SendError(client *Client, code string, message string) {
|
|||
|
|
errMsg := ErrorMessage{
|
|||
|
|
Type: "error",
|
|||
|
|
}
|
|||
|
|
errMsg.Payload.Code = code
|
|||
|
|
errMsg.Payload.Message = message
|
|||
|
|
|
|||
|
|
data, err := json.Marshal(errMsg)
|
|||
|
|
if err != nil {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case <-client.Quit:
|
|||
|
|
case client.Send <- data:
|
|||
|
|
default:
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Broadcast 广播消息给所有连接
|
|||
|
|
func (h *Hub) Broadcast(eventType string, payload interface{}) {
|
|||
|
|
h.mu.RLock()
|
|||
|
|
userIDs := make([]string, 0, len(h.clients))
|
|||
|
|
for uid := range h.clients {
|
|||
|
|
userIDs = append(userIDs, uid)
|
|||
|
|
}
|
|||
|
|
h.mu.RUnlock()
|
|||
|
|
|
|||
|
|
h.PublishToUsers(userIDs, eventType, payload)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// GetOnlineUsers 获取在线用户列表
|
|||
|
|
func (h *Hub) GetOnlineUsers() []string {
|
|||
|
|
h.mu.RLock()
|
|||
|
|
defer h.mu.RUnlock()
|
|||
|
|
|
|||
|
|
users := make([]string, 0, len(h.clients))
|
|||
|
|
for uid := range h.clients {
|
|||
|
|
users = append(users, uid)
|
|||
|
|
}
|
|||
|
|
return users
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// GetOnlineCount 获取在线用户数
|
|||
|
|
func (h *Hub) GetOnlineCount() int {
|
|||
|
|
h.mu.RLock()
|
|||
|
|
defer h.mu.RUnlock()
|
|||
|
|
return len(h.clients)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Subscribe 订阅事件
|
|||
|
|
func (h *Hub) Subscribe(userID string, afterID uint64) (chan Event, func(), []Event) {
|
|||
|
|
subID := h.NextID()
|
|||
|
|
sub := &subscriber{
|
|||
|
|
id: subID,
|
|||
|
|
ch: make(chan Event, defaultUserBufferSize),
|
|||
|
|
quit: make(chan struct{}),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
h.mu.Lock()
|
|||
|
|
if _, ok := h.subscribers[userID]; !ok {
|
|||
|
|
h.subscribers[userID] = make(map[uint64]*subscriber)
|
|||
|
|
}
|
|||
|
|
h.subscribers[userID][subID] = sub
|
|||
|
|
|
|||
|
|
replay := make([]Event, 0)
|
|||
|
|
for _, e := range h.history[userID] {
|
|||
|
|
if e.ID > afterID {
|
|||
|
|
replay = append(replay, e)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
h.mu.Unlock()
|
|||
|
|
|
|||
|
|
cancel := func() {
|
|||
|
|
h.mu.Lock()
|
|||
|
|
defer h.mu.Unlock()
|
|||
|
|
if userSubs, ok := h.subscribers[userID]; ok {
|
|||
|
|
if s, exists := userSubs[subID]; exists {
|
|||
|
|
close(s.quit)
|
|||
|
|
delete(userSubs, subID)
|
|||
|
|
close(s.ch)
|
|||
|
|
}
|
|||
|
|
if len(userSubs) == 0 {
|
|||
|
|
delete(h.subscribers, userID)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return sub.ch, cancel, replay
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// EncodeData 编码事件数据
|
|||
|
|
func EncodeData(ev Event) (string, error) {
|
|||
|
|
body, err := json.Marshal(ev.Payload)
|
|||
|
|
if err != nil {
|
|||
|
|
return "", err
|
|||
|
|
}
|
|||
|
|
return string(body), nil
|
|||
|
|
}
|