feat(ws): WebSocket 核心模块（Hub + Client + PubSub + Handler）
- 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
This commit is contained in:
bujinyuan
123
backend/go-service/pkg/ws/client.go
Normal file
123
backend/go-service/pkg/ws/client.go
Normal file
@@ -0,0 +1,123 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
|
||||
"github.com/echochat/backend/pkg/logs"
|
||||
"github.com/gorilla/websocket"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
writeWait = 10 * time.Second // 写超时
|
||||
pongWait = 60 * time.Second // 等待 pong 的超时
|
||||
pingPeriod = 30 * time.Second // 心跳发送间隔（必须小于 pongWait）
|
||||
maxMessageSize = 4096 // 单条消息最大字节数
|
||||
sendBufSize = 256 // 发送缓冲区大小
|
||||
)
|
||||
|
||||
// Client 封装单个 WebSocket 客户端连接
|
||||
// 每个连接持有两个 goroutine：readPump（读取客户端消息）和 writePump（写入消息到客户端）
|
||||
type Client struct {
|
||||
hub *Hub
|
||||
conn *websocket.Conn
|
||||
send chan []byte // 待发送消息缓冲队列
|
||||
UserID int64 // 关联的用户 ID
|
||||
}
|
||||
|
||||
// NewClient 创建客户端实例
|
||||
func NewClient(hub *Hub, conn *websocket.Conn, userID int64) *Client {
|
||||
return &Client{
|
||||
hub: hub,
|
||||
conn: conn,
|
||||
send: make(chan []byte, sendBufSize),
|
||||
UserID: userID,
|
||||
}
|
||||
}
|
||||
|
||||
// MessageHandler 消息处理回调函数类型
|
||||
type MessageHandler func(client *Client, msg *Message)
|
||||
|
||||
// ReadPump 读取客户端消息的循环
|
||||
// 当连接断开或出错时退出，并触发注销流程
|
||||
func (c *Client) ReadPump(onMessage MessageHandler) {
|
||||
defer func() {
|
||||
c.hub.Unregister(c)
|
||||
c.conn.Close()
|
||||
}()
|
||||
|
||||
c.conn.SetReadLimit(maxMessageSize)
|
||||
c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
c.conn.SetPongHandler(func(string) error {
|
||||
c.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
return nil
|
||||
})
|
||||
|
||||
for {
|
||||
_, rawMsg, err := c.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
|
||||
logs.Warn(nil, "ws.client.ReadPump", "WebSocket 异常关闭",
|
||||
zap.Int64("user_id", c.UserID), zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var msg Message
|
||||
if err := json.Unmarshal(rawMsg, &msg); err != nil {
|
||||
logs.Warn(nil, "ws.client.ReadPump", "消息格式解析失败",
|
||||
zap.Int64("user_id", c.UserID), zap.Error(err))
|
||||
continue
|
||||
}
|
||||
|
||||
if onMessage != nil {
|
||||
onMessage(c, &msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// WritePump 向客户端写入消息的循环
|
||||
// 从 send channel 读取消息并写入 WebSocket，同时负责心跳
|
||||
func (c *Client) WritePump() {
|
||||
ticker := time.NewTicker(pingPeriod)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
c.conn.Close()
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case message, ok := <-c.send:
|
||||
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if !ok {
|
||||
c.conn.WriteMessage(websocket.CloseMessage, []byte{})
|
||||
return
|
||||
}
|
||||
|
||||
if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
|
||||
logs.Warn(nil, "ws.client.WritePump", "写入消息失败",
|
||||
zap.Int64("user_id", c.UserID), zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
case <-ticker.C:
|
||||
c.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Send 向客户端发送消息（非阻塞，缓冲区满时丢弃）
|
||||
func (c *Client) Send(data []byte) bool {
|
||||
select {
|
||||
case c.send <- data:
|
||||
return true
|
||||
default:
|
||||
logs.Warn(nil, "ws.client.Send", "发送缓冲区已满，丢弃消息",
|
||||
zap.Int64("user_id", c.UserID))
|
||||
return false
|
||||
}
|
||||
}
|
||||
123
backend/go-service/pkg/ws/hub.go
Normal file
123
backend/go-service/pkg/ws/hub.go
Normal file
@@ -0,0 +1,123 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/echochat/backend/pkg/logs"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// Hub 管理所有活跃的 WebSocket 客户端连接
|
||||
// 提供按 userID 注册/注销/查找连接的能力
|
||||
type Hub struct {
|
||||
clients map[int64]*Client // userID -> Client 映射
|
||||
register chan *Client // 注册通道
|
||||
unregister chan *Client // 注销通道
|
||||
mu sync.RWMutex // 保护 clients map 的读写锁
|
||||
}
|
||||
|
||||
// NewHub 创建 Hub 实例
|
||||
func NewHub() *Hub {
|
||||
return &Hub{
|
||||
clients: make(map[int64]*Client),
|
||||
register: make(chan *Client, 256),
|
||||
unregister: make(chan *Client, 256),
|
||||
}
|
||||
}
|
||||
|
||||
// Run 启动 Hub 主循环，处理连接注册和注销
|
||||
// 应在单独的 goroutine 中运行
|
||||
func (h *Hub) Run() {
|
||||
for {
|
||||
select {
|
||||
case client := <-h.register:
|
||||
h.mu.Lock()
|
||||
if old, ok := h.clients[client.UserID]; ok {
|
||||
logs.Info(nil, "ws.hub.Run", "用户重复连接，关闭旧连接",
|
||||
zap.Int64("user_id", client.UserID))
|
||||
close(old.send)
|
||||
}
|
||||
h.clients[client.UserID] = client
|
||||
h.mu.Unlock()
|
||||
|
||||
logs.Info(nil, "ws.hub.Run", "客户端已注册",
|
||||
zap.Int64("user_id", client.UserID),
|
||||
zap.Int("online_count", h.OnlineCount()))
|
||||
|
||||
case client := <-h.unregister:
|
||||
h.mu.Lock()
|
||||
if existing, ok := h.clients[client.UserID]; ok && existing == client {
|
||||
delete(h.clients, client.UserID)
|
||||
close(client.send)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
logs.Info(nil, "ws.hub.Run", "客户端已注销",
|
||||
zap.Int64("user_id", client.UserID),
|
||||
zap.Int("online_count", h.OnlineCount()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Register 注册客户端连接
|
||||
func (h *Hub) Register(client *Client) {
|
||||
h.register <- client
|
||||
}
|
||||
|
||||
// Unregister 注销客户端连接
|
||||
func (h *Hub) Unregister(client *Client) {
|
||||
h.unregister <- client
|
||||
}
|
||||
|
||||
// GetClient 根据 userID 获取客户端连接（线程安全）
|
||||
func (h *Hub) GetClient(userID int64) (*Client, bool) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
client, ok := h.clients[userID]
|
||||
return client, ok
|
||||
}
|
||||
|
||||
// SendToUser 向指定用户发送消息（仅本地 Hub）
|
||||
// 如果用户不在本实例，返回 false
|
||||
func (h *Hub) SendToUser(userID int64, data []byte) bool {
|
||||
h.mu.RLock()
|
||||
client, ok := h.clients[userID]
|
||||
h.mu.RUnlock()
|
||||
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
|
||||
select {
|
||||
case client.send <- data:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// OnlineCount 返回当前在线连接数
|
||||
func (h *Hub) OnlineCount() int {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return len(h.clients)
|
||||
}
|
||||
|
||||
// OnlineUserIDs 返回所有在线用户 ID 列表
|
||||
func (h *Hub) OnlineUserIDs() []int64 {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
ids := make([]int64, 0, len(h.clients))
|
||||
for id := range h.clients {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
// IsOnline 检查指定用户是否在线（本地 Hub）
|
||||
func (h *Hub) IsOnline(userID int64) bool {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
_, ok := h.clients[userID]
|
||||
return ok
|
||||
}
|
||||
63
backend/go-service/pkg/ws/message.go
Normal file
63
backend/go-service/pkg/ws/message.go
Normal file
@@ -0,0 +1,63 @@
|
||||
// Package ws 提供 WebSocket 通讯基础设施
|
||||
// 包含消息协议定义、连接管理、Redis Pub/Sub 消息路由
|
||||
package ws
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Message 客户端 → 服务端的 WebSocket 消息
|
||||
// 事件命名规范：{模块}.{对象}.{动作}，如 contact.request.send
|
||||
type Message struct {
|
||||
Event string `json:"event"` // 事件类型
|
||||
Seq int64 `json:"seq"` // 客户端消息序号，用于 ACK 匹配
|
||||
Data json.RawMessage `json:"data,omitempty"` // 业务数据（延迟解析）
|
||||
Time string `json:"time"` // 发送时间
|
||||
}
|
||||
|
||||
// Response 服务端 → 客户端的 ACK 响应
|
||||
type Response struct {
|
||||
Event string `json:"event"` // 事件类型（原事件 + ".ack" 后缀）
|
||||
Seq int64 `json:"seq"` // 对应请求的序号
|
||||
Code int `json:"code"` // 状态码，0=成功
|
||||
Message string `json:"message"` // 状态描述
|
||||
Data interface{} `json:"data,omitempty"` // 响应数据
|
||||
}
|
||||
|
||||
// PushMessage 服务端主动推送给客户端的消息
|
||||
type PushMessage struct {
|
||||
Event string `json:"event"` // 事件类型
|
||||
Data interface{} `json:"data,omitempty"` // 推送数据
|
||||
Time string `json:"time"` // 推送时间
|
||||
}
|
||||
|
||||
// NewPushMessage 创建一条推送消息
|
||||
func NewPushMessage(event string, data interface{}) *PushMessage {
|
||||
return &PushMessage{
|
||||
Event: event,
|
||||
Data: data,
|
||||
Time: time.Now().Format("2006-01-02 15:04:05"),
|
||||
}
|
||||
}
|
||||
|
||||
// NewResponse 创建一条 ACK 响应
|
||||
func NewResponse(event string, seq int64, code int, message string, data interface{}) *Response {
|
||||
return &Response{
|
||||
Event: event + ".ack",
|
||||
Seq: seq,
|
||||
Code: code,
|
||||
Message: message,
|
||||
Data: data,
|
||||
}
|
||||
}
|
||||
|
||||
// MarshalPush 将推送消息序列化为 JSON 字节
|
||||
func MarshalPush(msg *PushMessage) ([]byte, error) {
|
||||
return json.Marshal(msg)
|
||||
}
|
||||
|
||||
// MarshalResponse 将响应消息序列化为 JSON 字节
|
||||
func MarshalResponse(resp *Response) ([]byte, error) {
|
||||
return json.Marshal(resp)
|
||||
}
|
||||
98
backend/go-service/pkg/ws/pubsub.go
Normal file
98
backend/go-service/pkg/ws/pubsub.go
Normal file
@@ -0,0 +1,98 @@
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user