SSE 推送替代轮训

2026-05-27 06:18:55 +08:00
parent 1458fc279c
commit a002987784
9 changed files with 493 additions and 42 deletions
@@ -0,0 +1,87 @@
+package chathub
+
+import (
+	"fmt"
+	"time"
+
+	"hfb_sys/backend/internal/middleware"
+	"hfb_sys/backend/pkg/response"
+
+	"github.com/gin-gonic/gin"
+)
+
+const heartbeatInterval = 30 * time.Second
+
+type Handler struct {
+	hub *Hub
+}
+
+func NewHandler(hub *Hub) *Handler {
+	return &Handler{hub: hub}
+}
+
+// UserEvents 处理用户端 SSE 连接: GET /api/chats/events
+func (h *Handler) UserEvents(c *gin.Context) {
+	value, ok := c.Get(middleware.ContextUserID)
+	if !ok {
+		response.Unauthorized(c, "缺少用户上下文")
+		return
+	}
+	userID, ok := value.(uint64)
+	if !ok {
+		response.Unauthorized(c, "用户上下文无效")
+		return
+	}
+	h.serveSSE(c, "user", userID)
+}
+
+// AdminEvents 处理管理端 SSE 连接: GET /api/admin/chats/events
+func (h *Handler) AdminEvents(c *gin.Context) {
+	value, ok := c.Get(middleware.ContextAdminID)
+	if !ok {
+		response.Unauthorized(c, "缺少管理员上下文")
+		return
+	}
+	adminID, ok := value.(uint64)
+	if !ok {
+		response.Unauthorized(c, "管理员上下文无效")
+		return
+	}
+	h.serveSSE(c, "admin", adminID)
+}
+
+func (h *Handler) serveSSE(c *gin.Context, pType string, pID uint64) {
+	ch := h.hub.Subscribe(pType, pID)
+	defer h.hub.Unsubscribe(pType, pID, ch)
+
+	c.Header("Content-Type", "text/event-stream")
+	c.Header("Cache-Control", "no-cache")
+	c.Header("Connection", "keep-alive")
+	c.Header("X-Accel-Buffering", "no")
+
+	// 发送初始连接确认
+	fmt.Fprintf(c.Writer, "event: connected\ndata: {\"ok\":true}\n\n")
+	c.Writer.Flush()
+
+	heartbeat := time.NewTicker(heartbeatInterval)
+	defer heartbeat.Stop()
+
+	clientGone := c.Request.Context().Done()
+
+	for {
+		select {
+		case <-clientGone:
+			return
+		case <-heartbeat.C:
+			fmt.Fprintf(c.Writer, ":heartbeat\n\n")
+			c.Writer.Flush()
+		case event, ok := <-ch:
+			if !ok {
+				return
+			}
+			data := MarshalEvent(event)
+			fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", event.Type, data)
+			c.Writer.Flush()
+		}
+	}
+}
@@ -0,0 +1,153 @@
+package chathub
+
+import (
+	"encoding/json"
+	"sync"
+
+	"gorm.io/gorm"
+)
+
+// ChatEvent 是通过 SSE 推送给客户端的事件。
+type ChatEvent struct {
+	Type           string       `json:"type"` // "new_message" | "conversation_updated"
+	ConversationID uint64       `json:"conversation_id"`
+	Message        *MessageData `json:"message,omitempty"`
+}
+
+// MessageData 是事件中携带的消息数据，与 chat.MessageDTO 对齐。
+type MessageData struct {
+	ID             uint64   `json:"id"`
+	ConversationID uint64   `json:"conversation_id"`
+	SenderType     string   `json:"sender_type"`
+	SenderID       uint64   `json:"sender_id"`
+	SenderRole     string   `json:"sender_role"`
+	SenderName     string   `json:"sender_name"`
+	ContentType    string   `json:"content_type"`
+	Content        string   `json:"content"`
+	AttachmentURLS []string `json:"attachment_urls"`
+	CreatedAt      string   `json:"created_at"`
+}
+
+// principal 标识一个连接方。
+type principal struct {
+	Type string // "user" or "admin"
+	ID   uint64
+}
+
+// Hub 管理所有 SSE 连接。
+type Hub struct {
+	mu      sync.RWMutex
+	clients map[principal]map[chan *ChatEvent]struct{}
+	db      *gorm.DB
+}
+
+// NewHub 创建 Hub 实例。
+func NewHub(db *gorm.DB) *Hub {
+	return &Hub{
+		clients: make(map[principal]map[chan *ChatEvent]struct{}),
+		db:      db,
+	}
+}
+
+// Subscribe 注册一个 SSE 连接，返回事件 channel。
+func (h *Hub) Subscribe(pType string, pID uint64) <-chan *ChatEvent {
+	ch := make(chan *ChatEvent, 16)
+	key := principal{Type: pType, ID: pID}
+	h.mu.Lock()
+	if h.clients[key] == nil {
+		h.clients[key] = make(map[chan *ChatEvent]struct{})
+	}
+	h.clients[key][ch] = struct{}{}
+	h.mu.Unlock()
+	return ch
+}
+
+// Unsubscribe 注销 SSE 连接。
+func (h *Hub) Unsubscribe(pType string, pID uint64, ch <-chan *ChatEvent) {
+	key := principal{Type: pType, ID: pID}
+	h.mu.Lock()
+	if clients, ok := h.clients[key]; ok {
+		// 找到对应的发送 channel 并删除
+		for sendCh := range clients {
+			if sendCh == ch {
+				delete(clients, sendCh)
+				close(sendCh)
+				break
+			}
+		}
+		if len(clients) == 0 {
+			delete(h.clients, key)
+		}
+	}
+	h.mu.Unlock()
+}
+
+// NotifyConversation 查询会话参与者，向所有在线参与者推送事件。
+func (h *Hub) NotifyConversation(conversationID uint64, event *ChatEvent) {
+	if h.db == nil {
+		return
+	}
+	// 查询会话参与者
+	type participantRow struct {
+		ParticipantType string
+		ParticipantID   uint64
+	}
+	var rows []participantRow
+	h.db.Table("chat_participants").
+		Select("participant_type, participant_id").
+		Where("conversation_id = ?", conversationID).
+		Find(&rows)
+
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	seen := make(map[principal]struct{})
+	for _, row := range rows {
+		key := principal{Type: row.ParticipantType, ID: row.ParticipantID}
+		if _, ok := seen[key]; ok {
+			continue
+		}
+		seen[key] = struct{}{}
+		if clients, ok := h.clients[key]; ok {
+			for ch := range clients {
+				select {
+				case ch <- event:
+				default:
+					// channel 满了，丢弃事件避免阻塞
+				}
+			}
+		}
+	}
+}
+
+// NotifyUser 直接向指定用户/Admin 推送事件。
+func (h *Hub) NotifyUser(pType string, pID uint64, event *ChatEvent) {
+	key := principal{Type: pType, ID: pID}
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	if clients, ok := h.clients[key]; ok {
+		for ch := range clients {
+			select {
+			case ch <- event:
+			default:
+			}
+		}
+	}
+}
+
+// MarshalEvent 将事件序列化为 SSE data 行。
+func MarshalEvent(event *ChatEvent) string {
+	raw, _ := json.Marshal(event)
+	return string(raw)
+}
+
+// OnlineCount 返回当前在线连接数。
+func (h *Hub) OnlineCount() int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	count := 0
+	for _, clients := range h.clients {
+		count += len(clients)
+	}
+	return count
+}