internal/server/api/snapshot_hub.go
Ref: Size: 6.8 KiB History
package api
import (
"bytes"
"sync"
"time"
)
// snapshotHub computes the SSE snapshot centrally and fans it to connected
// clients, PER TENANT. Before this, each SSE client goroutine ran its own 1s
// ticker and marshalled its own snapshot (a DB tx + full build + json.Marshal),
// so M clients meant M full snapshots every second. The hub collapses that to
// one store read per tick plus one marshal per DISTINCT SUBSCRIBED TENANT,
// regardless of how many clients each tenant has.
//
// Isolation: a subscriber registers with its tenant and only ever receives that
// tenant's bytes — the hub filters the fleet snapshot into a separate payload
// per tenant, so no cross-tenant host/VM/metric can reach a connection.
//
// Fan-out is latest-wins / non-blocking: each subscriber holds a buffer-1
// channel carrying only the newest snapshot for its tenant. A lagging client
// drops intermediate snapshots (correct — SSE state is a full snapshot, so the
// newest supersedes any it missed), and the central loop NEVER blocks on a slow
// reader, so a stuck client can neither wedge the hub nor grow memory unbounded.
type snapshotHub struct {
// build reads the store snapshot ONCE and returns one marshalled payload per
// requested tenant (a.marshalSnapshots). The hub passes it the distinct set
// of currently-subscribed tenants.
build func(tenants []string) (map[string][]byte, error)
wake <-chan struct{} // desired-state wake source (single subscription)
notifUnsub func() // releases the notifier subscription on Close
mu sync.Mutex
// current holds the latest marshalled snapshot per tenant, delivered to new
// subscribers of that tenant immediately.
current map[string][]byte
// subs maps each subscriber channel to the tenant it is scoped to.
subs map[chan []byte]string
stop chan struct{}
done chan struct{}
stopOnce sync.Once
}
// newSnapshotHub builds the hub and subscribes to the notifier SYNCHRONOUSLY (so
// no wake can be lost in the window before run's goroutine is scheduled). Unlike
// the pre-tenant hub it computes NO initial snapshot — payloads are per-tenant
// and no tenant is known until a client subscribes, so the initial build for a
// tenant happens on its first subscribe. The caller must launch run() (once).
func newSnapshotHub(build func(tenants []string) (map[string][]byte, error), notif *notifier) *snapshotHub {
wake, unsub := notif.subscribe()
return &snapshotHub{
build: build,
wake: wake,
notifUnsub: unsub,
current: make(map[string][]byte),
subs: make(map[chan []byte]string),
stop: make(chan struct{}),
done: make(chan struct{}),
}
}
// run is the hub's single goroutine: it recomputes on a 1s tick (to catch
// agent-reported actual-state changes) or on a desired-state wake, at most once
// per event. It returns when Close is called.
func (h *snapshotHub) run() {
defer close(h.done)
tick := time.NewTicker(time.Second)
defer tick.Stop()
for {
select {
case <-h.stop:
return
case <-h.wake:
h.recompute()
case <-tick.C:
h.recompute()
}
}
}
// subscribedTenants returns the distinct set of tenants with at least one live
// subscriber. Caller must hold h.mu.
func (h *snapshotHub) subscribedTenants() []string {
seen := make(map[string]struct{}, len(h.subs))
out := make([]string, 0, len(h.subs))
for _, tn := range h.subs {
if _, ok := seen[tn]; ok {
continue
}
seen[tn] = struct{}{}
out = append(out, tn)
}
return out
}
// recompute reads the store ONCE, marshals a payload per subscribed tenant, and
// for each tenant whose bytes changed, stores them and fans them to that tenant's
// subscribers. Unchanged bytes are suppressed (no push), matching the prior
// per-client diff behaviour.
func (h *snapshotHub) recompute() {
h.mu.Lock()
tenants := h.subscribedTenants()
h.mu.Unlock()
if len(tenants) == 0 {
return // no subscribers: nothing to build
}
payloads, err := h.build(tenants)
if err != nil {
return // transient store error: keep serving the last good snapshots
}
h.mu.Lock()
defer h.mu.Unlock()
// A tenant can lose its last subscriber between snapshotting tenants above and
// re-taking the lock; only cache/fan payloads whose tenant still has a live
// subscriber, so h.current never accumulates entries for departed tenants.
live := make(map[string]struct{})
for _, tn := range h.subs {
live[tn] = struct{}{}
}
changed := make(map[string]struct{}, len(payloads))
for tn, b := range payloads {
if _, ok := live[tn]; !ok {
continue
}
if !bytes.Equal(b, h.current[tn]) {
h.current[tn] = b
changed[tn] = struct{}{}
}
}
for ch, tn := range h.subs {
if _, ok := changed[tn]; !ok {
continue
}
b := h.current[tn]
// Latest-wins: drain any stale pending snapshot, then send the newest.
// Both sends are non-blocking; because subscribers only ever receive
// (never send) and we hold h.mu, the post-drain buffer has room and the
// hub can never block on a slow client.
select {
case <-ch:
default:
}
select {
case ch <- b:
default:
}
}
}
// subscribe registers a client scoped to tenant and immediately delivers that
// tenant's CURRENT snapshot so a new connection gets initial state without doing
// its own marshal. The first subscriber for a tenant computes that tenant's
// payload on demand (outside the lock — build reads the store). It returns the
// client's buffer-1 channel and an unsubscribe func.
func (h *snapshotHub) subscribe(tenant string) (<-chan []byte, func()) {
ch := make(chan []byte, 1)
h.mu.Lock()
cur, ok := h.current[tenant]
h.mu.Unlock()
if !ok {
// First subscriber for this tenant: build its initial payload now so the
// connection gets state without waiting for the next tick. A build error
// just means no initial frame — the next recompute will deliver one.
if payloads, err := h.build([]string{tenant}); err == nil {
cur = payloads[tenant]
}
}
h.mu.Lock()
if cur != nil {
h.current[tenant] = cur
ch <- cur // buffer-1 and empty ⇒ never blocks under the lock
}
h.subs[ch] = tenant
h.mu.Unlock()
return ch, func() {
h.mu.Lock()
delete(h.subs, ch)
// Drop the tenant's cached payload once its last subscriber leaves, so a
// fleet of transient tenants cannot grow h.current unboundedly.
if !h.tenantHasSubs(tenant) {
delete(h.current, tenant)
}
h.mu.Unlock()
}
}
// tenantHasSubs reports whether any live subscriber is scoped to tenant. Caller
// must hold h.mu.
func (h *snapshotHub) tenantHasSubs(tenant string) bool {
for _, tn := range h.subs {
if tn == tenant {
return true
}
}
return false
}
// Close stops the hub goroutine, waits for it to exit, and releases the notifier
// subscription. Idempotent — safe to call more than once (e.g. a shutdown path
// plus a test cleanup).
func (h *snapshotHub) Close() {
h.stopOnce.Do(func() {
close(h.stop)
<-h.done
h.notifUnsub()
})
}