From 1f157bcb90a4e4908e55109998263bb5c9a0f6d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20Skocze=C5=84?= <mskoczen@google.com>
Date: Mon, 2 Sep 2024 12:12:19 +0000
Subject: [PATCH] Don't panic when popping from empty scheduling queue

---
 pkg/scheduler/backend/heap/heap.go      | 6 ++++++
 pkg/scheduler/backend/heap/heap_test.go | 5 +++++
 2 files changed, 11 insertions(+)

diff --git a/pkg/scheduler/backend/heap/heap.go b/pkg/scheduler/backend/heap/heap.go
index df72db1e13e..182954e2300 100644
--- a/pkg/scheduler/backend/heap/heap.go
+++ b/pkg/scheduler/backend/heap/heap.go
@@ -85,6 +85,9 @@ func (h *data[T]) Len() int { return len(h.queue) }
 // Swap implements swapping of two elements in the heap. This is a part of standard
 // heap interface and should never be called directly.
 func (h *data[T]) Swap(i, j int) {
+	if i < 0 || j < 0 {
+		return
+	}
 	h.queue[i], h.queue[j] = h.queue[j], h.queue[i]
 	item := h.items[h.queue[i]]
 	item.index = i
@@ -102,6 +105,9 @@ func (h *data[T]) Push(kv interface{}) {
 
 // Pop is supposed to be called by container/heap.Pop only.
 func (h *data[T]) Pop() interface{} {
+	if len(h.queue) == 0 {
+		return nil
+	}
 	key := h.queue[len(h.queue)-1]
 	h.queue = h.queue[0 : len(h.queue)-1]
 	item, ok := h.items[key]
diff --git a/pkg/scheduler/backend/heap/heap_test.go b/pkg/scheduler/backend/heap/heap_test.go
index 35e010747d1..24f1123712b 100644
--- a/pkg/scheduler/backend/heap/heap_test.go
+++ b/pkg/scheduler/backend/heap/heap_test.go
@@ -94,6 +94,11 @@ func TestHeapBasic(t *testing.T) {
 		}
 		prevNum = num
 	}
+
+	_, err := h.Pop()
+	if err == nil {
+		t.Errorf("expected Pop() to error on empty heap")
+	}
 }
 
 // TestHeap_AddOrUpdate_Add tests add capabilities of Heap.AddOrUpdate