lan
6bf87fec46
Introduce several performance and synchronization enhancements: - Implement `SeqBufferManager` to allow sequence number pre-allocation via Redis Lua scripts, reducing atomic increment overhead. - Add `PushWorker` to handle asynchronous message pushing using Redis Streams. - Implement incremental conversation synchronization via `ConversationVersionLog` to allow clients to fetch only recent changes. - Add support for Gzip compression in WebSocket communications to reduce bandwidth usage. - Update dependency injection and configuration to support these new components.
655 lines
15 KiB
Go
655 lines
15 KiB
Go
package ws
|
||
|
||
import (
|
||
"encoding/json"
|
||
"errors"
|
||
"sync"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
const (
|
||
defaultUserBufferSize = 128
|
||
groupPushConcurrency = 64
|
||
groupPushQueueSize = 512
|
||
maxTotalConnections = 100000
|
||
maxConnectionsPerUser = 10
|
||
ackTimeout = 30 * time.Second
|
||
ackCheckInterval = 10 * time.Second
|
||
)
|
||
|
||
// Event 表示一个WebSocket事件
|
||
type Event struct {
|
||
ID uint64 `json:"event_id"`
|
||
Type string `json:"type"`
|
||
Event string `json:"event"` // 事件名称（兼容多种协议）
|
||
TS int64 `json:"ts"`
|
||
Payload any `json:"payload"`
|
||
}
|
||
|
||
// Client 表示一个WebSocket客户端连接
|
||
type Client struct {
|
||
ID uint64
|
||
UserID string
|
||
Send chan []byte
|
||
Quit chan struct{}
|
||
PendingAcks map[string]time.Time
|
||
CompressEnabled bool // 客户端是否启用 gzip 压缩
|
||
}
|
||
|
||
// ErrConnectionLimit 连接数限制错误
|
||
var ErrConnectionLimit = errors.New("connection limit reached")
|
||
|
||
// 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 any `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{}
|
||
}
|
||
|
||
// DisconnectHandler 用户断开连接时的回调函数
|
||
// userID: 断开的用户ID
|
||
// remainingCount: 该用户剩余的连接数（0表示用户完全离线）
|
||
type DisconnectHandler func(userID string, remainingCount int)
|
||
|
||
// ConnectHandler 用户连接时的回调函数
|
||
// userID: 连接的用户ID
|
||
// isFirstConnection: 是否是该用户的第一个连接（之前完全离线）
|
||
type ConnectHandler func(userID string, isFirstConnection bool)
|
||
|
||
// Hub WebSocket连接管理器
|
||
type Hub struct {
|
||
seq uint64
|
||
connCount atomic.Int64
|
||
|
||
mu sync.RWMutex
|
||
clients map[string]map[uint64]*Client
|
||
subscribers map[string]map[uint64]*subscriber
|
||
disconnectHandlers []DisconnectHandler
|
||
connectHandlers []ConnectHandler
|
||
}
|
||
|
||
// NewHub 创建WebSocket Hub
|
||
func NewHub() *Hub {
|
||
return &Hub{
|
||
clients: make(map[string]map[uint64]*Client),
|
||
subscribers: make(map[string]map[uint64]*subscriber),
|
||
disconnectHandlers: make([]DisconnectHandler, 0),
|
||
connectHandlers: make([]ConnectHandler, 0),
|
||
}
|
||
}
|
||
|
||
// NextID 获取下一个事件ID
|
||
func (h *Hub) NextID() uint64 {
|
||
return atomic.AddUint64(&h.seq, 1)
|
||
}
|
||
|
||
// Register 注册客户端连接，返回可能的错误（连接数超限）
|
||
func (h *Hub) Register(client *Client) error {
|
||
if h.connCount.Load() >= int64(maxTotalConnections) {
|
||
return ErrConnectionLimit
|
||
}
|
||
|
||
h.mu.Lock()
|
||
|
||
if len(h.clients[client.UserID]) >= maxConnectionsPerUser {
|
||
h.mu.Unlock()
|
||
return ErrConnectionLimit
|
||
}
|
||
|
||
_, existed := h.clients[client.UserID]
|
||
isFirstConnection := !existed
|
||
|
||
if !existed {
|
||
h.clients[client.UserID] = make(map[uint64]*Client)
|
||
}
|
||
h.clients[client.UserID][client.ID] = client
|
||
h.connCount.Add(1)
|
||
|
||
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])),
|
||
zap.Int64("total_connections", h.connCount.Load()),
|
||
zap.Bool("is_first_connection", isFirstConnection),
|
||
)
|
||
|
||
handlers := make([]ConnectHandler, len(h.connectHandlers))
|
||
copy(handlers, h.connectHandlers)
|
||
h.mu.Unlock()
|
||
|
||
if len(handlers) > 0 {
|
||
go func() {
|
||
for _, handler := range handlers {
|
||
handler(client.UserID, isFirstConnection)
|
||
}
|
||
}()
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
// Unregister 注销客户端连接
|
||
func (h *Hub) Unregister(client *Client) {
|
||
h.mu.Lock()
|
||
|
||
var remainingCount int
|
||
if userClients, ok := h.clients[client.UserID]; ok {
|
||
if _, exists := userClients[client.ID]; exists {
|
||
delete(userClients, client.ID)
|
||
close(client.Quit)
|
||
h.connCount.Add(-1)
|
||
remainingCount = len(userClients)
|
||
|
||
if remainingCount == 0 {
|
||
delete(h.clients, client.UserID)
|
||
}
|
||
|
||
zap.L().Debug("WebSocket client unregistered",
|
||
zap.String("user_id", client.UserID),
|
||
zap.Uint64("client_id", client.ID),
|
||
zap.Int("remaining_connections", remainingCount),
|
||
)
|
||
}
|
||
}
|
||
|
||
// 复制处理器列表避免锁竞争
|
||
handlers := make([]DisconnectHandler, len(h.disconnectHandlers))
|
||
copy(handlers, h.disconnectHandlers)
|
||
h.mu.Unlock()
|
||
|
||
// 异步调用断开处理器，避免阻塞 Unregister
|
||
if len(handlers) > 0 {
|
||
go func() {
|
||
for _, handler := range handlers {
|
||
handler(client.UserID, remainingCount)
|
||
}
|
||
}()
|
||
}
|
||
}
|
||
|
||
// HasClients 检查用户是否有活跃连接
|
||
func (h *Hub) HasClients(userID string) bool {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
return len(h.clients[userID]) > 0
|
||
}
|
||
|
||
// FilterOnline 批量检查用户在线状态，一次加锁返回在线/离线列表
|
||
func (h *Hub) FilterOnline(userIDs []string) (online []string, offline []string) {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
for _, uid := range userIDs {
|
||
if len(h.clients[uid]) > 0 {
|
||
online = append(online, uid)
|
||
} else {
|
||
offline = append(offline, uid)
|
||
}
|
||
}
|
||
return
|
||
}
|
||
|
||
// OnDisconnect 注册断开连接事件处理器
|
||
// 当用户断开 WebSocket 连接时，会调用所有注册的处理器
|
||
// 处理器会收到用户ID和该用户剩余的连接数
|
||
func (h *Hub) OnDisconnect(handler DisconnectHandler) {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
h.disconnectHandlers = append(h.disconnectHandlers, handler)
|
||
}
|
||
|
||
// OnConnect 注册连接事件处理器
|
||
// 当用户建立 WebSocket 连接时，会调用所有注册的处理器
|
||
// 处理器会收到用户ID和是否是该用户的第一个连接
|
||
func (h *Hub) OnConnect(handler ConnectHandler) {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
h.connectHandlers = append(h.connectHandlers, handler)
|
||
}
|
||
|
||
// 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 any) 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 any) {
|
||
if len(userIDs) == 0 {
|
||
return
|
||
}
|
||
|
||
ev := Event{
|
||
ID: h.NextID(),
|
||
Type: eventType,
|
||
TS: time.Now().UnixMilli(),
|
||
Payload: payload,
|
||
}
|
||
|
||
msg := ResponseMessage{
|
||
EventID: ev.ID,
|
||
Type: eventType,
|
||
TS: ev.TS,
|
||
Payload: payload,
|
||
}
|
||
data, err := json.Marshal(msg)
|
||
if err != nil {
|
||
zap.L().Error("Failed to marshal WebSocket message",
|
||
zap.String("event_type", eventType),
|
||
zap.Error(err),
|
||
)
|
||
return
|
||
}
|
||
|
||
for _, uid := range userIDs {
|
||
h.publishPreMarshaled(uid, ev, data)
|
||
}
|
||
}
|
||
|
||
// publishPreMarshaled 使用预序列化的数据发送事件
|
||
func (h *Hub) publishPreMarshaled(userID string, ev Event, data []byte) {
|
||
h.mu.RLock()
|
||
targets := make([]*Client, 0, len(h.clients[userID]))
|
||
for _, c := range h.clients[userID] {
|
||
targets = append(targets, c)
|
||
}
|
||
h.mu.RUnlock()
|
||
|
||
for _, c := range targets {
|
||
select {
|
||
case <-c.Quit:
|
||
case c.Send <- data:
|
||
default:
|
||
zap.L().Warn("WebSocket client buffer full, closing slow client",
|
||
zap.String("user_id", userID),
|
||
zap.Uint64("client_id", c.ID),
|
||
)
|
||
close(c.Quit)
|
||
}
|
||
}
|
||
}
|
||
|
||
// PublishToUserOnline 仅在用户在线时发送事件，不存入历史回放
|
||
// 适用于通话信令等实时性消息，避免断线重连时重放过期消息
|
||
func (h *Hub) PublishToUserOnline(userID string, eventType string, payload any) bool {
|
||
h.mu.RLock()
|
||
targets := make([]*Client, 0, len(h.clients[userID]))
|
||
for _, c := range h.clients[userID] {
|
||
targets = append(targets, c)
|
||
}
|
||
h.mu.RUnlock()
|
||
|
||
if len(targets) == 0 {
|
||
return false
|
||
}
|
||
|
||
msg := ResponseMessage{
|
||
EventID: h.NextID(),
|
||
Type: eventType,
|
||
TS: time.Now().UnixMilli(),
|
||
Payload: payload,
|
||
}
|
||
data, err := json.Marshal(msg)
|
||
if err != nil {
|
||
return false
|
||
}
|
||
|
||
for _, c := range targets {
|
||
select {
|
||
case <-c.Quit:
|
||
case c.Send <- data:
|
||
default:
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// PublishToUserOnlineReliable 可靠地发送事件给在线用户（阻塞发送）
|
||
// 用于重要的信令消息（如 SDP, ICE candidate），确保消息送达
|
||
// 如果用户不在线，返回 false
|
||
func (h *Hub) PublishToUserOnlineReliable(userID string, eventType string, payload any) bool {
|
||
h.mu.RLock()
|
||
targets := make([]*Client, 0, len(h.clients[userID]))
|
||
for _, c := range h.clients[userID] {
|
||
targets = append(targets, c)
|
||
}
|
||
h.mu.RUnlock()
|
||
|
||
if len(targets) == 0 {
|
||
return false
|
||
}
|
||
|
||
msg := ResponseMessage{
|
||
EventID: h.NextID(),
|
||
Type: eventType,
|
||
TS: time.Now().UnixMilli(),
|
||
Payload: payload,
|
||
}
|
||
data, err := json.Marshal(msg)
|
||
if err != nil {
|
||
return false
|
||
}
|
||
|
||
// 阻塞发送，确保消息送达
|
||
for _, c := range targets {
|
||
select {
|
||
case <-c.Quit:
|
||
// 客户端已断开，跳过
|
||
case c.Send <- data:
|
||
// 消息已发送
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// publish 内部发布方法
|
||
func (h *Hub) publish(userID string, ev Event) {
|
||
h.mu.RLock()
|
||
targets := make([]*Client, 0, len(h.clients[userID]))
|
||
for _, c := range h.clients[userID] {
|
||
targets = append(targets, c)
|
||
}
|
||
h.mu.RUnlock()
|
||
|
||
// 构建消息
|
||
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, closing slow client",
|
||
zap.String("user_id", userID),
|
||
zap.Uint64("client_id", c.ID),
|
||
)
|
||
close(c.Quit)
|
||
}
|
||
}
|
||
}
|
||
|
||
// 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 any) {
|
||
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)
|
||
}
|
||
|
||
// PublishToUsersConcurrent 并发推送消息给多个用户（适用于群聊等大列表场景）
|
||
// 使用有界并发控制，避免 goroutine 爆炸
|
||
func (h *Hub) PublishToUsersConcurrent(userIDs []string, eventType string, payload any) {
|
||
if len(userIDs) == 0 {
|
||
return
|
||
}
|
||
|
||
ev := Event{
|
||
ID: h.NextID(),
|
||
Type: eventType,
|
||
TS: time.Now().UnixMilli(),
|
||
Payload: payload,
|
||
}
|
||
|
||
msg := ResponseMessage{
|
||
EventID: ev.ID,
|
||
Type: eventType,
|
||
TS: ev.TS,
|
||
Payload: payload,
|
||
}
|
||
data, err := json.Marshal(msg)
|
||
if err != nil {
|
||
zap.L().Error("Failed to marshal WebSocket message for batch push",
|
||
zap.String("event_type", eventType),
|
||
zap.Error(err),
|
||
)
|
||
return
|
||
}
|
||
|
||
sem := make(chan struct{}, groupPushConcurrency)
|
||
var wg sync.WaitGroup
|
||
wg.Add(len(userIDs))
|
||
|
||
for _, uid := range userIDs {
|
||
sem <- struct{}{}
|
||
go func(userID string) {
|
||
defer wg.Done()
|
||
defer func() { <-sem }()
|
||
|
||
h.mu.RLock()
|
||
targets := make([]*Client, 0, len(h.clients[userID]))
|
||
for _, c := range h.clients[userID] {
|
||
targets = append(targets, c)
|
||
}
|
||
h.mu.RUnlock()
|
||
|
||
for _, c := range targets {
|
||
select {
|
||
case <-c.Quit:
|
||
case c.Send <- data:
|
||
default:
|
||
zap.L().Warn("WebSocket client buffer full, closing slow client",
|
||
zap.String("user_id", userID),
|
||
zap.Uint64("client_id", c.ID),
|
||
)
|
||
close(c.Quit)
|
||
}
|
||
}
|
||
}(uid)
|
||
}
|
||
|
||
wg.Wait()
|
||
}
|
||
|
||
// 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)
|
||
}
|
||
|
||
// GetConnectionCount 获取总连接数
|
||
func (h *Hub) GetConnectionCount() int64 {
|
||
return h.connCount.Load()
|
||
}
|
||
|
||
// CloseAllConnections 关闭所有WebSocket连接（用于优雅关闭）
|
||
func (h *Hub) CloseAllConnections() {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
|
||
for _, userClients := range h.clients {
|
||
for _, c := range userClients {
|
||
select {
|
||
case <-c.Quit:
|
||
default:
|
||
close(c.Quit)
|
||
}
|
||
}
|
||
}
|
||
h.clients = make(map[string]map[uint64]*Client)
|
||
h.connCount.Store(0)
|
||
}
|
||
|
||
// Subscribe 订阅事件
|
||
func (h *Hub) Subscribe(userID string, afterID uint64) (chan Event, func()) {
|
||
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
|
||
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
|
||
}
|
||
|
||
// AckMessage 确认消息已送达，从 pendingAcks 中移除
|
||
func (h *Hub) AckMessage(userID, messageID string) bool {
|
||
h.mu.RLock()
|
||
defer h.mu.RUnlock()
|
||
|
||
for _, client := range h.clients[userID] {
|
||
if _, ok := client.PendingAcks[messageID]; ok {
|
||
delete(client.PendingAcks, messageID)
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// startAckChecker 启动超时 ACK 检查器
|
||
func (h *Hub) startAckChecker(onTimeout func(userID, messageID string)) {
|
||
ticker := time.NewTicker(ackCheckInterval)
|
||
defer ticker.Stop()
|
||
|
||
for range ticker.C {
|
||
h.mu.RLock()
|
||
// 收集超时的 ack
|
||
type timeoutEntry struct {
|
||
userID string
|
||
messageID string
|
||
}
|
||
var timeouts []timeoutEntry
|
||
|
||
now := time.Now()
|
||
for uid, userClients := range h.clients {
|
||
for _, client := range userClients {
|
||
for msgID, sentAt := range client.PendingAcks {
|
||
if now.Sub(sentAt) > ackTimeout {
|
||
timeouts = append(timeouts, timeoutEntry{uid, msgID})
|
||
delete(client.PendingAcks, msgID)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
h.mu.RUnlock()
|
||
|
||
for _, entry := range timeouts {
|
||
onTimeout(entry.userID, entry.messageID)
|
||
}
|
||
}
|
||
}
|
||
|
||
// StartAckChecker 启动 ACK 检查器（公开方法）
|
||
func (h *Hub) StartAckChecker(onTimeout func(userID, messageID string)) {
|
||
go h.startAckChecker(onTimeout)
|
||
}
|
||
|
||
// EncodeData 编码事件数据
|
||
func EncodeData(ev Event) (string, error) {
|
||
body, err := json.Marshal(ev.Payload)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
return string(body), nil
|
||
}
|