bfb4fb0ccb
- pkg/ws/message.go: 统一消息协议（Message/Response/PushMessage） - pkg/ws/hub.go: Hub 连接管理（注册/注销/按 userID 查找/在线计数） - pkg/ws/client.go: 客户端连接封装（readPump/writePump/心跳 30s） - pkg/ws/pubsub.go: Redis Pub/Sub 消息路由（按用户频道发布/订阅） - app/ws/handler.go: WebSocket 升级处理（JWT 认证 + 消息分发） - app/ws/router.go: GET /ws 路由注册 - app/ws/provider.go: Wire Provider Set - 更新 provider/router 集成 WebSocket 模块 Made-with: Cursor
99 lines
2.3 KiB
Go
99 lines
2.3 KiB
Go
package ws
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"sync"
|
||
|
||
"github.com/echochat/backend/pkg/logs"
|
||
"github.com/redis/go-redis/v9"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
const channelPrefix = "echo:ws:channel:"
|
||
|
||
// PubSub 封装 Redis Pub/Sub，提供按用户频道的消息发布和订阅
|
||
type PubSub struct {
|
||
rdb *redis.Client
|
||
hub *Hub
|
||
subs map[int64]context.CancelFunc // userID -> 订阅取消函数
|
||
mu sync.Mutex
|
||
}
|
||
|
||
// NewPubSub 创建 PubSub 实例
|
||
func NewPubSub(rdb *redis.Client, hub *Hub) *PubSub {
|
||
return &PubSub{
|
||
rdb: rdb,
|
||
hub: hub,
|
||
subs: make(map[int64]context.CancelFunc),
|
||
}
|
||
}
|
||
|
||
// channelName 生成用户专属的 Redis 频道名
|
||
func channelName(userID int64) string {
|
||
return fmt.Sprintf("%s%d", channelPrefix, userID)
|
||
}
|
||
|
||
// Publish 向指定用户的频道发布消息
|
||
func (ps *PubSub) Publish(ctx context.Context, userID int64, data []byte) error {
|
||
channel := channelName(userID)
|
||
return ps.rdb.Publish(ctx, channel, data).Err()
|
||
}
|
||
|
||
// Subscribe 订阅指定用户的频道
|
||
// 收到消息后自动转发给本地 Hub 中对应的 Client
|
||
func (ps *PubSub) Subscribe(userID int64) {
|
||
ps.mu.Lock()
|
||
if cancel, ok := ps.subs[userID]; ok {
|
||
cancel()
|
||
}
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
ps.subs[userID] = cancel
|
||
ps.mu.Unlock()
|
||
|
||
channel := channelName(userID)
|
||
sub := ps.rdb.Subscribe(ctx, channel)
|
||
|
||
go func() {
|
||
defer sub.Close()
|
||
ch := sub.Channel()
|
||
|
||
logs.Info(nil, "ws.pubsub.Subscribe", "开始订阅用户频道",
|
||
zap.Int64("user_id", userID), zap.String("channel", channel))
|
||
|
||
for {
|
||
select {
|
||
case <-ctx.Done():
|
||
logs.Info(nil, "ws.pubsub.Subscribe", "取消订阅用户频道",
|
||
zap.Int64("user_id", userID))
|
||
return
|
||
case msg, ok := <-ch:
|
||
if !ok {
|
||
return
|
||
}
|
||
ps.hub.SendToUser(userID, []byte(msg.Payload))
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// Unsubscribe 取消订阅指定用户的频道
|
||
func (ps *PubSub) Unsubscribe(userID int64) {
|
||
ps.mu.Lock()
|
||
defer ps.mu.Unlock()
|
||
|
||
if cancel, ok := ps.subs[userID]; ok {
|
||
cancel()
|
||
delete(ps.subs, userID)
|
||
}
|
||
}
|
||
|
||
// PublishToUser 便捷方法：序列化推送消息并发布到用户频道
|
||
func (ps *PubSub) PublishToUser(ctx context.Context, userID int64, msg *PushMessage) error {
|
||
data, err := MarshalPush(msg)
|
||
if err != nil {
|
||
return fmt.Errorf("序列化推送消息失败: %w", err)
|
||
}
|
||
return ps.Publish(ctx, userID, data)
|
||
}
|