Repository: kubernetes-sigs/controller-runtime

Issue:
Priority queue: priority not always respected
I suspect that there may be conditions under which the item returned by `GetWithPriority()` is not always the correct/expected element.
Both `GetWithPriority()`  and  `AddWithOpts` trigger a `spin()`: https://github.com/kubernetes-sigs/controller-runtime/blob/fb2beabd029b3d083b453ded8446cd43366c0db9/pkg/controller/priorityqueue/priorityqueue.go#L206

In the `GetWithPriority()`, the `w.waiters` are raised `w.waiters.Add(1)`, and in the `spin` items are skipped if none is waiting:
```golang
				if w.waiters.Load() == 0 {
					// Have to keep iterating here to ensure we update metrics
					// for further items that became ready and set nextReady.
					return true
				}
```

If a `GetWithPriority()` is executed right after an `AddWithOpts`, there could be the possibility that the `Ascend` of the Btree is not completed yet, causing not all the items of the BTree to evaluate the same `w.waiters.Load()` value.

This will result in the non-obtaining of the higher-priority element.

In particular, you can see that this test https://github.com/kubevirt/kubevirt/blob/66f09757e8707ae4a09c660119ec1ab3658b0cdb/pkg/virt-controller/watch/migration/migration_test.go#L2342 is flaky (a failure example: https://prow.ci.kubevirt.io/view/gs/kubevirt-prow/pr-logs/pull/kubevirt_kubevirt/15887/pull-kubevirt-unit-test/1981424912933326848).

A possible solution could be adding a lock before increasing the waiters:
https://github.com/kubernetes-sigs/controller-runtime/blob/fb2beabd029b3d083b453ded8446cd43366c0db9/pkg/controller/priorityqueue/priorityqueue.go#L313-L319

might be:
```golang
 func (w *priorityqueue[T]) GetWithPriority() (_ T, priority int, shutdown bool) { 
 	if w.shutdown.Load() { 
 		var zero T 
 		return zero, 0, true 
 	} 
    w.lock.Lock()
 	w.waiters.Add(1) 
    w.lock.Unlock()
```

What do you think?
Thank you!


Edit: Checked the above solution. The flaky disappeared, but IDK the implications

Hints from issue comments:
@alvaroaleman: Acquiring the lock in `GetWithPriority` is fine I think as it is blocking anyways. No need to use an atomic then anymore. It would be good to have a test for this, as this is the kind of thing that is easy to regress in the future

Solution patch:
diff --git a/pkg/controller/priorityqueue/priorityqueue.go b/pkg/controller/priorityqueue/priorityqueue.go
index 98df84c5..71363f0d 100644
--- a/pkg/controller/priorityqueue/priorityqueue.go
+++ b/pkg/controller/priorityqueue/priorityqueue.go
@@ -124,8 +124,8 @@ type priorityqueue[T comparable] struct {
 	get chan item[T]
 
 	// waiters is the number of routines blocked in Get, we use it to determine
-	// if we can push items.
-	waiters atomic.Int64
+	// if we can push items. Every manipulation has to be protected with the lock.
+	waiters int64
 
 	// Configurable for testing
 	now  func() time.Time
@@ -269,7 +269,7 @@ func (w *priorityqueue[T]) spin() {
 						}
 					}
 
-					if w.waiters.Load() == 0 {
+					if w.waiters == 0 {
 						// Have to keep iterating here to ensure we update metrics
 						// for further items that became ready and set nextReady.
 						return true
@@ -277,7 +277,7 @@ func (w *priorityqueue[T]) spin() {
 
 					w.metrics.get(item.Key, item.Priority)
 					w.locked.Insert(item.Key)
-					w.waiters.Add(-1)
+					w.waiters--
 					delete(w.items, item.Key)
 					toDelete = append(toDelete, item)
 					w.becameReady.Delete(item.Key)
@@ -316,7 +316,9 @@ func (w *priorityqueue[T]) GetWithPriority() (_ T, priority int, shutdown bool)
 		return zero, 0, true
 	}
 
-	w.waiters.Add(1)
+	w.lock.Lock()
+	w.waiters++
+	w.lock.Unlock()
 
 	w.notifyItemOrWaiterAdded()
 



Test patch:
diff --git a/pkg/controller/priorityqueue/priorityqueue_test.go b/pkg/controller/priorityqueue/priorityqueue_test.go
index fb186944..9c708e98 100644
--- a/pkg/controller/priorityqueue/priorityqueue_test.go
+++ b/pkg/controller/priorityqueue/priorityqueue_test.go
@@ -441,6 +441,60 @@ func newQueueWithTimeForwarder() (_ *priorityqueue[string], _ *fakeMetricsProvid
 	}
 }
 
+func TestHighPriorityItemsAreReturnedBeforeLowPriorityItemMultipleTimes(t *testing.T) {
+	t.Parallel()
+	synctest.Test(t, func(t *testing.T) {
+		g := NewWithT(t)
+
+		q, metrics := newQueue()
+		defer q.ShutDown()
+
+		const itemsPerPriority = 1000
+		lowPriority := 0
+		lowMiddlePriority := 5
+		middlePriority := 10
+		upperMiddlePriority := 15
+		highPriority := 20
+		for i := range itemsPerPriority {
+			q.AddWithOpts(AddOpts{Priority: &highPriority}, fmt.Sprintf("high-%d", i))
+			q.AddWithOpts(AddOpts{Priority: &upperMiddlePriority}, fmt.Sprintf("upperMiddle-%d", i))
+			q.AddWithOpts(AddOpts{Priority: &middlePriority}, fmt.Sprintf("middle-%d", i))
+			q.AddWithOpts(AddOpts{Priority: &lowMiddlePriority}, fmt.Sprintf("lowMiddle-%d", i))
+			q.AddWithOpts(AddOpts{Priority: &lowPriority}, fmt.Sprintf("low-%d", i))
+		}
+		synctest.Wait()
+
+		for range itemsPerPriority {
+			key, prio, _ := q.GetWithPriority()
+			g.Expect(prio).To(Equal(highPriority))
+			g.Expect(key).To(HavePrefix("high-"))
+		}
+		for range itemsPerPriority {
+			key, prio, _ := q.GetWithPriority()
+			g.Expect(prio).To(Equal(upperMiddlePriority))
+			g.Expect(key).To(HavePrefix("upperMiddle-"))
+		}
+		for range itemsPerPriority {
+			key, prio, _ := q.GetWithPriority()
+			g.Expect(prio).To(Equal(middlePriority))
+			g.Expect(key).To(HavePrefix("middle-"))
+		}
+		for range itemsPerPriority {
+			key, prio, _ := q.GetWithPriority()
+			g.Expect(prio).To(Equal(lowMiddlePriority))
+			g.Expect(key).To(HavePrefix("lowMiddle-"))
+		}
+		for range itemsPerPriority {
+			key, prio, _ := q.GetWithPriority()
+			g.Expect(prio).To(Equal(lowPriority))
+			g.Expect(key).To(HavePrefix("low-"))
+		}
+		g.Expect(metrics.depth["test"]).To(Equal(map[int]int{10: 0, 5: 0, 0: 0, 20: 0, 15: 0}))
+		g.Expect(metrics.adds["test"]).To(Equal(itemsPerPriority * 5))
+		g.Expect(metrics.retries["test"]).To(Equal(0))
+	})
+}
+
 func newQueue() (*priorityqueue[string], *fakeMetricsProvider) {
 	metrics := newFakeMetricsProvider()
 	q := New("test", func(o *Opts[string]) {