package sse import ( "bufio" "context" "io" "net/http" "strings" "sync" "sync/atomic" "time" ) // Event represents a Server-Sent Event type Event struct { ID string Event string Data string Retry int } // ConnectionMeta holds metadata about an SSE connection type ConnectionMeta struct { EventChan <-chan Event ErrorChan <-chan error lastPingAt atomic.Value // stores time.Time cancel context.CancelFunc } // LastPing returns the last ping time func (c *ConnectionMeta) LastPing() time.Time { if t, ok := c.lastPingAt.Load().(time.Time); ok { return t } return time.Time{} } // UpdateLastPing atomically updates the last ping time func (c *ConnectionMeta) UpdateLastPing() { c.lastPingAt.Store(time.Now()) } // Cancel terminates the connection func (c *ConnectionMeta) Cancel() { if c.cancel != nil { c.cancel() } } // ConnectionRegistry manages SSE connections with high performance type ConnectionRegistry struct { mu sync.RWMutex conn *ConnectionMeta monitorOnce sync.Once monitorStop chan struct{} onPing func() // Callback for ping events } // NewConnectionRegistry creates a new high-performance connection registry func NewConnectionRegistry(ctx context.Context) *ConnectionRegistry { r := &ConnectionRegistry{ monitorStop: make(chan struct{}), } // Configure ping handler r.onPing = func() { r.UpdateLastPing() } return r } // Subscribe provides SSE events, creating a connection if needed func (r *ConnectionRegistry) Subscribe(request func() (*http.Response, error)) (<-chan Event, <-chan error, error) { // Fast path: check if connection exists if conn := r.getConnection(); conn != nil { return conn.EventChan, conn.ErrorChan, nil } // Slow path: create new connection under lock r.mu.Lock() defer r.mu.Unlock() // Double-check after acquiring lock if r.conn != nil { return r.conn.EventChan, r.conn.ErrorChan, nil } res, err := request() if err != nil { return nil, nil, err } conn, err := r.createStream(res) if err != nil { return nil, nil, err } r.conn = conn // Start monitor once r.monitorOnce.Do(func() { go r.monitorConnections() }) return conn.EventChan, conn.ErrorChan, nil } // Get retrieves the current connection func (r *ConnectionRegistry) Get() (*ConnectionMeta, bool) { conn := r.getConnection() return conn, conn != nil } // IsConnected checks if there's an active connection func (r *ConnectionRegistry) IsConnected() bool { return r.getConnection() != nil } // UpdateLastPing updates the last ping time for the current connection func (r *ConnectionRegistry) UpdateLastPing() { if conn := r.getConnection(); conn != nil { conn.UpdateLastPing() } } // Close gracefully shuts down the registry func (r *ConnectionRegistry) Close() { // Stop monitor first select { case <-r.monitorStop: // Already closed default: close(r.monitorStop) } // Close connection r.mu.Lock() if r.conn != nil { r.conn.Cancel() r.conn = nil } r.mu.Unlock() } // getConnection returns the current connection without locking func (r *ConnectionRegistry) getConnection() *ConnectionMeta { r.mu.RLock() conn := r.conn r.mu.RUnlock() return conn } func (r *ConnectionRegistry) createStream(res *http.Response) (*ConnectionMeta, error) { ctx, cancel := context.WithCancel(context.Background()) eventChan, errorChan, err := r.stream(ctx, res) if err != nil { cancel() return nil, err } meta := &ConnectionMeta{ EventChan: eventChan, ErrorChan: errorChan, cancel: cancel, } meta.UpdateLastPing() return meta, nil } // stream processes SSE data from an HTTP response func (r *ConnectionRegistry) stream(ctx context.Context, res *http.Response) (<-chan Event, <-chan error, error) { eventChan := make(chan Event, 10) errorChan := make(chan error, 1) go r.processStream(ctx, res.Body, eventChan, errorChan) return eventChan, errorChan, nil } // processStream reads and parses SSE events from the response body func (r *ConnectionRegistry) processStream(ctx context.Context, body io.ReadCloser, eventChan chan<- Event, errorChan chan<- error) { defer body.Close() defer close(eventChan) defer close(errorChan) scanner := bufio.NewScanner(body) var currentEvent Event var dataBuilder strings.Builder for scanner.Scan() { select { case <-ctx.Done(): return default: } line := scanner.Text() // Empty line indicates end of event if len(line) == 0 { if currentEvent.Data != "" || currentEvent.Event != "" { // Finalize data if dataBuilder.Len() > 0 { currentEvent.Data = dataBuilder.String() dataBuilder.Reset() } // Handle ping events if r.isPingEvent(currentEvent) { if r.onPing != nil { r.onPing() } } else { // Send non-ping events select { case eventChan <- currentEvent: case <-ctx.Done(): return } } // Reset for next event currentEvent = Event{} } continue } // Parse line efficiently r.parseLine(line, ¤tEvent, &dataBuilder) } if err := scanner.Err(); err != nil { select { case errorChan <- err: case <-ctx.Done(): } } } // parseLine efficiently parses SSE protocol lines func (r *ConnectionRegistry) parseLine(line string, event *Event, dataBuilder *strings.Builder) { colonIndex := strings.IndexByte(line, ':') if colonIndex == -1 { return // Invalid line format } field := line[:colonIndex] value := line[colonIndex+1:] // Trim leading space from value (SSE spec) if len(value) > 0 && value[0] == ' ' { value = value[1:] } switch field { case "data": if dataBuilder.Len() > 0 { dataBuilder.WriteByte('\n') } dataBuilder.WriteString(value) case "event": event.Event = value case "id": event.ID = value case "retry": // Parse retry value if needed // This could be used to configure reconnection delay case "": // Comment line, ignore } } // isPingEvent checks if an event is a ping/keepalive func (r *ConnectionRegistry) isPingEvent(event Event) bool { // Check for common ping patterns if event.Event == "ping" { return true } // Check for heartbeat data (common pattern is "1" or similar) if event.Event == "" && strings.TrimSpace(event.Data) == "1" { return true } return false } // monitorConnections checks connection health periodically func (r *ConnectionRegistry) monitorConnections() { const ( checkInterval = 30 * time.Second pingTimeout = 2 * time.Minute ) ticker := time.NewTicker(checkInterval) defer ticker.Stop() for { select { case <-r.monitorStop: return case <-ticker.C: r.checkConnectionHealth(pingTimeout) } } } // checkConnectionHealth verifies connection is still alive func (r *ConnectionRegistry) checkConnectionHealth(timeout time.Duration) { conn := r.getConnection() if conn == nil { return } if time.Since(conn.LastPing()) > timeout { // Connection is stale, close it r.mu.Lock() if r.conn == conn { // Verify it's still the same connection r.conn.Cancel() r.monitorStop <- struct{}{} r.conn = nil } r.mu.Unlock() } }