base-golang/pkg/sse/server.go

165 lines
3.4 KiB
Go
Raw Normal View History

2024-07-23 10:23:43 +08:00
package sse
import (
2024-09-07 11:25:30 +08:00
"gitea.bvbej.com/bvbej/base-golang/pkg/mux"
2024-07-23 10:23:43 +08:00
"github.com/gin-gonic/gin"
"io"
"net/http"
"sync"
"sync/atomic"
2024-09-07 15:55:14 +08:00
"time"
2024-07-23 10:23:43 +08:00
)
var _ Server = (*event)(nil)
type Server interface {
2024-09-07 11:25:30 +08:00
HandlerFunc() mux.HandlerFunc
2024-09-07 15:13:38 +08:00
GinHandlerFunc(auth func(c *gin.Context) (string, error)) gin.HandlerFunc
2024-09-07 13:41:57 +08:00
Push(user any, name, msg string) bool
2024-07-23 10:23:43 +08:00
Broadcast(name, msg string)
2024-09-07 17:13:43 +08:00
ClientCount() int32
2024-07-23 10:23:43 +08:00
}
type clientChan struct {
2024-09-07 11:25:30 +08:00
User any
2024-07-23 10:23:43 +08:00
Chan chan msgChan
}
type msgChan struct {
Name string
Message string
}
type event struct {
SessionList sync.Map
Count atomic.Int32
2024-09-07 15:44:13 +08:00
Register chan clientChan
Unregister chan any
2024-07-23 10:23:43 +08:00
}
func NewServer() Server {
e := &event{
SessionList: sync.Map{},
Count: atomic.Int32{},
Register: make(chan clientChan),
2024-09-07 11:25:30 +08:00
Unregister: make(chan any),
2024-07-23 10:23:43 +08:00
}
go e.listen()
return e
}
func (stream *event) listen() {
for {
select {
case client := <-stream.Register:
stream.SessionList.Store(client.User, client.Chan)
stream.Count.Add(1)
case user := <-stream.Unregister:
value, ok := stream.SessionList.Load(user)
if ok {
event := value.(chan msgChan)
close(event)
stream.SessionList.Delete(user)
stream.Count.Add(-1)
}
}
}
}
2024-09-07 15:13:38 +08:00
func (stream *event) HandlerFunc() mux.HandlerFunc {
return func(c mux.Context) {
auth := c.Auth()
if auth == nil {
c.Context().AbortWithStatus(http.StatusBadRequest)
return
2024-07-23 10:23:43 +08:00
}
2024-09-07 15:13:38 +08:00
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")
2024-09-07 16:00:30 +08:00
time.AfterFunc(time.Second, func() {
e <- msgChan{Name: "message", Message: "success"}
})
2024-09-07 15:13:38 +08:00
c.Context().Stream(func(w io.Writer) bool {
if msg, ok := <-e; ok {
c.Context().SSEvent(msg.Name, msg.Message)
return true
}
return false
})
}
2024-09-07 11:25:30 +08:00
}
2024-07-23 10:23:43 +08:00
2024-09-07 15:13:38 +08:00
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,
Chan: e,
}
stream.Register <- client
defer func() {
stream.Unregister <- user
}()
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")
2024-09-07 15:55:14 +08:00
time.AfterFunc(time.Second, func() {
e <- msgChan{Name: "message", Message: "success"}
})
2024-09-07 15:13:38 +08:00
c.Stream(func(w io.Writer) bool {
if msg, ok := <-e; ok {
c.SSEvent(msg.Name, msg.Message)
return true
}
return false
})
2024-07-23 10:23:43 +08:00
}
}
2024-09-07 13:41:57 +08:00
func (stream *event) Push(user any, name, msg string) bool {
2024-07-23 10:23:43 +08:00
value, ok := stream.SessionList.Load(user)
if ok {
val := value.(chan msgChan)
val <- msgChan{Name: name, Message: msg}
}
return false
}
func (stream *event) Broadcast(name, msg string) {
stream.SessionList.Range(func(user, value any) bool {
val := value.(chan msgChan)
val <- msgChan{Name: name, Message: msg}
return true
})
}
2024-09-07 17:13:43 +08:00
func (stream *event) ClientCount() int32 {
return stream.Count.Load()
}