package sse import ( "context" "encoding/json" "fmt" "log/slog" "sync" "time" "github.com/google/uuid" "golang.org/x/time/rate" ) // Event represents an SSE event type Event struct { ID string `json:"id"` Type string `json:"type"` InstanceName string `json:"instanceName"` Timestamp time.Time `json:"timestamp"` Data any `json:"data"` Metadata map[string]any `json:"metadata,omitempty"` } // Client represents an SSE client connection type Client struct { ID string InstanceName string Channel chan *Event Filters EventFilters Context context.Context Cancel context.CancelFunc } // EventFilters for client-specific filtering type EventFilters struct { EventTypes []string `json:"eventTypes,omitempty"` Namespaces []string `json:"namespaces,omitempty"` Apps []string `json:"apps,omitempty"` } // Manager manages all SSE connections type Manager struct { clients map[string]map[string]*Client // instanceName -> clientID -> Client unregister chan *Client broadcast chan *Event mu sync.RWMutex rateLimiters map[string]*rate.Limiter } // NewManager creates a new SSE manager func NewManager() *Manager { m := &Manager{ clients: make(map[string]map[string]*Client), unregister: make(chan *Client, 100), broadcast: make(chan *Event, 1000), rateLimiters: make(map[string]*rate.Limiter), } go m.run() return m } // run processes client unregistration and event broadcasting func (m *Manager) run() { for { select { case client := <-m.unregister: m.mu.Lock() if clients, ok := m.clients[client.InstanceName]; ok { delete(clients, client.ID) if len(clients) == 0 { delete(m.clients, client.InstanceName) } } close(client.Channel) m.mu.Unlock() slog.Info("client unregistered", "component", "sse", "client", client.ID) case event := <-m.broadcast: m.mu.RLock() // Send to clients registered for this specific instance instanceClients := m.clients[event.InstanceName] // Also send to clients registered with wildcard "*" to receive all events globalClients := m.clients["*"] m.mu.RUnlock() // Send to instance-specific clients for _, client := range instanceClients { if m.shouldSendToClient(event, client) { select { case client.Channel <- event: default: // Client channel full, skip slog.Info("client channel full, skipping event", "component", "sse", "client", client.ID) } } } // Send to global clients (registered with "*") for _, client := range globalClients { if m.shouldSendToClient(event, client) { select { case client.Channel <- event: default: // Client channel full, skip slog.Info("client channel full, skipping event", "component", "sse", "client", client.ID) } } } } } } // shouldSendToClient checks if event matches client filters func (m *Manager) shouldSendToClient(event *Event, client *Client) bool { // Rate limiting per client limiter := m.getRateLimiter(client.ID) if !limiter.Allow() { return false } // Check event type filter if len(client.Filters.EventTypes) > 0 { matched := false for _, eventType := range client.Filters.EventTypes { if event.Type == eventType { matched = true break } } if !matched { return false } } // Check namespace filter for k8s events if len(client.Filters.Namespaces) > 0 { if namespace, ok := event.Metadata["namespace"].(string); ok { matched := false for _, ns := range client.Filters.Namespaces { if namespace == ns { matched = true break } } if !matched { return false } } } // Check app filter if len(client.Filters.Apps) > 0 { if appName, ok := event.Metadata["app"].(string); ok { matched := false for _, app := range client.Filters.Apps { if appName == app { matched = true break } } if !matched { return false } } } return true } // getRateLimiter gets or creates a rate limiter for a client func (m *Manager) getRateLimiter(clientID string) *rate.Limiter { m.mu.Lock() defer m.mu.Unlock() if limiter, exists := m.rateLimiters[clientID]; exists { return limiter } // 100 events per second with burst of 200 limiter := rate.NewLimiter(100, 200) m.rateLimiters[clientID] = limiter return limiter } // RegisterClient registers a new SSE client func (m *Manager) RegisterClient(instanceName string, filters EventFilters) *Client { ctx, cancel := context.WithCancel(context.Background()) client := &Client{ ID: generateClientID(), InstanceName: instanceName, Channel: make(chan *Event, 100), Filters: filters, Context: ctx, Cancel: cancel, } m.mu.Lock() if m.clients[instanceName] == nil { m.clients[instanceName] = make(map[string]*Client) } m.clients[instanceName][client.ID] = client m.mu.Unlock() slog.Info("client registered", "component", "sse", "client", client.ID, "instance", instanceName) return client } // UnregisterClient removes a client func (m *Manager) UnregisterClient(client *Client) { client.Cancel() m.unregister <- client // Clean up rate limiter m.mu.Lock() delete(m.rateLimiters, client.ID) m.mu.Unlock() } // Broadcast sends an event to all matching clients func (m *Manager) Broadcast(event *Event) { event.ID = generateEventID() event.Timestamp = time.Now() select { case m.broadcast <- event: default: slog.Error("broadcast channel full, dropping event", "component", "sse", "event", event.ID, "type", event.Type, "instance", event.InstanceName) } } // GetConnectionCount returns the number of active connections func (m *Manager) GetConnectionCount() int { m.mu.RLock() defer m.mu.RUnlock() count := 0 for _, clients := range m.clients { count += len(clients) } return count } // GetInstanceConnectionCount returns the number of connections for a specific instance func (m *Manager) GetInstanceConnectionCount(instanceName string) int { m.mu.RLock() defer m.mu.RUnlock() if clients, ok := m.clients[instanceName]; ok { return len(clients) } return 0 } // Helper functions func generateClientID() string { return fmt.Sprintf("client-%s", uuid.New().String()[:8]) } func generateEventID() string { return fmt.Sprintf("event-%s", uuid.New().String()[:8]) } // JSON marshals the event to JSON func (e *Event) JSON() ([]byte, error) { return json.Marshal(e) }