127 lines
3.5 KiB
Go
127 lines
3.5 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"log"
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// EventBroadcaster manages WebSocket connections and broadcasts events
|
|
type EventBroadcaster struct {
|
|
mu sync.RWMutex
|
|
connections map[*websocket.Conn]*sync.Mutex // Each connection has its own write mutex
|
|
maxConnections int
|
|
}
|
|
|
|
// NewEventBroadcaster creates a new event broadcaster
|
|
func NewEventBroadcaster() *EventBroadcaster {
|
|
return &EventBroadcaster{
|
|
connections: make(map[*websocket.Conn]*sync.Mutex),
|
|
maxConnections: 100, // Limit to 100 concurrent WebSocket connections
|
|
}
|
|
}
|
|
|
|
// AddConnection registers a new WebSocket connection
|
|
func (eb *EventBroadcaster) AddConnection(conn *websocket.Conn) error {
|
|
eb.mu.Lock()
|
|
defer eb.mu.Unlock()
|
|
|
|
// Check connection limit
|
|
if len(eb.connections) >= eb.maxConnections {
|
|
log.Printf("WebSocket connection limit reached (%d), rejecting new connection", eb.maxConnections)
|
|
return &websocket.CloseError{Code: websocket.ClosePolicyViolation, Text: "Connection limit reached"}
|
|
}
|
|
|
|
eb.connections[conn] = &sync.Mutex{} // Create a dedicated mutex for this connection
|
|
log.Printf("WebSocket connected to events. Total connections: %d", len(eb.connections))
|
|
return nil
|
|
}
|
|
|
|
// RemoveConnection unregisters a WebSocket connection
|
|
func (eb *EventBroadcaster) RemoveConnection(conn *websocket.Conn) {
|
|
eb.mu.Lock()
|
|
defer eb.mu.Unlock()
|
|
delete(eb.connections, conn)
|
|
log.Printf("WebSocket disconnected from events. Total connections: %d", len(eb.connections))
|
|
}
|
|
|
|
// Broadcast sends an event to all connected WebSocket clients
|
|
func (eb *EventBroadcaster) Broadcast(event DockerEvent) {
|
|
// Marshal event to JSON
|
|
data, err := json.Marshal(event)
|
|
if err != nil {
|
|
log.Printf("Error marshaling event: %v", err)
|
|
return
|
|
}
|
|
|
|
// Track dead connections
|
|
var deadConnections []*websocket.Conn
|
|
|
|
// Get snapshot of connections with their mutexes
|
|
eb.mu.RLock()
|
|
connMutexes := make(map[*websocket.Conn]*sync.Mutex, len(eb.connections))
|
|
for conn, mu := range eb.connections {
|
|
connMutexes[conn] = mu
|
|
}
|
|
eb.mu.RUnlock()
|
|
|
|
// Send to all connections (with per-connection write lock)
|
|
for conn, mu := range connMutexes {
|
|
mu.Lock()
|
|
err := conn.WriteMessage(websocket.TextMessage, data)
|
|
mu.Unlock()
|
|
|
|
if err != nil {
|
|
log.Printf("Error sending event to WebSocket: %v", err)
|
|
deadConnections = append(deadConnections, conn)
|
|
}
|
|
}
|
|
|
|
// Clean up dead connections
|
|
if len(deadConnections) > 0 {
|
|
// Remove from map first (fast, under lock)
|
|
eb.mu.Lock()
|
|
var connectionsToClose []*websocket.Conn
|
|
for _, conn := range deadConnections {
|
|
// Only delete if connection still exists in map
|
|
if _, exists := eb.connections[conn]; exists {
|
|
delete(eb.connections, conn)
|
|
connectionsToClose = append(connectionsToClose, conn)
|
|
}
|
|
}
|
|
eb.mu.Unlock()
|
|
|
|
// Close connections outside lock (slow, can block)
|
|
for _, conn := range connectionsToClose {
|
|
conn.Close()
|
|
}
|
|
}
|
|
}
|
|
|
|
// GetConnectionCount returns the number of active WebSocket connections
|
|
func (eb *EventBroadcaster) GetConnectionCount() int {
|
|
eb.mu.RLock()
|
|
defer eb.mu.RUnlock()
|
|
return len(eb.connections)
|
|
}
|
|
|
|
// CloseAll closes all WebSocket connections
|
|
func (eb *EventBroadcaster) CloseAll() {
|
|
eb.mu.Lock()
|
|
var connectionsToClose []*websocket.Conn
|
|
for conn := range eb.connections {
|
|
connectionsToClose = append(connectionsToClose, conn)
|
|
}
|
|
eb.connections = make(map[*websocket.Conn]*sync.Mutex)
|
|
eb.mu.Unlock()
|
|
|
|
// Close connections outside lock (can block on network I/O)
|
|
for _, conn := range connectionsToClose {
|
|
conn.Close()
|
|
}
|
|
|
|
log.Println("Closed all event WebSocket connections")
|
|
}
|