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,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
+}