pkg/sse/server.go

@@ -1,27 +1,23 @@
 package sse
 
 import (
-	"gitea.bvbej.com/bvbej/base-golang/pkg/mux"
 	"github.com/gin-gonic/gin"
 	"io"
 	"net/http"
 	"sync"
 	"sync/atomic"
-	"time"
 )
 
 var _ Server = (*event)(nil)
 
 type Server interface {
-	HandlerFunc() mux.HandlerFunc
-	GinHandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc
-	Push(user any, name, msg string) bool
+	HandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc
+	Push(user, name, msg string) bool
 	Broadcast(name, msg string)
-	ClientCount() int32
 }
 
 type clientChan struct {
-	User any
+	User string
 	Chan chan msgChan
 }
 
@@ -33,8 +29,9 @@ type msgChan struct {
 type event struct {
 	SessionList sync.Map
 	Count       atomic.Int32
-	Register    chan clientChan
-	Unregister  chan any
+
+	Register   chan clientChan
+	Unregister chan string
 }
 
 func NewServer() Server {
@@ -42,7 +39,7 @@ func NewServer() Server {
 		SessionList: sync.Map{},
 		Count:       atomic.Int32{},
 		Register:    make(chan clientChan),
-		Unregister:  make(chan any),
+		Unregister:  make(chan string),
 	}
 
 	go e.listen()
@@ -68,44 +65,7 @@ func (stream *event) listen() {
 	}
 }
 
-func (stream *event) HandlerFunc() mux.HandlerFunc {
-	return func(c mux.Context) {
-		auth := c.Auth()
-		if auth == nil {
-			c.Context().AbortWithStatus(http.StatusBadRequest)
-			return
-		}
-
-		e := make(chan msgChan)
-		client := clientChan{
-			User: auth,
-			Chan: e,
-		}
-		stream.Register <- client
-		defer func() {
-			stream.Unregister <- auth
-		}()
-
-		c.Context().Writer.Header().Set("Content-Type", "text/event-stream")
-		c.Context().Writer.Header().Set("Cache-Control", "no-cache")
-		c.Context().Writer.Header().Set("Connection", "keep-alive")
-		c.Context().Writer.Header().Set("Transfer-Encoding", "chunked")
-
-		time.AfterFunc(time.Second, func() {
-			e <- msgChan{Name: "message", Message: "success"}
-		})
-
-		c.Context().Stream(func(w io.Writer) bool {
-			if msg, ok := <-e; ok {
-				c.Context().SSEvent(msg.Name, msg.Message)
-				return true
-			}
-			return false
-		})
-	}
-}
-
-func (stream *event) GinHandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc {
+func (stream *event) HandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		user, err := auth(c)
 		if err != nil {
@@ -128,10 +88,6 @@ func (stream *event) GinHandlerFunc(auth func(c *gin.Context) (string, error)) g
 		c.Writer.Header().Set("Connection", "keep-alive")
 		c.Writer.Header().Set("Transfer-Encoding", "chunked")
 
-		time.AfterFunc(time.Second, func() {
-			e <- msgChan{Name: "message", Message: "success"}
-		})
-
 		c.Stream(func(w io.Writer) bool {
 			if msg, ok := <-e; ok {
 				c.SSEvent(msg.Name, msg.Message)
@@ -139,10 +95,12 @@ func (stream *event) GinHandlerFunc(auth func(c *gin.Context) (string, error)) g
 			}
 			return false
 		})
+
+		c.Next()
 	}
 }
 
-func (stream *event) Push(user any, name, msg string) bool {
+func (stream *event) Push(user, name, msg string) bool {
 	value, ok := stream.SessionList.Load(user)
 	if ok {
 		val := value.(chan msgChan)
@@ -158,7 +116,3 @@ func (stream *event) Broadcast(name, msg string) {
 		return true
 	})
 }
-
-func (stream *event) ClientCount() int32 {
-	return stream.Count.Load()
-}