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

cmd/atelet/statspoller.go

@@ -68,9 +68,9 @@ 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, so a slow apiserver must not stall the
+// sweep it is merely enriching.
 const workerPoolListTimeout = 10 * time.Second
 
 // clampActorStatsPollInterval enforces the floor on a nonzero configured
@@ -118,13 +118,13 @@ 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 by resolveWorkerPools, which layers cachedPools over it. A nil
+	// return means the fetch failed. 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 +133,15 @@ type statsPoller struct {
 	// Nil disables emission; emit is nil-safe.
 	eventEmitter *statsEventEmitter
 
+	// cachedPools carries pool resolutions across sweeps, so one failed pod
+	// list cannot re-home a tick's samples -- and the CPU counter's
+	// increments, which can never be re-attributed -- onto a pool-less label
+	// set. Safe because a pod's pool is immutable for the pod's lifetime: an
+	// entry can be stale, never wrong. Only the sweep loop touches it;
+	// 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
@@ -175,9 +184,9 @@ type templateKey struct {
 	templateName      string
 	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.
+	// workerPool is zero-valued when no fetch has resolved the pod yet (see
+	// resolveWorkerPools): those samples group together without pool labels
+	// rather than vanish.
 	workerPool workerPoolRef
 }
 
@@ -244,10 +253,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 +268,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 +389,46 @@ 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: fresh
+// resolutions win, a failed or partial list falls back to cachedPools, and
+// the result -- rebuilt restricted to the pods whose ateom directories exist
+// -- becomes the new cache, pruning departed pods and bounding its size. The
+// residual unlabeled case is a pod no fetch has resolved yet, whether first
+// seen during an outage or omitted by a 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 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: one
+// field-selected, label-selected LIST of nodeName's worker pods per call.
+// Deliberately not a standing informer: at the poll cadence the apiserver
+// cost is negligible, there is no informer cache to sync before the first
+// sweep, and resolveWorkerPools already absorbs failed lists.
+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.
+		// Bounded so a hung apiserver cannot stall the sweep: past the
+		// deadline the fetch returns nil and resolveWorkerPools falls back
+		// to the cache.
 		listCtx, cancel := context.WithTimeout(ctx, workerPoolListTimeout)
 		defer cancel()
 		pods, err := client.CoreV1().Pods(metav1.NamespaceAll).List(listCtx, metav1.ListOptions{
@@ -567,7 +603,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,121 @@ 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 failed list must answer from the cache and keep that tick's
+// samples 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
+// promised fallback ("a failed OR PARTIAL list"): a fetch that succeeds but
+// omits a cached pod must not strand that pod's samples. Distinct from the
+// all-nil flap test above -- a whole-map fallback 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)
 	}
 }