123 lines
2.3 KiB
Go
123 lines
2.3 KiB
Go
package sse
|
|
|
|
import (
|
|
"gitea.bvbej.com/bvbej/base-golang/pkg/mux"
|
|
"github.com/gin-gonic/gin"
|
|
"io"
|
|
"net/http"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
var _ Server = (*event)(nil)
|
|
|
|
type Server interface {
|
|
GinHandle(ctx *gin.Context, user any)
|
|
HandlerFunc() mux.HandlerFunc
|
|
Push(user any, name, msg string) bool
|
|
Broadcast(name, msg string)
|
|
}
|
|
|
|
type clientChan struct {
|
|
User any
|
|
Chan chan msgChan
|
|
}
|
|
|
|
type msgChan struct {
|
|
Name string
|
|
Message string
|
|
}
|
|
|
|
type event struct {
|
|
SessionList sync.Map
|
|
Count atomic.Int32
|
|
|
|
Register chan clientChan
|
|
Unregister chan any
|
|
}
|
|
|
|
func NewServer() Server {
|
|
e := &event{
|
|
SessionList: sync.Map{},
|
|
Count: atomic.Int32{},
|
|
Register: make(chan clientChan),
|
|
Unregister: make(chan any),
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (stream *event) GinHandle(ctx *gin.Context, user any) {
|
|
if user == nil {
|
|
ctx.AbortWithStatus(http.StatusUnauthorized)
|
|
return
|
|
}
|
|
e := make(chan msgChan)
|
|
client := clientChan{
|
|
User: user,
|
|
Chan: e,
|
|
}
|
|
stream.Register <- client
|
|
defer func() {
|
|
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")
|
|
|
|
ctx.Stream(func(w io.Writer) bool {
|
|
if msg, ok := <-e; ok {
|
|
ctx.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())
|
|
}
|
|
}
|
|
|
|
func (stream *event) Push(user any, name, msg string) bool {
|
|
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
|
|
})
|
|
}
|