Add lock to TestAsyncPreemption to prevent races

macsko · macsko · commit 1024b21ea115 · 2025-06-30T13:05:33.000Z
diff --git a/pkg/scheduler/backend/queue/scheduling_queue.go b/pkg/scheduler/backend/queue/scheduling_queue.go
@@ -629,11 +629,11 @@ func (p *PriorityQueue) moveToActiveQ(logger klog.Logger, pInfo *framework.Queue
 			pInfo.InitialAttemptTimestamp = &now
 		}
 
-		unlockedActiveQ.add(pInfo, event)
-		added = true
-
 		p.unschedulablePods.delete(pInfo.Pod, gatedBefore)
 		p.backoffQ.delete(pInfo)
+
+		unlockedActiveQ.add(pInfo, event)
+		added = true
 		logger.V(5).Info("Pod moved to an internal scheduling queue", "pod", klog.KObj(pInfo.Pod), "event", event, "queue", activeQ)
 		if event == framework.EventUnscheduledPodAdd.Label() || event == framework.EventUnscheduledPodUpdate.Label() {
 			p.AddNominatedPod(logger, pInfo.PodInfo, nil)
@@ -660,8 +660,9 @@ func (p *PriorityQueue) moveToBackoffQ(logger klog.Logger, pInfo *framework.Queu
 			return false
 		}
 	}
-	p.backoffQ.add(logger, pInfo, event)
 	p.unschedulablePods.delete(pInfo.Pod, gatedBefore)
+
+	p.backoffQ.add(logger, pInfo, event)
 	logger.V(5).Info("Pod moved to an internal scheduling queue", "pod", klog.KObj(pInfo.Pod), "event", event, "queue", backoffQ)
 	return true
 }
@@ -1008,7 +1009,6 @@ func (p *PriorityQueue) Update(logger klog.Logger, oldPod, newPod *v1.Pod) {
 	if pInfo := p.unschedulablePods.get(newPod); pInfo != nil {
 		_ = pInfo.Update(newPod)
 		p.UpdateNominatedPod(logger, oldPod, pInfo.PodInfo)
-		gated := pInfo.Gated()
 		if p.isSchedulingQueueHintEnabled {
 			// When unscheduled Pods are updated, we check with QueueingHint
 			// whether the update may make the pods schedulable.
@@ -1032,7 +1032,6 @@ func (p *PriorityQueue) Update(logger klog.Logger, oldPod, newPod *v1.Pod) {
 			// so we should check isPodBackingoff before moving the pod to backoffQ.
 			if p.backoffQ.isPodBackingoff(pInfo) {
 				if added := p.moveToBackoffQ(logger, pInfo, framework.EventUnscheduledPodUpdate.Label()); added {
-					p.unschedulablePods.delete(pInfo.Pod, gated)
 					if p.isPopFromBackoffQEnabled {
 						p.activeQ.broadcast()
 					}
diff --git a/test/integration/scheduler/preemption/preemption_test.go b/test/integration/scheduler/preemption/preemption_test.go
@@ -22,6 +22,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -813,9 +814,12 @@ func TestAsyncPreemption(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			// We need to use a custom preemption plugin to test async preemption behavior
 			delayedPreemptionPluginName := "delay-preemption"
+			var lock sync.Mutex
 			// keyed by the pod name
 			preemptionDoneChannels := make(map[string]chan struct{})
 			defer func() {
+				lock.Lock()
+				defer lock.Unlock()
 				for _, ch := range preemptionDoneChannels {
 					close(ch)
 				}
@@ -841,7 +845,10 @@ func TestAsyncPreemption(t *testing.T) {
 				preemptPodFn := preemptionPlugin.Evaluator.PreemptPod
 				preemptionPlugin.Evaluator.PreemptPod = func(ctx context.Context, c preemption.Candidate, preemptor, victim *v1.Pod, pluginName string) error {
 					// block the preemption goroutine to complete until the test case allows it to proceed.
-					if ch, ok := preemptionDoneChannels[preemptor.Name]; ok {
+					lock.Lock()
+					ch, ok := preemptionDoneChannels[preemptor.Name]
+					lock.Unlock()
+					if ok {
 						<-ch
 					}
 					return preemptPodFn(ctx, c, preemptor, victim, pluginName)
@@ -941,7 +948,9 @@ func TestAsyncPreemption(t *testing.T) {
 						t.Fatal(lastFailure)
 					}
 
+					lock.Lock()
 					preemptionDoneChannels[scenario.schedulePod.podName] = make(chan struct{})
+					lock.Unlock()
 					testCtx.Scheduler.ScheduleOne(testCtx.Ctx)
 					if scenario.schedulePod.expectSuccess {
 						if err := wait.PollUntilContextTimeout(testCtx.Ctx, 200*time.Millisecond, wait.ForeverTestTimeout, false, testutils.PodScheduled(cs, testCtx.NS.Name, scenario.schedulePod.podName)); err != nil {
@@ -953,12 +962,14 @@ func TestAsyncPreemption(t *testing.T) {
 						}
 					}
 				case scenario.completePreemption != "":
+					lock.Lock()
 					if _, ok := preemptionDoneChannels[scenario.completePreemption]; !ok {
 						t.Fatalf("The preemptor Pod %q is not running preemption", scenario.completePreemption)
 					}
 
 					close(preemptionDoneChannels[scenario.completePreemption])
 					delete(preemptionDoneChannels, scenario.completePreemption)
+					lock.Unlock()
 				case scenario.podGatedInQueue != "":
 					// make sure the Pod is in the queue in the first place.
 					if !podInUnschedulablePodPool(t, testCtx.Scheduler.SchedulingQueue, scenario.podGatedInQueue) {
