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

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

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: AlexanderYastrebov

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 {