websocket.go 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. package notification
  2. import (
  3. "sync"
  4. "github.com/0xJacky/Nginx-UI/model"
  5. )
  6. // WebSocketNotificationManager manages WebSocket notification subscriptions
  7. type WebSocketNotificationManager struct {
  8. subscribers map[chan *model.Notification]struct{}
  9. mutex sync.RWMutex
  10. }
  11. var (
  12. wsManager *WebSocketNotificationManager
  13. wsManagerOnce sync.Once
  14. )
  15. // GetWebSocketManager returns the singleton WebSocket notification manager
  16. func GetWebSocketManager() *WebSocketNotificationManager {
  17. wsManagerOnce.Do(func() {
  18. wsManager = &WebSocketNotificationManager{
  19. subscribers: make(map[chan *model.Notification]struct{}),
  20. }
  21. })
  22. return wsManager
  23. }
  24. // Subscribe adds a channel to receive notifications
  25. func (m *WebSocketNotificationManager) Subscribe(ch chan *model.Notification) {
  26. m.mutex.Lock()
  27. defer m.mutex.Unlock()
  28. m.subscribers[ch] = struct{}{}
  29. }
  30. // Unsubscribe removes a channel from receiving notifications
  31. func (m *WebSocketNotificationManager) Unsubscribe(ch chan *model.Notification) {
  32. m.mutex.Lock()
  33. defer m.mutex.Unlock()
  34. delete(m.subscribers, ch)
  35. close(ch)
  36. }
  37. // Broadcast sends a notification to all subscribers
  38. func (m *WebSocketNotificationManager) Broadcast(data *model.Notification) {
  39. m.mutex.RLock()
  40. defer m.mutex.RUnlock()
  41. for ch := range m.subscribers {
  42. select {
  43. case ch <- data:
  44. default:
  45. // Skip if channel buffer is full
  46. }
  47. }
  48. }
  49. // BroadcastToWebSocket is a convenience function to broadcast notifications
  50. func BroadcastToWebSocket(data *model.Notification) {
  51. GetWebSocketManager().Broadcast(data)
  52. }