internal/server/api/snapshot_hub_test.go
Ref: Size: 8.8 KiB History
package api
import (
"fmt"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// waitFor polls cond until it is true or the deadline elapses.
func waitFor(t *testing.T, d time.Duration, cond func() bool) {
t.Helper()
deadline := time.Now().Add(d)
for time.Now().Before(deadline) {
if cond() {
return
}
time.Sleep(time.Millisecond)
}
require.True(t, cond(), "condition not met within %s", d)
}
// constBuild returns a build func that hands every requested tenant the same
// bytes b (ignoring the tenant), counting how many times it is called.
func constBuild(b string, calls *atomic.Int64) func([]string) (map[string][]byte, error) {
return func(tenants []string) (map[string][]byte, error) {
calls.Add(1)
out := make(map[string][]byte, len(tenants))
for _, tn := range tenants {
out[tn] = []byte(b)
}
return out, nil
}
}
// TestSnapshotHubCloseIdempotent proves Close can be called more than once
// without panicking (close-of-closed-channel) — a shutdown path plus a test
// cleanup must both be safe.
func TestSnapshotHubCloseIdempotent(t *testing.T) {
var calls atomic.Int64
h := newSnapshotHub(constBuild("x", &calls), newNotifier())
go h.run()
h.Close()
h.Close() // must not panic
}
// TestSnapshotHubSingleBuildFanout proves the whole point of the hub: one
// underlying build per tick is fanned to every subscriber OF A TENANT as
// identical bytes, and a wake drives exactly one build regardless of how many
// clients are attached (M clients on one tenant must NOT cause M builds).
func TestSnapshotHubSingleBuildFanout(t *testing.T) {
var builds atomic.Int64
build := func(tenants []string) (map[string][]byte, error) {
n := builds.Add(1)
out := make(map[string][]byte, len(tenants))
for _, tn := range tenants {
out[tn] = []byte(fmt.Sprintf("snap-%d", n))
}
return out, nil
}
notif := newNotifier()
h := newSnapshotHub(build, notif)
go h.run()
defer h.Close()
// Three subscribers on the SAME tenant. The first triggers the initial build
// for that tenant; the other two reuse the cached current.
const n = 3
chans := make([]<-chan []byte, n)
for i := range n {
ch, unsub := h.subscribe("default")
defer unsub()
chans[i] = ch
}
require.Equal(t, int64(1), builds.Load(), "only the first subscriber of a tenant builds")
for i, ch := range chans {
select {
case b := <-ch:
assert.Equal(t, "snap-1", string(b), "subscriber %d initial snapshot", i)
case <-time.After(time.Second):
t.Fatalf("subscriber %d got no initial snapshot", i)
}
}
// One wake ⇒ exactly one build (the single subscribed tenant), fanned to all three.
notif.notify()
waitFor(t, time.Second, func() bool { return builds.Load() == 2 })
for i, ch := range chans {
select {
case b := <-ch:
assert.Equal(t, "snap-2", string(b), "subscriber %d wake snapshot", i)
case <-time.After(time.Second):
t.Fatalf("subscriber %d got no wake snapshot", i)
}
}
assert.Equal(t, int64(2), builds.Load(),
"one wake must drive exactly one build regardless of client count")
}
// TestSnapshotHubPerTenantPayloads proves each subscriber only ever receives its
// OWN tenant's bytes: a single recompute marshals a distinct payload per
// subscribed tenant and fans each to the right connection.
func TestSnapshotHubPerTenantPayloads(t *testing.T) {
// build echoes the tenant name into the payload so a mixup is visible.
build := func(tenants []string) (map[string][]byte, error) {
out := make(map[string][]byte, len(tenants))
for _, tn := range tenants {
out[tn] = []byte("payload-for-" + tn)
}
return out, nil
}
h := newSnapshotHub(build, newNotifier())
go h.run()
defer h.Close()
chA, unsubA := h.subscribe("alpha")
defer unsubA()
chB, unsubB := h.subscribe("beta")
defer unsubB()
for _, tc := range []struct {
ch <-chan []byte
want string
}{{chA, "payload-for-alpha"}, {chB, "payload-for-beta"}} {
select {
case b := <-tc.ch:
assert.Equal(t, tc.want, string(b))
case <-time.After(time.Second):
t.Fatalf("no initial snapshot for %s", tc.want)
}
}
}
// TestSnapshotHubSuppressesUnchanged pins change-suppression: identical bytes on
// recompute produce no push to a tenant's subscribers.
func TestSnapshotHubSuppressesUnchanged(t *testing.T) {
var builds atomic.Int64
notif := newNotifier()
h := newSnapshotHub(constBuild("constant", &builds), notif)
go h.run()
defer h.Close()
ch, unsub := h.subscribe("default")
defer unsub()
// Drain the immediate initial delivery.
select {
case b := <-ch:
assert.Equal(t, "constant", string(b))
case <-time.After(time.Second):
t.Fatal("no initial snapshot")
}
notif.notify()
waitFor(t, time.Second, func() bool { return builds.Load() >= 2 })
// Bytes are unchanged, so nothing new must be delivered.
select {
case b := <-ch:
t.Fatalf("unexpected push of unchanged snapshot: %q", b)
case <-time.After(100 * time.Millisecond):
}
}
// TestSnapshotHubLatestWins proves the fan-out is non-blocking / latest-wins: a
// slow (never-draining) subscriber's channel holds only the newest snapshot and
// the central loop never blocks on it.
func TestSnapshotHubLatestWins(t *testing.T) {
var n atomic.Int64
build := func(tenants []string) (map[string][]byte, error) {
v := n.Add(1)
out := make(map[string][]byte, len(tenants))
for _, tn := range tenants {
out[tn] = []byte(fmt.Sprintf("v%d", v))
}
return out, nil
}
notif := newNotifier()
h := newSnapshotHub(build, notif)
go h.run()
defer h.Close()
// Subscribe but never read: the buffer-1 channel already holds the initial
// snapshot. Further recomputes must overwrite it (drain-stale-then-send),
// never block the hub.
ch, unsub := h.subscribe("default")
defer unsub()
for i := range 5 {
notif.notify()
waitFor(t, time.Second, func() bool { return n.Load() >= int64(i+2) })
}
// The lone buffered value must be the LATEST, not a stale early one.
last := fmt.Sprintf("v%d", n.Load())
select {
case b := <-ch:
assert.Equal(t, last, string(b), "slow subscriber must hold the newest snapshot")
case <-time.After(time.Second):
t.Fatal("slow subscriber holds nothing")
}
}
// TestSnapshotHubDropsTenantOnLastUnsub proves current is pruned when a tenant's
// last subscriber leaves, so transient tenants cannot grow the map unboundedly.
func TestSnapshotHubDropsTenantOnLastUnsub(t *testing.T) {
var calls atomic.Int64
h := newSnapshotHub(constBuild("x", &calls), newNotifier())
go h.run()
defer h.Close()
_, unsub := h.subscribe("ephemeral")
h.mu.Lock()
_, present := h.current["ephemeral"]
h.mu.Unlock()
require.True(t, present, "subscribing must cache the tenant's payload")
unsub()
h.mu.Lock()
_, present = h.current["ephemeral"]
h.mu.Unlock()
assert.False(t, present, "last unsubscribe must drop the tenant's cached payload")
}
// TestSnapshotHubRecomputeFansEachTenantItsOwnBytes pins tenant isolation on the
// path that carries every snapshot after the first one.
//
// TestSnapshotHubPerTenantPayloads covers the subscribe() delivery, which is a
// different code path: subscribe hands back h.current[tenant] directly, while
// recompute walks h.subs and picks a payload per channel. A mixup in that fan
// loop is a live cross-tenant fleet disclosure — every VM, host and address the
// other tenant owns — and it survives the whole hub suite today.
//
// The payload carries both halves of the claim: the tenant it was built for,
// and the build generation, so the assertion proves the bytes came from THIS
// recompute rather than a stale initial delivery.
func TestSnapshotHubRecomputeFansEachTenantItsOwnBytes(t *testing.T) {
var builds atomic.Int64
build := func(tenants []string) (map[string][]byte, error) {
n := builds.Add(1)
out := make(map[string][]byte, len(tenants))
for _, tn := range tenants {
out[tn] = fmt.Appendf(nil, "%s-v%d", tn, n)
}
return out, nil
}
notif := newNotifier()
h := newSnapshotHub(build, notif)
go h.run()
defer h.Close()
chA, unsubA := h.subscribe("alpha")
defer unsubA()
chB, unsubB := h.subscribe("beta")
defer unsubB()
// Drain the two first-subscriber builds so what follows can only be the
// recompute delivery.
drainOne(t, chA, "alpha-v1")
drainOne(t, chB, "beta-v2")
notif.notify()
waitFor(t, time.Second, func() bool { return builds.Load() == 3 })
drainOne(t, chA, "alpha-v3")
drainOne(t, chB, "beta-v3")
}
// drainOne reads one snapshot and asserts its exact bytes, naming the leak in
// the failure so the tenant that received the wrong payload is in the message.
func drainOne(t *testing.T, ch <-chan []byte, want string) {
t.Helper()
select {
case b := <-ch:
assert.Equal(t, want, string(b),
"recompute delivered %q where %q was due — the fan-out must route each tenant its own payload; crossing them is a cross-tenant fleet disclosure", b, want)
case <-time.After(time.Second):
t.Fatalf("no snapshot delivered; expected %q", want)
}
}