7474a92e3b
P5 阶段交付物： - services/msg: 消息通知服务（NestJS） - notifications: 发送通知 + ES 全文检索 + search - config/elasticsearch.ts: ES Client 单例 - package.json: 补充 @opentelemetry/sdk-node + exporter-trace-otlp-http - services/push-gateway: WebSocket 推送网关（Go Gin） - internal/hub/hub.go: WebSocket 连接池管理（Register/Unregister/SendToUser） - internal/ws/handler.go: JWT 鉴权 + WebSocket 升级 + 内部推送 API - services/ai: AI 辅助服务（Python FastAPI） - /chat + /chat/stream（SSE 流式） - /generate/question + /optimize/expression - config.py: OpenAI 兼容 API 配置 - packages/shared-proto/proto/msg.proto: NotificationService 契约（send/search） - packages/shared-proto/proto/ai.proto: AiService 契约（含 stream 方法）
60 lines
1.1 KiB
Go
60 lines
1.1 KiB
Go
package hub
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// Hub 管理 WebSocket 客户端连接
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
clients map[string]map[*websocket.Conn]bool // userID -> connections
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
return &Hub{
|
|
clients: make(map[string]map[*websocket.Conn]bool),
|
|
}
|
|
}
|
|
|
|
func (h *Hub) Register(userID string, conn *websocket.Conn) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
if h.clients[userID] == nil {
|
|
h.clients[userID] = make(map[*websocket.Conn]bool)
|
|
}
|
|
h.clients[userID][conn] = true
|
|
}
|
|
|
|
func (h *Hub) Unregister(userID string, conn *websocket.Conn) {
|
|
h.mu.Lock()
|
|
defer h.mu.Unlock()
|
|
|
|
if conns, ok := h.clients[userID]; ok {
|
|
delete(conns, conn)
|
|
if len(conns) == 0 {
|
|
delete(h.clients, userID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// SendToUser 向指定用户的所有连接推送消息
|
|
func (h *Hub) SendToUser(userID string, message []byte) error {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
|
|
conns, ok := h.clients[userID]
|
|
if !ok {
|
|
return nil
|
|
}
|
|
|
|
for conn := range conns {
|
|
if err := conn.WriteMessage(websocket.TextMessage, message); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|