atelet: cache pool resolutions so a list flap cannot split the CPU counter

cmd/atelet/statspoller.go

@@ -68,9 +68,10 @@ const statsRPCTimeout = 55 * time.Second
 // ceil(n/statsSweepConcurrency) timeouts instead of n.
 const statsSweepConcurrency = 8
 
-// workerPoolListTimeout bounds the per-sweep pod list that resolves worker
-// pools. Pool labels are enrichment: better one unlabeled tick than a sweep
-// blocked on the apiserver.
+// workerPoolListTimeout bounds the per-sweep pod list that fetches worker
+// pools. Pool labels are enrichment: better to fall back to cachedPools (or,
+// for a pod not yet cached, one unlabeled tick) than a sweep blocked on the
+// apiserver.
 const workerPoolListTimeout = 10 * time.Second
 
 // clampActorStatsPollInterval enforces the floor on a nonzero configured
@@ -118,13 +119,16 @@ type statsPoller struct {
 	// connections the lifecycle RPCs are using.
 	dial func(ctx context.Context, podUID string) (activeStatsClient, io.Closer, error)
 
-	// workerPools resolves this node's worker pod UIDs to the pool that owns
-	// them, called once per sweep. Nil (or a nil map, or a missing entry)
-	// degrades to samples grouped without pool labels rather than dropped:
-	// the pool is enrichment, the sample is the point. The real resolver
-	// lists the node's pods by workerPoolLabel; the ateom directory name IS
-	// the worker pod UID, which is the join key.
-	workerPools func(ctx context.Context) map[string]workerPoolRef
+	// fetchWorkerPools is the raw fetch: one apiserver list mapping this
+	// node's worker pod UIDs to the pool that owns them, called once per sweep
+	// through resolveWorkerPools, which backs it with cachedPools so one
+	// failed list does not strand that tick's samples on a pool-less label
+	// set. A nil return means this fetch failed and the cache alone answers;
+	// a pod neither resolves groups without pool labels rather than being
+	// dropped: the pool is enrichment, the sample is the point. The real
+	// fetcher lists the node's pods by workerPoolLabel; the ateom directory
+	// name IS the worker pod UID, which is the join key.
+	fetchWorkerPools func(ctx context.Context) map[string]workerPoolRef
 
 	inst *statsInstruments
 
@@ -133,6 +137,17 @@ type statsPoller struct {
 	// Nil disables emission; emit is nil-safe.
 	eventEmitter *statsEventEmitter
 
+	// cachedPools carries pool resolutions across sweeps: a pod's pool is
+	// immutable for the pod's lifetime (the label is stamped at creation and
+	// pods do not move between pools), so an entry can never go wrong, only
+	// stale-and-then-pruned. It is what keeps one failed pod list from
+	// re-homing that tick's samples -- and, worse, the CPU counter's
+	// increments, which can never be re-attributed -- onto a pool-less label
+	// set. Only the sweep loop touches it, and resolveWorkerPools prunes it to
+	// the pods whose ateom directories still exist, the same bound lastCPU
+	// keeps.
+	cachedPools map[string]workerPoolRef
+
 	// lastCPU is the previous sweep's cpu_usage_usec per actor uid, the
 	// baseline the next sweep's deltas are computed against. Only the sweep
 	// loop touches it (under collect's mutex), and entries for actors a sweep
@@ -176,8 +191,10 @@ type templateKey struct {
 	sandboxClass      string
 	source            string
 	// workerPool is zero-valued when the pod could not be resolved to a pool
-	// (resolver disabled, list failure, pod already gone): those samples group
-	// together without pool labels rather than vanish.
+	// (resolver disabled, or a pod no fetch has resolved yet -- first seen
+	// during an outage, or omitted by a successful list racing the directory
+	// scan; previously seen pods answer from cachedPools): those samples
+	// group together without pool labels rather than vanish.
 	workerPool workerPoolRef
 }
 
@@ -244,10 +261,13 @@ func (p *statsPoller) collect(ctx context.Context) map[templateKey]*templateAggr
 		return nil
 	}
 
-	var pools map[string]workerPoolRef
-	if p.workerPools != nil {
-		pools = p.workerPools(ctx)
+	podUIDs := make([]string, 0, len(entries))
+	for _, e := range entries {
+		if e.IsDir() {
+			podUIDs = append(podUIDs, e.Name())
+		}
 	}
+	pools := p.resolveWorkerPools(ctx, podUIDs)
 
 	var (
 		mu      sync.Mutex
@@ -256,11 +276,7 @@ func (p *statsPoller) collect(ctx context.Context) map[templateKey]*templateAggr
 		g       errgroup.Group
 	)
 	g.SetLimit(statsSweepConcurrency)
-	for _, e := range entries {
-		if !e.IsDir() {
-			continue
-		}
-		podUID := e.Name()
+	for _, podUID := range podUIDs {
 
 		g.Go(func() error {
 			// One deadline over the whole probe, dial included. The real dial
@@ -381,18 +397,52 @@ func statsSourceLabel(s ateompb.StatsSource) string {
 	}
 }
 
-// nodeWorkerPools returns a resolver that lists nodeName's worker pods once
-// per sweep and maps pod UID to the pool that owns it. One field-selected,
-// label-selected LIST per interval per node is deliberately chosen over a
-// standing informer: at the poll cadence the apiserver cost is negligible,
-// there is no cache to sync before the first sweep, and a failed list
-// degrades to unlabeled samples for one tick instead of blocking anything.
-func nodeWorkerPools(client kubernetes.Interface, nodeName string) func(ctx context.Context) map[string]workerPoolRef {
+// resolveWorkerPools returns this sweep's pod-UID-to-pool map, backed by
+// cachedPools: fresh resolutions win, a failed or partial list falls back to
+// the cache (safe: a pod's pool is immutable), and the cache is rebuilt
+// restricted to the pods whose ateom directories exist this sweep, which both
+// prunes departed pods and bounds its size. The residual unlabeled case is a
+// pod no fetch has resolved yet, whether first seen during an outage or
+// omitted by a successful list racing the directory scan -- it groups without
+// pool labels until a fetch returns it.
+func (p *statsPoller) resolveWorkerPools(ctx context.Context, podUIDs []string) map[string]workerPoolRef {
+	if p.fetchWorkerPools == nil {
+		return nil
+	}
+	if len(podUIDs) == 0 {
+		// Nothing to resolve: skip the fetch, and prune the cache to match
+		// the (empty) directory set as a normal sweep would.
+		p.cachedPools = nil
+		return nil
+	}
+	fresh := p.fetchWorkerPools(ctx)
+	merged := make(map[string]workerPoolRef, len(podUIDs))
+	for _, uid := range podUIDs {
+		if ref, ok := fresh[uid]; ok {
+			merged[uid] = ref
+		} else if ref, ok := p.cachedPools[uid]; ok {
+			merged[uid] = ref
+		}
+	}
+	p.cachedPools = merged
+	return merged
+}
+
+// newWorkerPoolFetcher returns the fetch func behind fetchWorkerPools: it
+// lists nodeName's worker pods and maps pod UID to the pool that owns it,
+// once per call. One field-selected, label-selected LIST per interval per
+// node is deliberately chosen over a standing informer: at the poll cadence
+// the apiserver cost is negligible, there is no informer cache to sync before
+// the first sweep, and a failed list costs nothing that matters --
+// resolveWorkerPools answers from cachedPools for every pod seen before, and
+// only a pod first seen during the outage goes unlabeled until a list
+// succeeds.
+func newWorkerPoolFetcher(client kubernetes.Interface, nodeName string) func(ctx context.Context) map[string]workerPoolRef {
 	return func(ctx context.Context) map[string]workerPoolRef {
 		// Bounded so a hung apiserver connection cannot stall the sweep it is
-		// merely enriching: past the deadline, this tick's samples group
-		// without pool labels, which is the same answer as any other failed
-		// list.
+		// merely enriching: past the deadline this fetch returns nil, which
+		// resolveWorkerPools treats like any other failed fetch -- cached
+		// pods keep their pools, unseen pods group without them.
 		listCtx, cancel := context.WithTimeout(ctx, workerPoolListTimeout)
 		defer cancel()
 		pods, err := client.CoreV1().Pods(metav1.NamespaceAll).List(listCtx, metav1.ListOptions{
@@ -567,7 +617,7 @@ func startStatsPoller(ctx context.Context, interval time.Duration, inst *statsIn
 	// NODE_NAME comes from the Downward API; without it the samples still
 	// flow, just grouped without pool labels.
 	if nodeName := os.Getenv("NODE_NAME"); nodeName != "" {
-		poller.workerPools = nodeWorkerPools(k8sClient, nodeName)
+		poller.fetchWorkerPools = newWorkerPoolFetcher(k8sClient, nodeName)
 	} else {
 		slog.WarnContext(ctx, "NODE_NAME not set; actor stats will carry no worker pool labels")
 	}

cmd/atelet/statspoller_test.go

@@ -347,7 +347,7 @@ func TestStatsPollerWorkerPoolLabels(t *testing.T) {
 		"uid-unpooled": {resp: executingResponse("ns-a", "tmpl-a", ateompb.SandboxClass_SANDBOX_CLASS_GVISOR, ateompb.StatsSource_STATS_SOURCE_CGROUP, 10, 8)},
 	}
 	p, _ := newPollerFixture(t, fakes)
-	p.workerPools = func(context.Context) map[string]workerPoolRef {
+	p.fetchWorkerPools = func(context.Context) map[string]workerPoolRef {
 		return map[string]workerPoolRef{"uid-pooled": {namespace: "pool-ns", name: "pool-a"}}
 	}
 
@@ -376,7 +376,7 @@ func TestStatsPollerPeriodicEvents(t *testing.T) {
 	p, _ := newPollerFixture(t, fakes)
 	var buf syncBuffer
 	p.eventEmitter = newBufferEmitter(&buf, false)
-	p.workerPools = func(context.Context) map[string]workerPoolRef {
+	p.fetchWorkerPools = func(context.Context) map[string]workerPoolRef {
 		return map[string]workerPoolRef{"uid-1": {namespace: "pool-ns", name: "pool-a"}}
 	}
 
@@ -470,12 +470,12 @@ func TestStatsPollerCPUDeltaSaturatesCorruptCounter(t *testing.T) {
 	}
 }
 
-// TestNodeWorkerPools pins the resolver's ingestion rules: a labeled worker
+// TestNewWorkerPoolFetcher pins the fetcher's ingestion rules: a labeled worker
 // maps by pod UID, an empty label value names no pool and never enters the
 // map (the presence-only selector matches it anyway), and unlabeled pods are
 // not workers at all. The fake clientset honors label selectors but not the
 // spec.nodeName field selector, so node scoping is not assertable here.
-func TestNodeWorkerPools(t *testing.T) {
+func TestNewWorkerPoolFetcher(t *testing.T) {
 	client := k8sfake.NewSimpleClientset(
 		&corev1.Pod{ObjectMeta: metav1.ObjectMeta{
 			Name: "worker-a", Namespace: "pool-ns", UID: "uid-a",
@@ -490,10 +490,123 @@ func TestNodeWorkerPools(t *testing.T) {
 		}},
 	)
 
-	got := nodeWorkerPools(client, "node-1")(context.Background())
+	got := newWorkerPoolFetcher(client, "node-1")(context.Background())
 
 	want := map[string]workerPoolRef{"uid-a": {namespace: "pool-ns", name: "pool-a"}}
 	if diff := cmp.Diff(want, got, cmp.AllowUnexported(workerPoolRef{})); diff != "" {
-		t.Errorf("nodeWorkerPools mismatch (-want +got):\n%s", diff)
+		t.Errorf("newWorkerPoolFetcher mismatch (-want +got):\n%s", diff)
+	}
+}
+
+// TestStatsPollerPoolCacheSurvivesListFlap pins the fix for the label-set
+// split: a pod's pool is immutable, so a failed list must answer from the
+// cache and keep that tick's samples -- and the CPU counter's increments,
+// which can never be re-attributed -- on the pooled label set.
+func TestStatsPollerPoolCacheSurvivesListFlap(t *testing.T) {
+	fakes := map[string]*fakeStatsAteom{
+		"uid-1": {resp: executingResponse("ns-a", "tmpl-a", ateompb.SandboxClass_SANDBOX_CLASS_GVISOR, ateompb.StatsSource_STATS_SOURCE_CGROUP, 100, 80)},
+	}
+	p, _ := newPollerFixture(t, fakes)
+	listOK := true
+	p.fetchWorkerPools = func(context.Context) map[string]workerPoolRef {
+		if !listOK {
+			return nil // the apiserver list failed this sweep
+		}
+		return map[string]workerPoolRef{"uid-1": {namespace: "pool-ns", name: "pool-a"}}
+	}
+
+	pooled := templateKey{templateNamespace: "ns-a", templateName: "tmpl-a", sandboxClass: "gvisor", source: "cgroup",
+		workerPool: workerPoolRef{namespace: "pool-ns", name: "pool-a"}}
+
+	// Sweep 1 resolves and seeds the cache.
+	if got := p.collect(context.Background()); got[pooled] == nil {
+		t.Fatalf("sweep 1: no pooled aggregate; got %v", got)
+	}
+
+	// Sweep 2's list fails; the samples must STILL group under the pool.
+	listOK = false
+	got := p.collect(context.Background())
+	if got[pooled] == nil {
+		t.Errorf("sweep 2 (list flap): samples left the pooled label set; got %v", got)
+	}
+	if len(got) != 1 {
+		t.Errorf("sweep 2 (list flap): %d label sets, want 1 (no pool-less split)", len(got))
+	}
+}
+
+// TestStatsPollerPoolCachePrunes: the cache is rebuilt against the pods whose
+// ateom directories exist, so a departed pod's entry does not linger.
+func TestStatsPollerPoolCachePrunes(t *testing.T) {
+	fakes := map[string]*fakeStatsAteom{
+		"uid-1": {resp: noSampleResponse(ateompb.NoSampleReason_NO_SAMPLE_REASON_NO_WORKLOAD)},
+	}
+	p, _ := newPollerFixture(t, fakes)
+	p.fetchWorkerPools = func(context.Context) map[string]workerPoolRef {
+		return map[string]workerPoolRef{
+			"uid-1":    {namespace: "pool-ns", name: "pool-a"},
+			"uid-gone": {namespace: "pool-ns", name: "pool-a"}, // no ateom dir
+		}
+	}
+
+	p.collect(context.Background())
+
+	if _, ok := p.cachedPools["uid-1"]; !ok {
+		t.Errorf("cachedPools lost the live pod's entry: %v", p.cachedPools)
+	}
+	if _, ok := p.cachedPools["uid-gone"]; ok {
+		t.Errorf("cachedPools kept an entry with no ateom directory: %v", p.cachedPools)
+	}
+}
+
+// TestStatsPollerPoolCacheMissDuringOutage: a pod first seen while the list
+// is failing has no cache entry to fall back to -- it groups without pool
+// labels (the residual, documented case) and heals on the next good list.
+func TestStatsPollerPoolCacheMissDuringOutage(t *testing.T) {
+	fakes := map[string]*fakeStatsAteom{
+		"uid-new": {resp: executingResponse("ns-a", "tmpl-a", ateompb.SandboxClass_SANDBOX_CLASS_GVISOR, ateompb.StatsSource_STATS_SOURCE_CGROUP, 10, 8)},
+	}
+	p, _ := newPollerFixture(t, fakes)
+	p.fetchWorkerPools = func(context.Context) map[string]workerPoolRef { return nil }
+
+	got := p.collect(context.Background())
+
+	bare := templateKey{templateNamespace: "ns-a", templateName: "tmpl-a", sandboxClass: "gvisor", source: "cgroup"}
+	if got[bare] == nil || len(got) != 1 {
+		t.Errorf("collect() during outage = %v, want the one pool-less group", got)
+	}
+}
+
+// TestStatsPollerPoolCachePartialListFallsBack pins the per-pod half of the
+// fallback the resolver promises ("a failed OR PARTIAL list"): a fetch that
+// succeeds but omits a previously cached pod must not strand that pod's
+// samples off the pooled label set. Distinct from the all-nil flap test above
+// -- a whole-map fallback (fresh == nil ? cache : fresh) would pass that test
+// and fail this one.
+func TestStatsPollerPoolCachePartialListFallsBack(t *testing.T) {
+	fakes := map[string]*fakeStatsAteom{
+		"uid-1": {resp: executingResponse("ns-a", "tmpl-a", ateompb.SandboxClass_SANDBOX_CLASS_GVISOR, ateompb.StatsSource_STATS_SOURCE_CGROUP, 100, 80)},
+	}
+	p, _ := newPollerFixture(t, fakes)
+	full := true
+	p.fetchWorkerPools = func(context.Context) map[string]workerPoolRef {
+		if !full {
+			// A successful list that no longer contains uid-1 (races between
+			// the pod store and the dir scan look exactly like this).
+			return map[string]workerPoolRef{"uid-other": {namespace: "pool-ns", name: "pool-b"}}
+		}
+		return map[string]workerPoolRef{"uid-1": {namespace: "pool-ns", name: "pool-a"}}
+	}
+
+	pooled := templateKey{templateNamespace: "ns-a", templateName: "tmpl-a", sandboxClass: "gvisor", source: "cgroup",
+		workerPool: workerPoolRef{namespace: "pool-ns", name: "pool-a"}}
+
+	if got := p.collect(context.Background()); got[pooled] == nil {
+		t.Fatalf("sweep 1: no pooled aggregate; got %v", got)
+	}
+
+	full = false
+	got := p.collect(context.Background())
+	if got[pooled] == nil || len(got) != 1 {
+		t.Errorf("sweep 2 (partial list): samples left the pooled label set; got %v", got)
 	}
 }