fathom/internal/web/events.go

package web

import (
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"git.z.gd/zzg/fathom/internal/store"
)

// Hub broadcasts ingest events to connected browsers over Server-Sent Events.
//
// SSE rather than WebSockets or polling. The traffic is one-directional and
// low-rate — the server tells the browser "something landed" and the browser
// decides what to refetch — which is exactly what SSE is for, and it is plain
// HTTP with no handshake, no framing, and no dependency. Browsers also
// reconnect on their own, so a restarted server heals without client code.
//
// Polling was the alternative and is what this replaces. A five-second poll is
// simultaneously too slow to feel live and too chatty when nothing is running;
// pushing on flush is both quieter and more immediate.
type Hub struct {
	mu      sync.Mutex
	clients map[chan []byte]struct{}
	// seq lets a client tell a genuinely new event from a duplicate delivered
	// after a reconnect.
	seq uint64

	// listeners are in-process consumers of the same signal, so the store keeps
	// one flush callback rather than a growing list.
	listeners []func()
}

// OnFlushed registers an in-process listener for flush notifications.
func (h *Hub) OnFlushed(fn func()) {
	h.mu.Lock()
	h.listeners = append(h.listeners, fn)
	h.mu.Unlock()
}

func NewHub() *Hub {
	return &Hub{clients: make(map[chan []byte]struct{})}
}

// Event is what the browser receives.
type Event struct {
	Seq      uint64   `json:"seq"`
	Type     string   `json:"type"`
	Spans    int      `json:"spans,omitempty"`
	Logs     int      `json:"logs,omitempty"`
	TraceIDs []string `json:"trace_ids,omitempty"`
}

// NotifyFlush is the store's OnFlush callback.
//
// The store calls this synchronously on the flushing goroutine, so it must not
// block. Every client send is non-blocking: a browser that has stopped reading
// gets its event dropped rather than stalling ingest. Dropping is safe because
// the event carries no state the client cannot recover — it is a hint to
// refetch, and the next flush will hint again.
func (h *Hub) NotifyFlush(info store.FlushInfo) {
	if info.Spans == 0 && info.Logs == 0 {
		return
	}
	h.mu.Lock()
	h.seq++
	ev := Event{
		Seq:      h.seq,
		Type:     "flush",
		Spans:    info.Spans,
		Logs:     info.Logs,
		TraceIDs: info.TraceIDs,
	}
	listeners := append([]func(){}, h.listeners...)
	h.mu.Unlock()

	h.broadcast(ev)

	// Listeners must not block the flusher either; each gets its own goroutine
	// and is expected to coalesce on its own side.
	for _, fn := range listeners {
		go fn()
	}
}

func (h *Hub) broadcast(ev Event) {
	payload, err := json.Marshal(ev)
	if err != nil {
		return
	}
	frame := []byte("data: " + string(payload) + "\n\n")

	h.mu.Lock()
	defer h.mu.Unlock()
	for ch := range h.clients {
		select {
		case ch <- frame:
		default:
			// Slow or dead client; drop this event rather than block ingest.
		}
	}
}

func (h *Hub) add(ch chan []byte) {
	h.mu.Lock()
	h.clients[ch] = struct{}{}
	h.mu.Unlock()
}

func (h *Hub) remove(ch chan []byte) {
	h.mu.Lock()
	delete(h.clients, ch)
	h.mu.Unlock()
}

// Clients reports the number of connected browsers, for tests and diagnostics.
func (h *Hub) Clients() int {
	h.mu.Lock()
	defer h.mu.Unlock()
	return len(h.clients)
}

const (
	// heartbeatInterval keeps the connection alive through proxies and, more
	// usefully, lets the browser notice a server that has gone away: no comment
	// for a while means the stream is dead, and EventSource reconnects.
	heartbeatInterval = 20 * time.Second

	// clientBuffer absorbs a burst without blocking the flusher. Events are
	// small and coalesced by the client anyway, so this is generous.
	clientBuffer = 16
)

// handleEvents streams events for the lifetime of the request.
func (h *Hub) handleEvents(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		httpError(w, http.StatusInternalServerError, "this server cannot stream")
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	// Without this, a reverse proxy in front of Fathom will happily buffer the
	// whole stream and deliver nothing until the connection closes.
	w.Header().Set("X-Accel-Buffering", "no")
	w.WriteHeader(http.StatusOK)

	// Tell the browser how soon to retry, and confirm the stream is open so the
	// UI can show a connected state immediately rather than after the first
	// trace arrives.
	fmt.Fprintf(w, "retry: 2000\n\n")
	fmt.Fprintf(w, "data: {\"type\":\"hello\"}\n\n")
	flusher.Flush()

	ch := make(chan []byte, clientBuffer)
	h.add(ch)
	defer h.remove(ch)

	ticker := time.NewTicker(heartbeatInterval)
	defer ticker.Stop()

	ctx := r.Context()
	for {
		select {
		case <-ctx.Done():
			return
		case frame := <-ch:
			if _, err := w.Write(frame); err != nil {
				return
			}
			flusher.Flush()
		case <-ticker.C:
			// A comment line: valid SSE, ignored by EventSource, enough to
			// detect a broken pipe.
			if _, err := fmt.Fprint(w, ": keepalive\n\n"); err != nil {
				return
			}
			flusher.Flush()
		}
	}
}