优化结构
This commit is contained in:
Yvan
@@ -0,0 +1,260 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
melody "github.com/olahol/melody"
|
||||
"kra/app/system/internal/conf"
|
||||
platformws "kra/pkg/websocket"
|
||||
)
|
||||
|
||||
// Server is the system-owned WebSocket endpoint. Business modules can attach
|
||||
// handlers and publish messages without depending on Gin or Melody directly.
|
||||
type Server struct {
|
||||
mu sync.RWMutex
|
||||
current *platformws.Server
|
||||
path string
|
||||
messageHandlers []func(*melody.Session, []byte)
|
||||
binaryHandlers []func(*melody.Session, []byte)
|
||||
connectHandlers []func(*melody.Session)
|
||||
disconnectHandlers []func(*melody.Session)
|
||||
closed bool
|
||||
}
|
||||
|
||||
func New(runtime *conf.Runtime) (*Server, func(), error) {
|
||||
s := &Server{}
|
||||
s.Replace(runtime)
|
||||
var unsubscribe func()
|
||||
if runtime != nil {
|
||||
unsubscribe = runtime.Subscribe(func(_ *conf.Data, _ *conf.AdminBackend) { s.Replace(runtime) })
|
||||
}
|
||||
return s, func() {
|
||||
if unsubscribe != nil {
|
||||
unsubscribe()
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.closed = true
|
||||
current := s.current
|
||||
s.current = nil
|
||||
s.mu.Unlock()
|
||||
if current != nil {
|
||||
_ = current.Close()
|
||||
}
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Server) Replace(runtime *conf.Runtime) {
|
||||
var config *conf.AdminBackend_WebSocket
|
||||
if runtime != nil && runtime.Admin() != nil {
|
||||
config = runtime.Admin().Websocket
|
||||
}
|
||||
if config == nil {
|
||||
config = &conf.AdminBackend_WebSocket{}
|
||||
}
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
previous := s.current
|
||||
if !config.Enabled {
|
||||
s.current = nil
|
||||
s.path = ""
|
||||
s.mu.Unlock()
|
||||
if previous != nil {
|
||||
_ = previous.Close()
|
||||
}
|
||||
return
|
||||
}
|
||||
next := platformws.New(platformws.Config{
|
||||
WriteWait: duration(config.WriteWait, 10*time.Second),
|
||||
PongWait: duration(config.PongWait, 60*time.Second),
|
||||
PingPeriod: duration(config.PingPeriod, 54*time.Second),
|
||||
MaxMessageSize: config.MaxMessageSize,
|
||||
MessageBufferSize: int(config.MessageBufferSize),
|
||||
ConcurrentMessageHandling: config.ConcurrentMessageHandling,
|
||||
AllowOrigins: config.AllowOrigins,
|
||||
})
|
||||
for _, handler := range s.messageHandlers {
|
||||
next.OnMessage(handler)
|
||||
}
|
||||
for _, handler := range s.binaryHandlers {
|
||||
next.OnBinaryMessage(handler)
|
||||
}
|
||||
for _, handler := range s.connectHandlers {
|
||||
next.OnConnect(handler)
|
||||
}
|
||||
for _, handler := range s.disconnectHandlers {
|
||||
next.OnDisconnect(handler)
|
||||
}
|
||||
s.path = strings.TrimSpace(config.Path)
|
||||
if s.path == "" {
|
||||
s.path = "/ws"
|
||||
} else if !strings.HasPrefix(s.path, "/") {
|
||||
s.path = "/" + s.path
|
||||
}
|
||||
s.current = next
|
||||
s.mu.Unlock()
|
||||
if previous != nil {
|
||||
_ = previous.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func duration(value interface{ AsDuration() time.Duration }, fallback time.Duration) time.Duration {
|
||||
if value == nil {
|
||||
return fallback
|
||||
}
|
||||
if result := value.AsDuration(); result > 0 {
|
||||
return result
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func (s *Server) Enabled() bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.current != nil
|
||||
}
|
||||
func (s *Server) Path() string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.path
|
||||
}
|
||||
func (s *Server) HandleRequest(w http.ResponseWriter, r *http.Request) error {
|
||||
if s == nil {
|
||||
return errors.New("websocket server is disabled")
|
||||
}
|
||||
s.mu.RLock()
|
||||
current := s.current
|
||||
s.mu.RUnlock()
|
||||
if current == nil {
|
||||
return errors.New("websocket server is disabled")
|
||||
}
|
||||
return current.HandleRequest(w, r)
|
||||
}
|
||||
func (s *Server) HandleRequestWithKeys(w http.ResponseWriter, r *http.Request, keys map[string]any) error {
|
||||
current, err := s.active()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return current.HandleRequestWithKeys(w, r, keys)
|
||||
}
|
||||
func (s *Server) Broadcast(message []byte) error {
|
||||
if s == nil {
|
||||
return errors.New("websocket server is disabled")
|
||||
}
|
||||
s.mu.RLock()
|
||||
current := s.current
|
||||
s.mu.RUnlock()
|
||||
if current == nil {
|
||||
return errors.New("websocket server is disabled")
|
||||
}
|
||||
return current.Broadcast(message)
|
||||
}
|
||||
func (s *Server) BroadcastBinary(message []byte) error {
|
||||
if s == nil {
|
||||
return errors.New("websocket server is disabled")
|
||||
}
|
||||
s.mu.RLock()
|
||||
current := s.current
|
||||
s.mu.RUnlock()
|
||||
if current == nil {
|
||||
return errors.New("websocket server is disabled")
|
||||
}
|
||||
return current.BroadcastBinary(message)
|
||||
}
|
||||
func (s *Server) Sessions() ([]*melody.Session, error) {
|
||||
current, err := s.active()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return current.Sessions()
|
||||
}
|
||||
func (s *Server) Len() int {
|
||||
current, err := s.active()
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return current.Len()
|
||||
}
|
||||
func (s *Server) Send(session *melody.Session, message []byte) error {
|
||||
current, err := s.active()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return current.Send(session, message)
|
||||
}
|
||||
func (s *Server) SendBinary(session *melody.Session, message []byte) error {
|
||||
current, err := s.active()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return current.SendBinary(session, message)
|
||||
}
|
||||
func (s *Server) active() (*platformws.Server, error) {
|
||||
if s == nil {
|
||||
return nil, errors.New("websocket server is disabled")
|
||||
}
|
||||
s.mu.RLock()
|
||||
current := s.current
|
||||
s.mu.RUnlock()
|
||||
if current == nil {
|
||||
return nil, errors.New("websocket server is disabled")
|
||||
}
|
||||
return current, nil
|
||||
}
|
||||
func (s *Server) OnMessage(handler func(*melody.Session, []byte)) {
|
||||
if s == nil || handler == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.messageHandlers = append(s.messageHandlers, handler)
|
||||
current := s.current
|
||||
if current != nil {
|
||||
current.OnMessage(handler)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
func (s *Server) OnBinaryMessage(handler func(*melody.Session, []byte)) {
|
||||
if s == nil || handler == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.binaryHandlers = append(s.binaryHandlers, handler)
|
||||
if s.current != nil {
|
||||
s.current.OnBinaryMessage(handler)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
func (s *Server) OnConnect(handler func(*melody.Session)) {
|
||||
if s == nil || handler == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.connectHandlers = append(s.connectHandlers, handler)
|
||||
if s.current != nil {
|
||||
s.current.OnConnect(handler)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
func (s *Server) OnDisconnect(handler func(*melody.Session)) {
|
||||
if s == nil || handler == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.disconnectHandlers = append(s.disconnectHandlers, handler)
|
||||
if s.current != nil {
|
||||
s.current.OnDisconnect(handler)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
Reference in New Issue
Block a user