Merge pull request '[🚀] Ticker' (#13) from dev into main

Reviewed-on: #13

pkg/sse/server.go

@@ -1,19 +1,22 @@
 package sse
 
 import (
+	"fmt"
 	"gitea.bvbej.com/bvbej/base-golang/pkg/mux"
+	"gitea.bvbej.com/bvbej/base-golang/pkg/ticker"
 	"github.com/gin-gonic/gin"
 	"io"
 	"net/http"
 	"sync"
 	"sync/atomic"
+	"time"
 )
 
 var _ Server = (*event)(nil)
 
 type Server interface {
-	GinHandle(ctx *gin.Context, user any)
 	HandlerFunc() mux.HandlerFunc
+	GinHandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc
 	Push(user any, name, msg string) bool
 	Broadcast(name, msg string)
 }
@@ -31,9 +34,9 @@ type msgChan struct {
 type event struct {
 	SessionList sync.Map
 	Count       atomic.Int32
-
 	Register    chan clientChan
 	Unregister  chan any
+	Ticker      ticker.Ticker
 }
 
 func NewServer() Server {
@@ -42,8 +45,12 @@ func NewServer() Server {
 		Count:       atomic.Int32{},
 		Register:    make(chan clientChan),
 		Unregister:  make(chan any),
+		Ticker:      ticker.New(time.Second * 3),
 	}
 
+	go e.Ticker.Process(func() {
+		e.Broadcast("ping", fmt.Sprintf("%d", time.Now().Unix()))
+	})
 	go e.listen()
 
 	return e
@@ -67,11 +74,47 @@ func (stream *event) listen() {
 	}
 }
 
-func (stream *event) GinHandle(ctx *gin.Context, user any) {
-	if user == nil {
-		ctx.AbortWithStatus(http.StatusUnauthorized)
+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")
+
+		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 {
+	return func(c *gin.Context) {
+		user, err := auth(c)
+		if err != nil {
+			c.AbortWithStatus(http.StatusBadRequest)
+			return
+		}
+
 		e := make(chan msgChan)
 		client := clientChan{
 			User: user,
@@ -82,25 +125,18 @@ func (stream *event) GinHandle(ctx *gin.Context, user any) {
 			stream.Unregister <- user
 		}()
 
-	ctx.Writer.Header().Set("Content-Type", "text/event-stream")
-	ctx.Writer.Header().Set("Cache-Control", "no-cache")
-	ctx.Writer.Header().Set("Connection", "keep-alive")
-	ctx.Writer.Header().Set("Transfer-Encoding", "chunked")
+		c.Writer.Header().Set("Content-Type", "text/event-stream")
+		c.Writer.Header().Set("Cache-Control", "no-cache")
+		c.Writer.Header().Set("Connection", "keep-alive")
+		c.Writer.Header().Set("Transfer-Encoding", "chunked")
 
-	ctx.Stream(func(w io.Writer) bool {
+		c.Stream(func(w io.Writer) bool {
 			if msg, ok := <-e; ok {
-			ctx.SSEvent(msg.Name, msg.Message)
+				c.SSEvent(msg.Name, msg.Message)
 				return true
 			}
 			return false
 		})
-
-	ctx.Next()
-}
-
-func (stream *event) HandlerFunc() mux.HandlerFunc {
-	return func(c mux.Context) {
-		stream.GinHandle(c.Context(), c.Auth())
 	}
 }