362 lines
10 KiB
Go
362 lines
10 KiB
Go
// Package http 暴露 im-service 的 HTTP/WebSocket 传输层。
|
||
package http
|
||
|
||
import (
|
||
"context"
|
||
"log"
|
||
"net/http"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gorilla/websocket"
|
||
imservice "hyapp/services/im-service/internal/service"
|
||
)
|
||
|
||
const (
|
||
serviceRestartReason = "service_restart"
|
||
closeWriteTimeout = 2 * time.Second
|
||
)
|
||
|
||
// clientFrame 是客户端到 im-service 的 WS JSON 帧契约。
|
||
type clientFrame struct {
|
||
// Type 决定帧路由，固定支持 auth、subscribe_room、unsubscribe_room、room_chat_send、peer_send、ack、pull_history。
|
||
Type string `json:"type"`
|
||
// Token 只在首帧 auth 中使用，鉴权通过后不再信任客户端自报身份。
|
||
Token string `json:"token,omitempty"`
|
||
// RoomID 用于房间订阅、公屏发送和历史拉取。
|
||
RoomID string `json:"room_id,omitempty"`
|
||
// ToUserID 用于单聊发送，FromUserID 由当前 session 推导。
|
||
ToUserID int64 `json:"to_user_id,omitempty"`
|
||
// Content 是 v1 文本消息体。
|
||
Content string `json:"content,omitempty"`
|
||
// MessageID 用于 ACK 指定待确认消息。
|
||
MessageID string `json:"message_id,omitempty"`
|
||
// Limit 用于 pull_history 裁剪房间最近窗口。
|
||
Limit int `json:"limit,omitempty"`
|
||
}
|
||
|
||
// WSHandler 负责承接客户端到 im-service 的真实 WebSocket 长连接。
|
||
type WSHandler struct {
|
||
// upgrader 执行 HTTP 到 WebSocket 的协议升级。
|
||
upgrader websocket.Upgrader
|
||
// svc 是连接、消息和房间频道逻辑的领域入口。
|
||
svc *imservice.Service
|
||
// mu 保护 draining 和 conns，保证进入 drain 后不会再登记新 WS 连接。
|
||
mu sync.Mutex
|
||
// draining 表示进程正在下线，新的 upgrade 必须被拒绝。
|
||
draining bool
|
||
// conns 保存当前本节点持有的真实 WS 连接，用于下线时发送 1012 close frame。
|
||
conns map[*trackedConn]struct{}
|
||
// wg 等待已登记连接的 handler 退出，避免下线时只关闭 listener 不等待 WS 清理。
|
||
wg sync.WaitGroup
|
||
}
|
||
|
||
// trackedConn 是传输层持有的真实 WS 连接及其日志身份。
|
||
type trackedConn struct {
|
||
// conn 是 gorilla websocket 的真实连接，close frame 和强制关闭都落到它上面。
|
||
conn *websocket.Conn
|
||
// mu 保护 connID/userID，身份会在 auth 成功后补齐。
|
||
mu sync.Mutex
|
||
// connID 是 im-service session id；auth 前连接使用 pending。
|
||
connID string
|
||
// userID 是 auth 后从 session manager 读出的真实用户 ID。
|
||
userID int64
|
||
// closeOnce 保证强制关闭同一底层连接最多执行一次。
|
||
closeOnce sync.Once
|
||
}
|
||
|
||
// NewWSHandler 创建一个 WS 入口处理器。
|
||
func NewWSHandler(svc *imservice.Service) *WSHandler {
|
||
return &WSHandler{
|
||
upgrader: websocket.Upgrader{
|
||
// 当前 v1 不在 WS 层做 Origin 收敛，外层网关或 CDN 可先承担来源控制。
|
||
CheckOrigin: func(*http.Request) bool { return true },
|
||
},
|
||
svc: svc,
|
||
conns: make(map[*trackedConn]struct{}),
|
||
}
|
||
}
|
||
|
||
// ServeHTTP 先完成 auth，再处理 subscribe/chat/peer/ack/pull_history 等帧。
|
||
func (h *WSHandler) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
|
||
if h.IsDraining() {
|
||
// drain 后新连接不能再进入 upgrade，客户端会按服务重启语义重新连接其他实例。
|
||
http.Error(writer, serviceRestartReason, http.StatusServiceUnavailable)
|
||
return
|
||
}
|
||
|
||
conn, err := h.upgrader.Upgrade(writer, request, nil)
|
||
if err != nil {
|
||
// 升级失败说明连接尚未进入 im-service session，不需要额外清理。
|
||
return
|
||
}
|
||
|
||
tracked, accepted := h.register(conn)
|
||
if !accepted {
|
||
// upgrade 和 drain 可能并发发生；已升级但未登记的连接也要收到稳定关闭语义。
|
||
writeServiceRestartClose(conn, "pending", 0)
|
||
_ = conn.Close()
|
||
return
|
||
}
|
||
defer h.unregister(tracked)
|
||
defer tracked.close()
|
||
|
||
if h.IsDraining() {
|
||
// 已登记连接在 auth 前遇到 drain，直接通知重连，不进入 session manager。
|
||
tracked.writeServiceRestartClose()
|
||
return
|
||
}
|
||
|
||
var first clientFrame
|
||
if err := conn.ReadJSON(&first); err != nil || first.Type != "auth" {
|
||
// 首帧必须是 auth，避免未鉴权连接先发送业务帧占用服务资源。
|
||
return
|
||
}
|
||
if h.IsDraining() {
|
||
// auth 读取和 drain 可以并发，drain 后不再创建新的 session。
|
||
tracked.writeServiceRestartClose()
|
||
return
|
||
}
|
||
|
||
connID, out, err := h.svc.Connect(first.Token)
|
||
if err != nil {
|
||
// token 无效时不下发错误细节，连接直接关闭，减少鉴权枚举面。
|
||
return
|
||
}
|
||
if userID, exists := h.svc.ConnectionUserID(connID); exists {
|
||
// close frame 日志需要真实连接身份，但身份必须来自已注册 session。
|
||
tracked.setIdentity(connID, userID)
|
||
}
|
||
|
||
// 连接退出时必须从 session manager 和 sink 表中同时移除。
|
||
defer h.svc.Disconnect(connID)
|
||
|
||
done := make(chan struct{})
|
||
|
||
go func() {
|
||
defer close(done)
|
||
|
||
// 写协程只消费 service 下行队列；业务 goroutine 不直接持有 websocket.Conn。
|
||
for frame := range out {
|
||
payload, err := imservice.MarshalFrame(frame)
|
||
if err != nil {
|
||
// 下行帧无法编码说明服务端数据异常，关闭连接让客户端重连恢复。
|
||
return
|
||
}
|
||
|
||
if err := conn.WriteMessage(websocket.TextMessage, payload); err != nil {
|
||
// 写失败通常代表客户端断线，读循环会随连接关闭退出。
|
||
return
|
||
}
|
||
}
|
||
}()
|
||
|
||
for {
|
||
var frame clientFrame
|
||
if err := conn.ReadJSON(&frame); err != nil {
|
||
// 读失败统一视为连接结束，defer Disconnect 会清理订阅和下行通道。
|
||
return
|
||
}
|
||
if h.IsDraining() {
|
||
// 下线期间不继续接受新的 subscribe_room 或发送帧，避免旧节点继续扩展连接态。
|
||
tracked.writeServiceRestartClose()
|
||
return
|
||
}
|
||
|
||
switch frame.Type {
|
||
case "subscribe_room":
|
||
// 订阅前会同步调用 room-service VerifyRoomPresence，防止越权收房间消息。
|
||
if err := h.svc.SubscribeRoom(request.Context(), connID, frame.RoomID); err != nil {
|
||
return
|
||
}
|
||
case "unsubscribe_room":
|
||
// 取消订阅是纯连接态操作，不改变 room-service presence。
|
||
h.svc.UnsubscribeRoom(connID, frame.RoomID)
|
||
case "room_chat_send":
|
||
// 公屏消息先经 room-service CheckSpeakPermission，再由 im-service 分配 room_seq。
|
||
if _, err := h.svc.SendRoomChat(request.Context(), connID, frame.RoomID, frame.Content); err != nil {
|
||
return
|
||
}
|
||
case "peer_send":
|
||
// 单聊 FromUserID 来自当前 session，客户端只提供目标用户和内容。
|
||
if _, err := h.svc.SendPeerMessage(connID, frame.ToUserID, frame.Content); err != nil {
|
||
return
|
||
}
|
||
case "ack":
|
||
// ACK 只清理当前用户的 peer_message/system_notice 待确认项。
|
||
if err := h.svc.Ack(connID, frame.MessageID); err != nil {
|
||
return
|
||
}
|
||
case "pull_history":
|
||
// 房间历史只返回 im-service 保存的最近窗口，用于重进房补拉。
|
||
history, err := h.svc.PullRoomHistory(connID, frame.RoomID, frame.Limit)
|
||
if err != nil {
|
||
return
|
||
}
|
||
|
||
payload, err := imservice.MarshalFrame(imservice.OutgoingFrame{
|
||
Type: "history",
|
||
RoomID: frame.RoomID,
|
||
History: history,
|
||
})
|
||
if err != nil {
|
||
// 历史帧编码失败说明存储数据异常，关闭连接比发送半包更安全。
|
||
return
|
||
}
|
||
|
||
if err := conn.WriteMessage(websocket.TextMessage, payload); err != nil {
|
||
// 同步历史响应写失败时连接已不可用，退出后交给客户端重连。
|
||
return
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Shutdown 进入 WS drain，发送 1012/service_restart，并等待客户端主动断开。
|
||
func (h *WSHandler) Shutdown(ctx context.Context) int {
|
||
conns := h.beginDraining()
|
||
for _, conn := range conns {
|
||
conn.writeServiceRestartClose()
|
||
}
|
||
|
||
done := make(chan struct{})
|
||
go func() {
|
||
h.wg.Wait()
|
||
close(done)
|
||
}()
|
||
|
||
select {
|
||
case <-done:
|
||
return 0
|
||
case <-ctx.Done():
|
||
remaining := h.snapshotConns()
|
||
for _, conn := range remaining {
|
||
conn.close()
|
||
}
|
||
return len(remaining)
|
||
}
|
||
}
|
||
|
||
// StartDraining 只切换 WS 入口状态，用于 HTTP listener 关闭前先拒绝并发 upgrade。
|
||
func (h *WSHandler) StartDraining() {
|
||
_ = h.beginDraining()
|
||
}
|
||
|
||
// ForceClose 立即关闭所有本机 WS 连接，用于第二次信号或 shutdown timeout 后加速退出。
|
||
func (h *WSHandler) ForceClose() int {
|
||
conns := h.beginDraining()
|
||
for _, conn := range conns {
|
||
conn.close()
|
||
}
|
||
|
||
return len(conns)
|
||
}
|
||
|
||
// ActiveConnections 返回当前仍由本节点持有的 WS 连接数量。
|
||
func (h *WSHandler) ActiveConnections() int {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
|
||
return len(h.conns)
|
||
}
|
||
|
||
// IsDraining 返回当前 handler 是否已经停止接收新连接。
|
||
func (h *WSHandler) IsDraining() bool {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
|
||
return h.draining
|
||
}
|
||
|
||
func (h *WSHandler) register(conn *websocket.Conn) (*trackedConn, bool) {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
|
||
if h.draining {
|
||
return nil, false
|
||
}
|
||
|
||
tracked := &trackedConn{conn: conn}
|
||
h.wg.Add(1)
|
||
h.conns[tracked] = struct{}{}
|
||
|
||
return tracked, true
|
||
}
|
||
|
||
func (h *WSHandler) unregister(conn *trackedConn) {
|
||
h.mu.Lock()
|
||
delete(h.conns, conn)
|
||
h.mu.Unlock()
|
||
|
||
h.wg.Done()
|
||
}
|
||
|
||
func (h *WSHandler) beginDraining() []*trackedConn {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
|
||
h.draining = true
|
||
return copyConns(h.conns)
|
||
}
|
||
|
||
func (h *WSHandler) snapshotConns() []*trackedConn {
|
||
h.mu.Lock()
|
||
defer h.mu.Unlock()
|
||
|
||
return copyConns(h.conns)
|
||
}
|
||
|
||
func copyConns(values map[*trackedConn]struct{}) []*trackedConn {
|
||
conns := make([]*trackedConn, 0, len(values))
|
||
for conn := range values {
|
||
conns = append(conns, conn)
|
||
}
|
||
|
||
return conns
|
||
}
|
||
|
||
func (c *trackedConn) setIdentity(connID string, userID int64) {
|
||
c.mu.Lock()
|
||
defer c.mu.Unlock()
|
||
|
||
c.connID = connID
|
||
c.userID = userID
|
||
}
|
||
|
||
func (c *trackedConn) identity() (string, int64) {
|
||
c.mu.Lock()
|
||
defer c.mu.Unlock()
|
||
|
||
connID := c.connID
|
||
if connID == "" {
|
||
connID = "pending"
|
||
}
|
||
|
||
return connID, c.userID
|
||
}
|
||
|
||
func (c *trackedConn) writeServiceRestartClose() {
|
||
connID, userID := c.identity()
|
||
writeServiceRestartClose(c.conn, connID, userID)
|
||
}
|
||
|
||
func (c *trackedConn) close() {
|
||
c.closeOnce.Do(func() {
|
||
_ = c.conn.Close()
|
||
})
|
||
}
|
||
|
||
func writeServiceRestartClose(conn *websocket.Conn, connID string, userID int64) {
|
||
err := conn.WriteControl(
|
||
websocket.CloseMessage,
|
||
websocket.FormatCloseMessage(websocket.CloseServiceRestart, serviceRestartReason),
|
||
time.Now().Add(closeWriteTimeout),
|
||
)
|
||
if err != nil {
|
||
log.Printf("service=im-service state=draining conn_id=%s user_id=%d close_code=%d close_reason=%s result=close_frame_failed error=%q", connID, userID, websocket.CloseServiceRestart, serviceRestartReason, err.Error())
|
||
return
|
||
}
|
||
|
||
log.Printf("service=im-service state=draining conn_id=%s user_id=%d close_code=%d close_reason=%s result=close_frame_sent", connID, userID, websocket.CloseServiceRestart, serviceRestartReason)
|
||
}
|