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()
+		}
+	}
+}