base-golang/pkg/sse/server.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
	})
}