Ensure consistent preemptor state across queue heads and second pass

pkg/cache/queue/cluster_queue.go

@@ -655,6 +655,15 @@ func (c *ClusterQueue) PendingInLocalQueue(lqRef utilqueue.LocalQueueReference)
 // Pop removes the head of the queue and returns it. It returns nil if the
 // queue is empty.
 func (c *ClusterQueue) Pop() *workload.Info {
+	if head := c.PopHead(); head != nil {
+		return &head.Info
+	}
+	return nil
+}
+
+// PopHead removes the head of the queue and returns it along with a flag
+// indicating whether it is a preemptor.
+func (c *ClusterQueue) PopHead() *Head {
 	c.rwm.Lock()
 	defer c.rwm.Unlock()
 
@@ -668,7 +677,15 @@ func (c *ClusterQueue) Pop() *workload.Info {
 	}
 
 	c.popCycle++
-	return c.workloads.PopActive()
+	wl := c.workloads.PopActive()
+	if wl == nil {
+		return nil
+	}
+
+	return &Head{
+		Info:        *wl,
+		IsPreemptor: c.IsPreemptor(wl),
+	}
 }
 
 func (c *ClusterQueue) hasPendingPenalties() bool {

pkg/cache/queue/cluster_queue_test.go

@@ -2390,3 +2390,79 @@ func TestClusterQueuePendingTrackers(t *testing.T) {
 		})
 	}
 }
+
+// TestPopHead ensures the right workload is popped and its preemptor status is
+// correctly identified.
+func TestPopHead(t *testing.T) {
+	ctx, _ := utiltesting.ContextWithLog(t)
+	cq, err := newClusterQueue(ctx, nil, utiltestingapi.MakeClusterQueue("cq").Obj(), nil, defaultOrdering, nil, nil)
+	if err != nil {
+		t.Fatalf("failed to create ClusterQueue: %v", err)
+	}
+
+	wl := utiltestingapi.MakeWorkload("wl", defaultNamespace).Obj()
+	wl.Generation = 1
+	wInfo := workload.NewInfo(wl)
+	wInfo.LastEvaluatedGeneration = 1
+
+	cq.PushOrUpdate(wInfo)
+	cq.RequeueIfNotPresent(ctx, wInfo, RequeueReasonPendingPreemption, "")
+
+	head := cq.PopHead()
+	if head == nil || !head.IsPreemptor {
+		t.Errorf("Unexpected popped head: %v", head)
+	}
+}
+
+// TestPopHeadConcurrentWithRequeue ensures PopHead can safely run concurrently with preemption requeues.
+func TestPopHeadConcurrentWithRequeue(t *testing.T) {
+	ctx, _ := utiltesting.ContextWithLog(t)
+	cq, err := newClusterQueue(ctx, nil, utiltestingapi.MakeClusterQueue("cq").QueueingStrategy(kueue.BestEffortFIFO).Obj(), nil, defaultOrdering, nil, nil)
+	if err != nil {
+		t.Fatalf("failed to create ClusterQueue: %v", err)
+	}
+
+	wInfo := workload.NewInfo(utiltestingapi.MakeWorkload("wl", defaultNamespace).Obj())
+
+	stop := make(chan struct{})
+	started := make(chan struct{})
+	done := make(chan struct{})
+
+	// Continuously requeue the workload with pending preemption in a background goroutine.
+	go func() {
+		defer close(done)
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				cq.RequeueIfNotPresent(ctx, wInfo, RequeueReasonPendingPreemption, "")
+				select {
+				case <-started:
+				default:
+					// Signal that the background goroutine has performed at least one requeue
+					close(started)
+				}
+			}
+		}
+	}()
+	// Wait until the background goroutine has performed its first requeue
+	<-started
+
+	// Concurrently pop heads and verify that any returned head correctly reflects
+	// the pending preemptor state
+	for range 1000 {
+		head := cq.PopHead()
+		if head != nil {
+			if !head.IsPreemptor {
+				t.Errorf("PopHead returned non-preemptor head for workload with pending preemption: %v", head)
+			}
+			if workload.Key(head.Obj) != workload.Key(wInfo.Obj) {
+				t.Errorf("PopHead returned unexpected workload: %v, want %v", head.Obj, wInfo.Obj)
+			}
+		}
+	}
+	// Signal the writer goroutine to stop and wait for it to exit
+	close(stop)
+	<-done
+}

pkg/cache/queue/manager.go

@@ -919,30 +919,35 @@ func (m *Manager) Heads(ctx context.Context) []Head {
 	}
 }
 
+// heads returns the heads of the queues and ready second-pass workloads.
 func (m *Manager) heads() []Head {
-	heads := m.secondPassQueue.takeAllReady()
+	var heads []Head
+	for wInfo := range m.secondPassQueue.takeAllReady() {
+		isPreemptor := false
+		if cq := m.getClusterQueueLockless(wInfo.ClusterQueue); cq != nil {
+			isPreemptor = cq.IsPreemptor(&wInfo)
+		}
+		heads = append(heads, Head{
+			Info:        wInfo,
+			IsPreemptor: isPreemptor,
+		})
+	}
 	for cqName, cq := range m.hm.ClusterQueues() {
 		// Cache might be nil in tests, if cache is nil, we'll skip the check.
 		if m.statusChecker != nil && !m.statusChecker.ClusterQueueActive(cqName) {
 			continue
 		}
-		wl := cq.Pop()
+		head := cq.PopHead()
 		reportCQPendingWorkloads(m, cq)
-		if wl == nil {
+		if head == nil {
 			continue
 		}
-		wlKey := workload.Key(wl.Obj)
-		wlCopy := *wl
-		wlCopy.ClusterQueue = cqName
-		heads = append(heads, Head{
-			Info:        wlCopy,
-			IsPreemptor: cq.IsPreemptor(wl),
-		})
-
+		head.ClusterQueue = cqName
+		heads = append(heads, *head)
+		wlKey := workload.Key(head.Obj)
 		qKey := m.workloadAssignedQueues[wlKey]
 		q := m.localQueues[qKey]
 		delete(q.items, wlKey)
-
 		reportLQPendingWorkloads(m, q)
 	}
 	return heads

pkg/cache/queue/manager_test.go

@@ -2183,7 +2183,7 @@ func TestQueueSecondPassIfNeeded(t *testing.T) {
 			fakeClock.Step(tc.passTime)
 
 			gotReady := sets.New[workload.Reference]()
-			for _, head := range manager.secondPassQueue.takeAllReady() {
+			for head := range manager.secondPassQueue.takeAllReady() {
 				gotReady.Insert(workload.Key(head.Obj))
 			}
 
@@ -2194,6 +2194,64 @@ func TestQueueSecondPassIfNeeded(t *testing.T) {
 	}
 }
 
+// TestSecondPassQueueIsPreemptor verifies that workloads queued in the second-pass queue retain their preemptor status.
+func TestSecondPassQueueIsPreemptor(t *testing.T) {
+	ctx, _ := utiltesting.ContextWithLog(t)
+	cq := utiltestingapi.MakeClusterQueue("cq").QueueingStrategy(kueue.BestEffortFIFO).Obj()
+	lq := utiltestingapi.MakeLocalQueue("lq", "ns").ClusterQueue("cq").Obj()
+
+	baseWorkloadBuilder := utiltestingapi.MakeWorkload("wl", "ns").
+		Queue("lq").
+		PodSets(*utiltestingapi.MakePodSet("one", 1).
+			RequiredTopologyRequest(corev1.LabelHostname).
+			Request(corev1.ResourceCPU, "1").
+			Obj())
+
+	wl := baseWorkloadBuilder.Clone().
+		ReserveQuotaAt(
+			utiltestingapi.MakeAdmission("cq").
+				PodSets(
+					utiltestingapi.MakePodSetAssignment("one").
+						Assignment(corev1.ResourceCPU, "tas-default", "1000m").
+						DelayedTopologyRequest(kueue.DelayedTopologyRequestStatePending).
+						Obj(),
+				).
+				Obj(),
+			time.Now(),
+		).
+		AdmissionCheck(kueue.AdmissionCheckState{
+			Name:  "prov-check",
+			State: kueue.CheckStateReady,
+		}).Obj()
+
+	fakeClock := testingclock.NewFakeClock(time.Now())
+	kClient := utiltesting.NewFakeClient(lq, cq, wl)
+	manager := NewManagerForUnitTests(kClient, nil, WithClock(fakeClock), WithPreemptionExpectations(preemptexpectations.New()))
+	if err := manager.AddClusterQueue(ctx, cq); err != nil {
+		t.Fatalf("Failed adding clusterQeueu: %v", err)
+	}
+	if err := manager.AddLocalQueue(ctx, lq); err != nil {
+		t.Fatalf("Failed adding localQueue: %v", err)
+	}
+	wInfo := workload.NewInfo(wl)
+	manager.getClusterQueue(kueue.ClusterQueueReference("cq")).RequeueIfNotPresent(ctx, wInfo, RequeueReasonPendingPreemption, "")
+
+	_ = manager.Heads(ctx)
+	got := manager.QueueSecondPassIfNeeded(ctx, wl, 0)
+	if !got {
+		t.Errorf("Expected workload to be queued")
+	}
+	fakeClock.Step(time.Second)
+
+	heads := manager.Heads(ctx)
+	if len(heads) != 1 {
+		t.Errorf("Expected 1 head, got %d", len(heads))
+	}
+	if !heads[0].IsPreemptor {
+		t.Errorf("Expected second pass workload to be IsPreemptor = true, got false")
+	}
+}
+
 func TestUpdateUnadmittedWorkload(t *testing.T) {
 	ctx, log := utiltesting.ContextWithLog(t)
 	cq := utiltestingapi.MakeClusterQueue("cq").Obj()

pkg/cache/queue/second_pass_queue.go

@@ -17,6 +17,7 @@ limitations under the License.
 package queue
 
 import (
+	"iter"
 	"sync"
 	"time"
 
@@ -50,16 +51,19 @@ func newSecondPassQueue() *secondPassQueue {
 	}
 }
 
-func (q *secondPassQueue) takeAllReady() []Head {
+// takeAllReady removes and returns all workloads currently queued for the second pass.
+func (q *secondPassQueue) takeAllReady() iter.Seq[workload.Info] {
 	q.Lock()
 	defer q.Unlock()
-
-	var result []Head
-	for _, v := range q.queued {
-		result = append(result, Head{Info: *v})
-	}
+	queued := q.queued
 	q.queued = make(map[workload.Reference]*workload.Info)
-	return result
+	return func(yield func(workload.Info) bool) {
+		for _, v := range queued {
+			if !yield(*v) {
+				return
+			}
+		}
+	}
 }
 
 func (q *secondPassQueue) prequeueIfAbsent(obj *kueue.Workload) bool {