From 998a3fe5051f05d1482a0db16831da18a70dc157 Mon Sep 17 00:00:00 2001
From: Jay Shane <327411586@qq.com>
Date: Tue, 26 Nov 2024 09:41:19 +0800
Subject: [PATCH] improve the nodeclaim sorting algorithm using heap

Signed-off-by: Jay Shane <327411586@qq.com>
---
 go.mod                                        |  4 ++
 .../provisioning/scheduling/heap.go           | 46 +++++++++++++++++
 .../provisioning/scheduling/heap_test.go      | 50 +++++++++++++++++++
 .../provisioning/scheduling/scheduler.go      | 10 ++--
 4 files changed, 105 insertions(+), 5 deletions(-)
 create mode 100644 pkg/controllers/provisioning/scheduling/heap.go
 create mode 100644 pkg/controllers/provisioning/scheduling/heap_test.go

diff --git a/go.mod b/go.mod
index e83aef1437..e43eda056e 100644
--- a/go.mod
+++ b/go.mod
@@ -86,6 +86,10 @@ require (
 	sigs.k8s.io/yaml v1.4.0 // indirect
 )
 
+require github.com/stretchr/testify v1.9.0
+
+require github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+
 retract (
 	v0.100.101-test // accidentally published testing version
 	v0.35.3 // accidentally published incomplete patch release
diff --git a/pkg/controllers/provisioning/scheduling/heap.go b/pkg/controllers/provisioning/scheduling/heap.go
new file mode 100644
index 0000000000..5ccfe10079
--- /dev/null
+++ b/pkg/controllers/provisioning/scheduling/heap.go
@@ -0,0 +1,46 @@
+/*
+Copyright The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package scheduling
+
+import "container/heap"
+
+type NodeClaimHeap []*NodeClaim
+
+var (
+	_ = heap.Interface(&NodeClaimHeap{}) // NodeClaimHeap is a standard heap
+)
+
+func NewNodeClaimHeap(nodeClaims []*NodeClaim) *NodeClaimHeap {
+	h := NodeClaimHeap(nodeClaims)
+	return &h
+}
+
+func (h NodeClaimHeap) Len() int           { return len(h) }
+func (h NodeClaimHeap) Less(i, j int) bool { return len(h[i].Pods) < len(h[j].Pods) }
+func (h NodeClaimHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+
+func (h *NodeClaimHeap) Push(x interface{}) {
+	*h = append(*h, x.(*NodeClaim))
+}
+
+func (h *NodeClaimHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	item := old[n-1]
+	*h = old[0 : n-1]
+	return item
+}
diff --git a/pkg/controllers/provisioning/scheduling/heap_test.go b/pkg/controllers/provisioning/scheduling/heap_test.go
new file mode 100644
index 0000000000..05b32a3fd9
--- /dev/null
+++ b/pkg/controllers/provisioning/scheduling/heap_test.go
@@ -0,0 +1,50 @@
+/*
+Copyright The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package scheduling
+
+import (
+	"container/heap"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+)
+
+func TestNodeClaimHeap_PopOrder(t *testing.T) {
+	// Create NodeClaims with different pod counts
+	nc1 := &NodeClaim{Pods: []*corev1.Pod{{}, {}}}     // 2 pods
+	nc2 := &NodeClaim{Pods: []*corev1.Pod{{}}}         // 1 pod
+	nc3 := &NodeClaim{Pods: []*corev1.Pod{{}, {}, {}}} // 3 pods
+	nc4 := &NodeClaim{Pods: []*corev1.Pod{}}           // 0 pods
+
+	// Initialize heap with NodeClaims
+	h := NewNodeClaimHeap([]*NodeClaim{nc1, nc2, nc3, nc4})
+	heap.Init(h)
+
+	// Pop items and verify they come out in ascending order of pod count
+	expected := []*NodeClaim{nc4, nc2, nc1, nc3}
+
+	for i := 0; i < len(expected); i++ {
+		item := heap.Pop(h).(*NodeClaim)
+		assert.Equal(t, len(expected[i].Pods), len(item.Pods),
+			"Expected NodeClaim with %d pods, got %d pods",
+			len(expected[i].Pods), len(item.Pods))
+	}
+
+	// Verify heap is empty
+	assert.Equal(t, 0, h.Len(), "Heap should be empty after popping all items")
+}
diff --git a/pkg/controllers/provisioning/scheduling/scheduler.go b/pkg/controllers/provisioning/scheduling/scheduler.go
index dfacff6936..3fd3d94900 100644
--- a/pkg/controllers/provisioning/scheduling/scheduler.go
+++ b/pkg/controllers/provisioning/scheduling/scheduler.go
@@ -18,6 +18,7 @@ package scheduling
 
 import (
 	"bytes"
+	"container/heap"
 	"context"
 	"fmt"
 	"sort"
@@ -272,11 +273,10 @@ func (s *Scheduler) add(ctx context.Context, pod *corev1.Pod) error {
 		}
 	}
 
-	// Consider using https://pkg.go.dev/container/heap
-	sort.Slice(s.newNodeClaims, func(a, b int) bool { return len(s.newNodeClaims[a].Pods) < len(s.newNodeClaims[b].Pods) })
-
-	// Pick existing node that we are about to create
-	for _, nodeClaim := range s.newNodeClaims {
+	h := NewNodeClaimHeap(s.newNodeClaims)
+	heap.Init(h)
+	for h.Len() > 0 {
+		nodeClaim := heap.Pop(h).(*NodeClaim)
 		if err := nodeClaim.Add(pod, s.cachedPodRequests[pod.UID]); err == nil {
 			return nil
 		}