sultan/pkg/stream/hub.go

package stream

import (
	"errors"
	"fmt"
	"net/http"
	"slices"
	"sync"
	"time"
)

var (
	// ErrNodeOffline: Für den Knoten ist kein Stream offen.
	ErrNodeOffline = errors.New("stream: Knoten nicht verbunden")
	// ErrBufferFull: Der Knoten liest nicht schnell genug; das Event wurde verworfen.
	ErrBufferFull = errors.New("stream: Sendepuffer voll")
	// ErrHubClosed: Der Hub wurde geschlossen.
	ErrHubClosed = errors.New("stream: Hub geschlossen")
	// ErrReplaced: Derselbe Knoten hat sich neu verbunden; der alte Stream endet.
	ErrReplaced = errors.New("stream: durch neue Verbindung ersetzt")
	// ErrStreamingUnsupported: Der ResponseWriter kann nicht flushen.
	ErrStreamingUnsupported = errors.New("stream: ResponseWriter unterstützt kein Streaming")
	// ErrInvalidNodeID: Serve wurde ohne Knoten-ID aufgerufen.
	ErrInvalidNodeID = errors.New("stream: Knoten-ID fehlt")
	// ErrDisconnected: Der Server hat den Stream per Disconnect beendet.
	ErrDisconnected = errors.New("stream: vom Server getrennt")
)

// Hub hält die offenen Streams aller Knoten. Der Nullwert ist einsatzbereit.
// Felder dürfen nach der ersten Verwendung nicht mehr geändert werden.
type Hub struct {
	// KeepAlive ist das Intervall der Keepalive-Kommentare (Default 45s).
	KeepAlive time.Duration
	// WriteTimeout begrenzt jeden einzelnen Schreibvorgang (Default 10s).
	// Hängt ein Schreibvorgang länger, gilt der Knoten als tot.
	WriteTimeout time.Duration
	// BufferSize ist die Zahl der Events, die pro Knoten gepuffert werden (Default 16).
	BufferSize int
	// OnConnect und OnDisconnect werden aufgerufen, wenn ein Knoten online
	// bzw. offline geht. Ersetzt ein Knoten seine eigene Verbindung, gibt es
	// nur ein weiteres OnConnect, kein OnDisconnect.
	OnConnect    func(nodeID string)
	OnDisconnect func(nodeID string)

	mu     sync.Mutex
	conns  map[string]*conn
	closed bool
}

type conn struct {
	frames chan []byte
	done   chan struct{}
	once   sync.Once
	reason error
}

func (c *conn) stop(reason error) {
	c.once.Do(func() {
		c.reason = reason
		close(c.done)
	})
}

func (h *Hub) keepAlive() time.Duration    { return orDefault(h.KeepAlive, 45*time.Second) }
func (h *Hub) writeTimeout() time.Duration { return orDefault(h.WriteTimeout, 10*time.Second) }

func (h *Hub) bufferSize() int {
	if h.BufferSize > 0 {
		return h.BufferSize
	}
	return 16
}

// Serve betreibt den Stream für nodeID auf w, bis der Client die Verbindung
// trennt, ein Schreibvorgang fehlschlägt, derselbe Knoten sich neu verbindet
// (ErrReplaced) oder der Hub geschlossen wird (ErrHubClosed). Authentifizierung
// ist Sache des Aufrufers: nodeID muss bereits geprüft sein.
func (h *Hub) Serve(w http.ResponseWriter, r *http.Request, nodeID string) error {
	if nodeID == "" {
		return ErrInvalidNodeID
	}
	h.mu.Lock()
	closed := h.closed
	h.mu.Unlock()
	if closed {
		return ErrHubClosed
	}

	rc := http.NewResponseController(w)
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("X-Accel-Buffering", "no") // Reverse-Proxies nicht puffern lassen
	w.WriteHeader(http.StatusOK)
	if err := rc.Flush(); err != nil {
		if errors.Is(err, http.ErrNotSupported) {
			return ErrStreamingUnsupported
		}
		return err
	}

	c := &conn{frames: make(chan []byte, h.bufferSize()), done: make(chan struct{})}
	if err := h.register(nodeID, c); err != nil {
		return err
	}
	defer h.unregister(nodeID, c)

	ticker := time.NewTicker(h.keepAlive())
	defer ticker.Stop()
	write := func(p []byte) error {
		err := rc.SetWriteDeadline(time.Now().Add(h.writeTimeout()))
		if err != nil && !errors.Is(err, http.ErrNotSupported) {
			return err
		}
		if _, err := w.Write(p); err != nil {
			return err
		}
		return rc.Flush()
	}

	keepalive := []byte(":\n\n")
	for {
		select {
		case <-r.Context().Done():
			return r.Context().Err()
		case <-c.done:
			return c.reason
		case frame := <-c.frames:
			if err := write(frame); err != nil {
				return err
			}
			ticker.Reset(h.keepAlive())
		case <-ticker.C:
			if err := write(keepalive); err != nil {
				return err
			}
		}
	}
}

func (h *Hub) register(nodeID string, c *conn) error {
	h.mu.Lock()
	if h.closed {
		h.mu.Unlock()
		return ErrHubClosed
	}
	if h.conns == nil {
		h.conns = map[string]*conn{}
	}
	if old := h.conns[nodeID]; old != nil {
		old.stop(ErrReplaced)
	}
	h.conns[nodeID] = c
	h.mu.Unlock()

	if h.OnConnect != nil {
		h.OnConnect(nodeID)
	}
	return nil
}

func (h *Hub) unregister(nodeID string, c *conn) {
	h.mu.Lock()
	current := h.conns[nodeID] == c
	if current {
		delete(h.conns, nodeID)
	}
	h.mu.Unlock()

	c.stop(ErrNodeOffline)
	if current && h.OnDisconnect != nil {
		h.OnDisconnect(nodeID)
	}
}

// Send stellt ev in den Stream von nodeID, ohne zu blockieren. NodeID wird
// immer auf nodeID gesetzt, ID und Timestamp nur ergänzt, wenn sie fehlen.
func (h *Hub) Send(nodeID string, ev Event) error {
	if ev.Type == "" {
		return fmt.Errorf("%w: Typ fehlt", ErrInvalidEvent)
	}
	if ev.ID == "" {
		ev.ID = newEventID()
	}
	if ev.Timestamp == 0 {
		ev.Timestamp = time.Now().Unix()
	}
	ev.NodeID = nodeID
	frame, err := encode(ev)
	if err != nil {
		return err
	}

	h.mu.Lock()
	closed, c := h.closed, h.conns[nodeID]
	h.mu.Unlock()
	if closed {
		return ErrHubClosed
	}
	if c == nil {
		return ErrNodeOffline
	}
	select {
	case <-c.done:
		return ErrNodeOffline
	default:
	}
	select {
	case c.frames <- frame:
		return nil
	case <-c.done:
		return ErrNodeOffline
	default:
		return ErrBufferFull
	}
}

// Disconnect beendet den Stream von nodeID sofort, z. B. wenn Mandant oder
// Knoten gesperrt wurden. Der Knoten gilt ab Rückkehr als offline,
// OnDisconnect wird aufgerufen, Serve endet mit ErrDisconnected.
func (h *Hub) Disconnect(nodeID string) error {
	h.mu.Lock()
	c := h.conns[nodeID]
	if c != nil {
		delete(h.conns, nodeID)
	}
	h.mu.Unlock()
	if c == nil {
		return ErrNodeOffline
	}
	c.stop(ErrDisconnected)
	if h.OnDisconnect != nil {
		h.OnDisconnect(nodeID)
	}
	return nil
}

// Connected meldet, ob für nodeID ein Stream offen ist.
func (h *Hub) Connected(nodeID string) bool {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.conns[nodeID] != nil
}

// Nodes liefert die verbundenen Knoten, sortiert.
func (h *Hub) Nodes() []string {
	h.mu.Lock()
	nodes := make([]string, 0, len(h.conns))
	for n := range h.conns {
		nodes = append(nodes, n)
	}
	h.mu.Unlock()
	slices.Sort(nodes)
	return nodes
}

// Close beendet alle Streams; danach liefern Serve und Send ErrHubClosed.
func (h *Hub) Close() {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.closed = true
	for _, c := range h.conns {
		c.stop(ErrHubClosed)
	}
}

func orDefault(d, def time.Duration) time.Duration {
	if d > 0 {
		return d
	}
	return def
}