TDD

Author: entlein

src/vizier/services/adaptive_export/internal/activeset/activeset.go

@@ -89,6 +89,11 @@ func New() *ActiveSet {
 // Upsert sets or extends a pod's t_end. Idempotent — if the pod is
 // already present with a >= t_end, no delta is emitted (caller-side
 // dedup of trivial extensions; saves debouncer churn).
+//
+// `version` is advanced ONLY on membership changes (new pod added).
+// A pure t_end extension does NOT bump version — subscribers use
+// version skips as their "membership might have changed" signal, and
+// spurious bumps force unnecessary re-snapshots.
 func (s *ActiveSet) Upsert(k Key, tEnd time.Time) {
 	s.mu.Lock()
 	prev, existed := s.members[k]
@@ -97,16 +102,16 @@ func (s *ActiveSet) Upsert(k Key, tEnd time.Time) {
 		return // no-op extension; quietly skip
 	}
 	s.members[k] = tEnd
+	if existed {
+		// Pure t_end extension: store new value, no version bump,
+		// no delta. Subscribers see no membership change.
+		s.mu.Unlock()
+		return
+	}
 	s.version++
 	v := s.version
 	s.mu.Unlock()
-
-	if !existed {
-		s.broadcast(Delta{Added: []Key{k}, Version: v})
-	}
-	// extension (existed && tEnd > prev) doesn't change membership;
-	// no delta needed — subscribers don't care about t_end shifts of
-	// already-present pods.
+	s.broadcast(Delta{Added: []Key{k}, Version: v})
 }
 
 // Remove drops a pod. No-op if not present. Always emits a delta on
@@ -172,6 +177,10 @@ func (s *ActiveSet) Size() int {
 // overflow and subscribers MUST re-snapshot if they detect a version
 // gap. Channel is closed when ctx-equivalent shutdown is signalled
 // via Unsubscribe.
+//
+// Race hazard: a caller that does `Snapshot()` then `Subscribe()`
+// can miss any membership change that lands between the two calls.
+// Prefer `SubscribeAndSnapshot()` which is atomic.
 func (s *ActiveSet) Subscribe(buffer int) <-chan Delta {
 	if buffer < 1 {
 		buffer = 1
@@ -183,6 +192,40 @@ func (s *ActiveSet) Subscribe(buffer int) <-chan Delta {
 	return ch
 }
 
+// SubscribeAndSnapshot atomically captures the current membership
+// AND registers the subscription, so the consumer is guaranteed to
+// see EVERY change that lands at or after the returned version
+// without losing changes in the race window between the two.
+//
+// Returned tuple:
+//
+//	keys      — current membership at snapshot time
+//	deltas    — channel that will receive every future delta
+//	version   — the version of `keys`; consumers can filter the
+//	            channel by `delta.Version > version`
+//
+// This is the recommended consumer API for bootstrapping.
+func (s *ActiveSet) SubscribeAndSnapshot(buffer int) ([]Key, <-chan Delta, uint64) {
+	if buffer < 1 {
+		buffer = 1
+	}
+	ch := make(chan Delta, buffer)
+	// Hold BOTH mutexes for the duration of {snapshot, register}.
+	// Order: s.mu first (membership), then s.subsMu (subscriber list).
+	// broadcast() takes only s.subsMu, so there's no ordering risk.
+	s.mu.Lock()
+	keys := make([]Key, 0, len(s.members))
+	for k := range s.members {
+		keys = append(keys, k)
+	}
+	version := s.version
+	s.subsMu.Lock()
+	s.subs = append(s.subs, ch)
+	s.subsMu.Unlock()
+	s.mu.Unlock()
+	return keys, ch, version
+}
+
 // Unsubscribe removes and closes a previously-returned channel.
 // Idempotent (no error on unknown chan).
 func (s *ActiveSet) Unsubscribe(ch <-chan Delta) {

src/vizier/services/adaptive_export/internal/activeset/activeset_test.go

@@ -94,6 +94,31 @@ func TestPruneExpiredBatchesRemovals(t *testing.T) {
 	}
 }
 
+func TestUpsertExtendDoesNotAdvanceVersion(t *testing.T) {
+	// Per CR feedback (activeset.go:110): pure extension shouldn't
+	// bump version, because the version is the consumer's "did
+	// membership change?" signal. Spurious bumps make subscribers
+	// re-snapshot for nothing.
+	s := New()
+	k := Key{Pod: "p"}
+	s.Upsert(k, time.Now().Add(time.Minute))
+	_, v1 := s.Snapshot()
+	// Extend the SAME pod's t_end repeatedly.
+	for i := 0; i < 10; i++ {
+		s.Upsert(k, time.Now().Add(time.Duration(i+2)*time.Minute))
+	}
+	_, v2 := s.Snapshot()
+	if v2 != v1 {
+		t.Fatalf("version advanced on pure extension: v1=%d v2=%d", v1, v2)
+	}
+	// But a new pod DOES advance.
+	s.Upsert(Key{Pod: "q"}, time.Now().Add(time.Minute))
+	_, v3 := s.Snapshot()
+	if v3 == v2 {
+		t.Fatalf("version did NOT advance on new pod add: v=%d", v3)
+	}
+}
+
 func TestSnapshotReturnsCurrentMembers(t *testing.T) {
 	s := New()
 	s.Upsert(Key{Namespace: "n1", Pod: "p1"}, time.Now().Add(time.Minute))
@@ -131,6 +156,48 @@ func TestSubscriberOverflowDropsOldest(t *testing.T) {
 	}
 }
 
+// TestSubscribeAndSnapshot_RaceFreeBootstrap — per CR (activeset.go:183):
+// a consumer that wants both "initial state" + "all future deltas"
+// must be able to do so without missing changes between Snapshot()
+// and Subscribe(). Verify the combined helper.
+func TestSubscribeAndSnapshot_RaceFreeBootstrap(t *testing.T) {
+	s := New()
+	s.Upsert(Key{Pod: "preexisting"}, time.Now().Add(time.Minute))
+
+	// Simulate a hostile interleaving: between when we'd call Snapshot
+	// and when we'd call Subscribe, a concurrent Upsert lands.
+	// Without a combined helper, we'd miss it. The combined helper
+	// must report the new pod EITHER in the initial set OR in the
+	// first delta — never lost.
+	keys, ch, version := s.SubscribeAndSnapshot(4)
+	// Concurrent upsert AFTER subscription.
+	go func() {
+		s.Upsert(Key{Pod: "racy"}, time.Now().Add(time.Minute))
+	}()
+
+	if len(keys) != 1 || keys[0].Pod != "preexisting" {
+		t.Fatalf("initial snapshot wrong: %+v", keys)
+	}
+	// Drain delta.
+	select {
+	case d := <-ch:
+		if d.Version <= version {
+			t.Fatalf("delta version %d <= snapshot version %d", d.Version, version)
+		}
+		seen := false
+		for _, k := range d.Added {
+			if k.Pod == "racy" {
+				seen = true
+			}
+		}
+		if !seen {
+			t.Fatalf("racy pod not in delta added=%v", d.Added)
+		}
+	case <-time.After(500 * time.Millisecond):
+		t.Fatalf("no delta within 500ms")
+	}
+}
+
 func TestConcurrentUpsertsAreSafe(t *testing.T) {
 	s := New()
 	var wg sync.WaitGroup

src/vizier/services/adaptive_export/internal/controller/controller.go

@@ -647,20 +647,44 @@ func (c *Controller) PruneExpired() int {
 	grace := 2 * c.cfg.After
 	// Collect under the lock; fire callbacks AFTER releasing so we
 	// don't hold the controller mutex across user code.
-	type prunedKey struct{ namespace, pod string }
-	var fired []prunedKey
+	//
+	// IMPORTANT (rev-3 streaming correctness): c.active is keyed by
+	// anomaly hash, but the streaming layer (ActiveSet) is keyed by
+	// (namespace, pod). One pod can host multiple distinct hashes
+	// (e.g. pgsql-server has hashes for postgres, pg_isready, runc:
+	// [2:INIT] processes). Firing OnPrune for every evicted hash
+	// would prematurely stop streaming for a pod that still has
+	// other active hashes. So: compute the set of pods that have
+	// NO remaining active hashes after this prune, and only fire
+	// OnPrune for those.
+	type podKey struct{ namespace, pod string }
+	prunedHashes := 0
+	var pruned []podKey
 	c.mu.Lock()
+	// Pass 1: delete expired hashes and remember which pods THEY
+	// belonged to.
+	candidatePods := map[podKey]struct{}{}
 	for h, row := range c.active {
 		if !row.TEnd.Add(grace).After(now) {
-			fired = append(fired, prunedKey{row.Namespace, row.Pod})
+			candidatePods[podKey{row.Namespace, row.Pod}] = struct{}{}
 			delete(c.active, h)
+			prunedHashes++
 		}
 	}
+	// Pass 2: from candidatePods, remove any pod that STILL has at
+	// least one surviving hash in c.active. What's left is the set
+	// of pods that lost their LAST hash — these get OnPrune.
+	for _, row := range c.active {
+		delete(candidatePods, podKey{row.Namespace, row.Pod})
+	}
+	for pk := range candidatePods {
+		pruned = append(pruned, pk)
+	}
 	c.mu.Unlock()
 	if c.cfg.OnPrune != nil {
-		for _, k := range fired {
+		for _, k := range pruned {
 			c.cfg.OnPrune(k.namespace, k.pod)
 		}
 	}
-	return len(fired)
+	return prunedHashes
 }

src/vizier/services/adaptive_export/internal/controller/controller_test.go

@@ -469,6 +469,125 @@ func TestController_OnPrune_NilIsNoop(t *testing.T) {
 	// No panic = pass.
 }
 
+// TestController_OnPrune_OnlyFiresWhenLastHashOnPodGone — multiple
+// anomaly hashes can share a single (namespace, pod) when distinct
+// PID×comm combinations on the same pod each get their own
+// kubescape rule firing. Real-world example (sweep observation):
+// pgsql-server has hashes for processes `postgres`, `pg_isready`,
+// and `runc:[2:INIT]` — three hashes, one pod.
+//
+// The streaming layer is pod-keyed, so OnPrune(ns, pod) must only
+// fire when the LAST hash for that pod is evicted. Premature firing
+// would stop the per-pod stream while other hashes are still active.
+// CR feedback (controller.go:156) caught this; see comment thread.
+func TestController_OnPrune_OnlyFiresWhenLastHashOnPodGone(t *testing.T) {
+	trig := newFakeTrigger()
+	snk := &fakeSink{}
+	clk := &fakeClock{t: canonicalEventTime}
+
+	var mu sync.Mutex
+	var prunedPods []string
+	cfg := Config{
+		Hostname: "node-1", Before: time.Minute, After: time.Minute,
+		OnPrune: func(ns, pod string) {
+			mu.Lock()
+			defer mu.Unlock()
+			prunedPods = append(prunedPods, ns+"/"+pod)
+		},
+	}
+	c := New(trig, snk, cfg, clk)
+	stop := runController(t, c, trig)
+	defer stop()
+
+	// Two events on the SAME pod but with different (PID, Comm) so
+	// anomaly.Hash returns two distinct hashes.
+	mkEvent := func(pid uint64, comm string) kubescape.Event {
+		return kubescape.Event{
+			Target: anomaly.Target{
+				PID: pid, Comm: comm, Pod: "pgsql-server-x", Namespace: "px",
+			},
+			EventTime: uint64(canonicalEventTime.UnixNano()),
+			RuleID:    "R1", Hostname: "node-1",
+		}
+	}
+	trig.push(mkEvent(100, "postgres"))
+	trig.push(mkEvent(200, "pg_isready"))
+	waitFor(t, "two distinct hashes active", 300*time.Millisecond, func() bool {
+		return c.Active() == 2
+	})
+
+	// Advance past TEnd + 2*After so BOTH hashes are evictable.
+	clk.advance(3*time.Minute + time.Second)
+	if r := c.PruneExpired(); r != 2 {
+		t.Fatalf("PruneExpired removed %d, want 2 hashes", r)
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	if len(prunedPods) != 1 {
+		t.Fatalf("OnPrune fired %d times for one pod with 2 hashes; want 1. Calls: %v",
+			len(prunedPods), prunedPods)
+	}
+	if prunedPods[0] != "px/pgsql-server-x" {
+		t.Fatalf("wrong pod pruned: %q", prunedPods[0])
+	}
+}
+
+// TestController_OnPrune_DoesNotFireWhileOtherHashesActive — inverse
+// case: only ONE hash on a pod expires; OnPrune must NOT fire for
+// that pod because other hashes for the same pod remain active.
+func TestController_OnPrune_DoesNotFireWhileOtherHashesActive(t *testing.T) {
+	trig := newFakeTrigger()
+	snk := &fakeSink{}
+	clk := &fakeClock{t: canonicalEventTime}
+
+	var mu sync.Mutex
+	var prunedPods []string
+	cfg := Config{
+		Hostname: "node-1", Before: time.Minute, After: time.Minute,
+		OnPrune: func(ns, pod string) {
+			mu.Lock()
+			defer mu.Unlock()
+			prunedPods = append(prunedPods, ns+"/"+pod)
+		},
+	}
+	c := New(trig, snk, cfg, clk)
+	stop := runController(t, c, trig)
+	defer stop()
+
+	mkEvent := func(pid uint64) kubescape.Event {
+		return kubescape.Event{
+			Target: anomaly.Target{
+				PID: pid, Comm: "c", Pod: "samepod", Namespace: "ns",
+			},
+			EventTime: uint64(canonicalEventTime.UnixNano()),
+			RuleID:    "R1", Hostname: "node-1",
+		}
+	}
+	trig.push(mkEvent(100))
+	waitFor(t, "1 hash", 300*time.Millisecond, func() bool { return c.Active() == 1 })
+
+	// Advance time so first hash's TEnd is in the past but not yet
+	// past the 2*After grace. Then push second hash on the same pod.
+	clk.advance(2 * time.Minute)
+	trig.push(mkEvent(200))
+	waitFor(t, "2 hashes", 300*time.Millisecond, func() bool { return c.Active() == 2 })
+
+	// Advance to where the FIRST hash is past grace (3m after its
+	// creation) but the SECOND is still alive (its TEnd is at
+	// canonical+3m; grace would be +5m). Total clock progression
+	// from canonical: 2m + 1m + 1s = 3m1s.
+	clk.advance(time.Minute + time.Second)
+	removed := c.PruneExpired()
+	if removed != 1 {
+		t.Fatalf("PruneExpired removed %d, want 1 (only the old hash)", removed)
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	if len(prunedPods) != 0 {
+		t.Fatalf("OnPrune fired for a pod that still has 1 active hash; calls: %v", prunedPods)
+	}
+}
+
 // TestController_OnAttribution_NotHeldUnderMutex — a slow callback
 // must NOT block PruneExpired's progress (the controller must not
 // be holding its own mutex while invoking user code).