pkg/clusteragent/autoscaling/cluster/spot: fix the race between pod updates and workload opt-out (#49880)

### What does this PR do?

Refactor podTracker to use up-to-date spot config and untrack workload in case it is not spot eligible anymore (opted-out).

### Motivation

This fixes the race between pod and workload updates where pod updates may arrive after workload was opted-out and silently recreate the pod set in the tracker.

### Describe how you validated your changes

Reproduced the race in main via:
```
while go test -tags "kubeapiserver test" -count=1 ./pkg/clusteragent/autoscaling/cluster/spot/... ; do echo OK ; done
```
and verified it does not reproduce with the changes. 

### Additional Notes

To reproduce the race in tests delay WLM updates.

Test scenarios use fake clock to speed-up on-demand fallback and rebalancing. The rebalancer skips an iteration when there are pending pods or in-flight admissions therefore update tests to step fake clock after workload updates have settled.

Also refactor manual pod creation/deletion into fakeDeployment Reconcile and ScaleDown methods.

Updates https://datadoghq.atlassian.net/browse/CASCL-1274

Co-authored-by: alexander.yastrebov <alexander.yastrebov@datadoghq.com>

Author: AlexanderYastrebov

pkg/clusteragent/autoscaling/cluster/spot/cluster_test.go

@@ -77,6 +77,7 @@ type fakeDeployment struct {
 	name               string
 	existingReplicaSet string
 	podSelector        map[string]string
+	replicas           int
 }
 
 // newFakeCluster creates a fakeCluster.
@@ -135,8 +136,8 @@ func (c *fakeCluster) OnPodDeleted(hook deletionHook) {
 	c.podDeletedHooks = append(c.podDeletedHooks, hook)
 }
 
-// CreatePod runs all registered admission hooks on the pod then creates it as Pending.
-func (c *fakeCluster) CreatePod(pod *corev1.Pod) {
+// createPod runs all registered admission hooks on the pod then creates it as Pending.
+func (c *fakeCluster) createPod(pod *corev1.Pod) {
 	unmodifiedCopy := pod.DeepCopy()
 	for _, hook := range c.podCreatedHooks {
 		updated, err := hook(pod)
@@ -218,9 +219,10 @@ func (c *fakeCluster) T() *testing.T {
 	return c.t
 }
 
-// WLM returns the underlying workloadmeta mock store.
-func (c *fakeCluster) WLM() workloadmetamock.Mock {
-	return c.wlm
+// WLM returns the workloadmeta component used by the scheduler.
+func (c *fakeCluster) WLM() workloadmeta.Component {
+	// Delay updates delivery to expose race conditions
+	return newDelayedWLM(c.wlm, 50*time.Millisecond)
 }
 
 func (c *fakeCluster) DynamicClient() dynamic.Interface {
@@ -416,16 +418,61 @@ func (d *fakeDeployment) rolloutWithDelay(replicas int) string {
 	// A new ReplicaSet created
 	newReplicaSet := replicaSetName(d.name)
 	for range replicas {
-		d.cluster.CreatePod(newPod(d.namespace, kubernetes.ReplicaSetKind, newReplicaSet, d.podSelector))
+		d.cluster.createPod(newPod(d.namespace, kubernetes.ReplicaSetKind, newReplicaSet, d.podSelector))
 	}
 	// Existing ReplicaSet is scaled down
 	if d.existingReplicaSet != "" {
 		d.cluster.DeleteOwnerPods(kubernetes.ReplicaSetKind, d.namespace, d.existingReplicaSet)
 	}
 	d.existingReplicaSet = newReplicaSet
+	d.replicas = replicas
 	return newReplicaSet
 }
 
+// Reconcile creates pods to bring the current ReplicaSet back to d.replicas,
+// counting existing non-terminal pods and creating only the missing number.
+func (d *fakeDeployment) Reconcile() {
+	if d.existingReplicaSet == "" {
+		return
+	}
+	pods := d.cluster.ListOwnerPods(kubernetes.ReplicaSetKind, d.namespace, d.existingReplicaSet)
+	active := 0
+	for _, pod := range pods {
+		phase := corev1.PodPhase(pod.Phase)
+		if phase != corev1.PodSucceeded && phase != corev1.PodFailed {
+			active++
+		}
+	}
+	for range max(0, d.replicas-active) {
+		d.cluster.createPod(newPod(d.namespace, kubernetes.ReplicaSetKind, d.existingReplicaSet, d.podSelector))
+	}
+}
+
+// ScaleDown deletes pods selected by deleteFilter and updates the replica count to the number of remaining pods.
+func (d *fakeDeployment) ScaleDown(deleteFilter func([]*workloadmeta.KubernetesPod) []*workloadmeta.KubernetesPod) {
+	t := d.cluster.t
+	t.Helper()
+	rs := d.ReplicaSet()
+	pods := d.cluster.ListOwnerPods(kubernetes.ReplicaSetKind, d.namespace, rs)
+	toDelete := deleteFilter(pods)
+
+	d.replicas = len(pods) - len(toDelete)
+
+	deleted := make(map[string]struct{}, len(toDelete))
+	for _, pod := range toDelete {
+		d.cluster.DeletePod(pod)
+		deleted[pod.ID] = struct{}{}
+	}
+	require.Eventually(t, func() bool {
+		for _, pod := range d.cluster.ListOwnerPods(kubernetes.ReplicaSetKind, d.namespace, rs) {
+			if _, ok := deleted[pod.ID]; ok {
+				return false
+			}
+		}
+		return true
+	}, 5*time.Second, 100*time.Millisecond)
+}
+
 func async(f func(workloadmeta.Entity), e workloadmeta.Entity) {
 	go func() {
 		time.Sleep(time.Duration(10+rand.N(40)) * time.Millisecond)

pkg/clusteragent/autoscaling/cluster/spot/delayed_wlm_test.go

@@ -0,0 +1,87 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2026-present Datadog, Inc.
+
+//go:build kubeapiserver
+
+package spot_test
+
+import (
+	"sync"
+	"time"
+
+	workloadmeta "github.com/DataDog/datadog-agent/comp/core/workloadmeta/def"
+)
+
+// delayedWLM wraps a workloadmeta.Component and introduces a fixed delay before
+// delivering each event bundle to subscribers.
+type delayedWLM struct {
+	workloadmeta.Component
+	delay time.Duration
+
+	mu       sync.Mutex
+	channels map[chan workloadmeta.EventBundle]delayedSubscription
+}
+
+type delayedSubscription struct {
+	realCh chan workloadmeta.EventBundle
+	done   chan struct{}
+}
+
+func newDelayedWLM(component workloadmeta.Component, delay time.Duration) *delayedWLM {
+	return &delayedWLM{
+		Component: component,
+		delay:     delay,
+		channels:  make(map[chan workloadmeta.EventBundle]delayedSubscription),
+	}
+}
+
+// Subscribe returns a channel that receives the same event bundles as the real
+// WLM channel but delayed by d.delay.
+func (d *delayedWLM) Subscribe(name string, priority workloadmeta.SubscriberPriority, filter *workloadmeta.Filter) chan workloadmeta.EventBundle {
+	realCh := d.Component.Subscribe(name, priority, filter)
+	wrappedCh := make(chan workloadmeta.EventBundle, 100)
+	done := make(chan struct{})
+
+	d.mu.Lock()
+	d.channels[wrappedCh] = delayedSubscription{realCh: realCh, done: done}
+	d.mu.Unlock()
+
+	go func() {
+		defer close(wrappedCh)
+		for {
+			select {
+			case bundle, ok := <-realCh:
+				if !ok {
+					return
+				}
+				time.Sleep(d.delay)
+				select {
+				case wrappedCh <- bundle:
+				case <-done:
+					return
+				}
+			case <-done:
+				return
+			}
+		}
+	}()
+
+	return wrappedCh
+}
+
+// Unsubscribe stops the forwarding goroutine and unsubscribes from the real WLM.
+func (d *delayedWLM) Unsubscribe(ch chan workloadmeta.EventBundle) {
+	d.mu.Lock()
+	sub, ok := d.channels[ch]
+	if ok {
+		delete(d.channels, ch)
+	}
+	d.mu.Unlock()
+
+	if ok {
+		close(sub.done)
+		d.Component.Unsubscribe(sub.realCh)
+	}
+}

pkg/clusteragent/autoscaling/cluster/spot/pod_tracker.go

@@ -76,8 +76,10 @@ func (t *podTracker) admitNewPod(o podOwnership) bool {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
-	ps := t.getOrCreatePodSetLocked(o)
-	t.refreshConfigLocked(o.topLevelOwner, ps)
+	ps, ok := t.getPodSetLocked(o)
+	if !ok {
+		return false
+	}
 
 	total := ps.totalCount()
 	spot := ps.spotCount()
@@ -103,7 +105,9 @@ func (t *podTracker) admitNewOnDemandPod(o podOwnership) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
-	t.getOrCreatePodSetLocked(o).admit(false)
+	if ps, ok := t.getPodSetLocked(o); ok {
+		ps.admit(false)
+	}
 }
 
 // addedOrUpdated updates tracking state when a pod is added or updated.
@@ -128,7 +132,11 @@ func (t *podTracker) addedOrUpdated(pod *workloadmeta.KubernetesPod) {
 		return
 	}
 
-	t.getOrCreatePodSetLocked(o).track(pod.ID, isSpot, podInfo{name: pod.Name, phase: pod.Phase}, t.clock.Now())
+	ps, ok := t.getPodSetLocked(o)
+	if !ok {
+		return
+	}
+	ps.track(pod.ID, isSpot, podInfo{name: pod.Name, phase: pod.Phase}, t.clock.Now())
 
 	if isSpot {
 		if pod.Phase == string(corev1.PodPending) {
@@ -182,28 +190,31 @@ func (t *podTracker) deletePodLocked(o podOwnership, uid string) {
 	delete(t.pendingSpotPods, uid)
 }
 
-// refreshConfigLocked refreshes the spot config for the ownerPodSet from the configSource.
+// getPodSetLocked returns the ownerPodSet for the given ownership, creating it if absent,
+// and always applies the latest config from configSource to the returned pod set.
+// When configSource has no entry for the top-level owner it removes
+// tracking state for the owner and returns (nil, false).
 // Must be called with t.mu held.
-func (t *podTracker) refreshConfigLocked(topLevelOwner objectRef, ps *ownerPodSet) {
-	if cfg, ok := t.configSource(topLevelOwner); ok {
-		ps.config = cfg
+func (t *podTracker) getPodSetLocked(o podOwnership) (*ownerPodSet, bool) {
+	cfg, hasConfig := t.configSource(o.topLevelOwner)
+	if !hasConfig {
+		t.untrackLocked(o.topLevelOwner)
+		return nil, false
 	}
-}
 
-// getOrCreatePodSetLocked returns the ownerPodSet for the given ownership, creating it if absent.
-// Must be called with t.mu held.
-func (t *podTracker) getOrCreatePodSetLocked(o podOwnership) *ownerPodSet {
 	owners, ok := t.podSets[o.topLevelOwner]
 	if !ok {
 		owners = make(map[objectRef]*ownerPodSet)
 		t.podSets[o.topLevelOwner] = owners
 	}
-	if ps, ok := owners[o.directOwner]; ok {
-		return ps
+
+	ps, exists := owners[o.directOwner]
+	if !exists {
+		ps = t.newOwnerPodSet()
+		owners[o.directOwner] = ps
 	}
-	ps := t.newOwnerPodSet()
-	owners[o.directOwner] = ps
-	return ps
+	ps.config = cfg
+	return ps, true
 }
 
 // getPodToDelete returns the uid, name, and namespace of a pod to delete to make progress toward
@@ -217,8 +228,13 @@ func (t *podTracker) getPodToDelete(rebalanceStabilizationPeriod time.Duration)
 	now := t.clock.Now()
 	lastUpdatedBefore := now.Add(-rebalanceStabilizationPeriod)
 	for topLevel, owners := range t.podSets {
+		cfg, ok := t.configSource(topLevel)
+		if !ok {
+			t.untrackLocked(topLevel)
+			continue
+		}
 		for owner, ps := range owners {
-			t.refreshConfigLocked(topLevel, ps)
+			ps.config = cfg
 			if ps.config.isDisabled(now) {
 				continue
 			}
@@ -256,7 +272,12 @@ func (t *podTracker) deletePendingSpotPod(uid string) {
 func (t *podTracker) untrack(topLevelOwner objectRef) {
 	t.mu.Lock()
 	defer t.mu.Unlock()
+	t.untrackLocked(topLevelOwner)
+}
 
+// untrackLocked removes all tracking state for the given top-level owner.
+// Must be called with t.mu held.
+func (t *podTracker) untrackLocked(topLevelOwner objectRef) {
 	delete(t.podSets, topLevelOwner)
 	for uid, p := range t.pendingSpotPods {
 		if p.topLevelOwner == topLevelOwner {
@@ -306,7 +327,7 @@ func (ps *ownerPodSet) track(uid string, isSpot bool, info podInfo, now time.Tim
 // getPodToDelete returns the uid and name of a pod to delete to make progress toward the desired config.
 // It returns empty strings if no deletion is needed.
 func (ps *ownerPodSet) getPodToDelete(lastUpdatedBefore time.Time) (string, string) {
-	if ps.admissionSpotCount > 0 || ps.admissionOnDemandCount > 0 {
+	if ps.hasAdmissions() {
 		return "", ""
 	}
 
@@ -340,6 +361,11 @@ func (ps *ownerPodSet) getPodToDelete(lastUpdatedBefore time.Time) (string, stri
 	return "", ""
 }
 
+// hasAdmissions returns true if pod set has admitted but not yet tracked pods.
+func (ps *ownerPodSet) hasAdmissions() bool {
+	return ps.admissionSpotCount > 0 || ps.admissionOnDemandCount > 0
+}
+
 // hasPending returns true if any tracked pod is in PodPending phase.
 func (ps *ownerPodSet) hasPending() bool {
 	const pending = string(corev1.PodPending)