optimize global strategy

strategy/global.go

@@ -1,17 +1,48 @@
 package strategy
 
 import (
+	"container/heap"
 	"context"
 	"fmt"
 	"sort"
 
+	"github.com/pkg/errors"
+
 	"github.com/projecteru2/core/log"
 	"github.com/projecteru2/core/types"
-	"github.com/projecteru2/core/utils"
-
-	"github.com/pkg/errors"
 )
 
+type infoHeapForGlobalStrategy []Info
+
+// Len .
+func (r infoHeapForGlobalStrategy) Len() int {
+	return len(r)
+}
+
+// Less .
+func (r infoHeapForGlobalStrategy) Less(i, j int) bool {
+	return (r[i].Usage + r[i].Rate) < (r[j].Usage + r[j].Rate)
+}
+
+// Swap .
+func (r infoHeapForGlobalStrategy) Swap(i, j int) {
+	r[i], r[j] = r[j], r[i]
+}
+
+// Push .
+func (r *infoHeapForGlobalStrategy) Push(x interface{}) {
+	*r = append(*r, x.(Info))
+}
+
+// Pop .
+func (r *infoHeapForGlobalStrategy) Pop() interface{} {
+	old := *r
+	n := len(old)
+	x := old[n-1]
+	*r = old[:n-1]
+	return x
+}
+
 // GlobalPlan 基于全局资源配额
 // 尽量使得资源消耗平均
 func GlobalPlan(ctx context.Context, infos []Info, need, total, _ int) (map[string]int, error) {
@@ -22,39 +53,27 @@ func GlobalPlan(ctx context.Context, infos []Info, need, total, _ int) (map[stri
 	strategyInfos := make([]Info, len(infos))
 	copy(strategyInfos, infos)
 	sort.Slice(infos, func(i, j int) bool { return infos[i].Capacity > infos[j].Capacity })
-	length := len(strategyInfos)
-	i := 0
-
-	deployMap := make(map[string]int)
-	for need > 0 {
-		p := i
-		deploy := 0
-		delta := 0.0
-		if i < length-1 {
-			delta = utils.Round(strategyInfos[i+1].Usage - strategyInfos[i].Usage)
-			i++
+	deployMap := map[string]int{}
+
+	infoHeap := &infoHeapForGlobalStrategy{}
+	for _, info := range strategyInfos {
+		if info.Capacity > 0 {
+			infoHeap.Push(info)
 		}
-		for j := 0; j <= p && need > 0 && delta >= 0; j++ {
-			// 减枝
-			if strategyInfos[j].Capacity == 0 {
-				continue
-			}
-			cost := utils.Round(strategyInfos[j].Rate)
-			deploy = int(delta / cost)
-			if deploy == 0 {
-				deploy = 1
-			}
-			if deploy > strategyInfos[j].Capacity {
-				deploy = strategyInfos[j].Capacity
-			}
-			if deploy > need {
-				deploy = need
-			}
-			strategyInfos[j].Capacity -= deploy
-			deployMap[strategyInfos[j].Nodename] += deploy
-			need -= deploy
+	}
+	heap.Init(infoHeap)
+
+	for i := 0; i < need; i++ {
+		infoWithMinUsage := heap.Pop(infoHeap).(Info)
+		deployMap[infoWithMinUsage.Nodename]++
+		infoWithMinUsage.Usage += infoWithMinUsage.Rate
+		infoWithMinUsage.Capacity--
+
+		if infoWithMinUsage.Capacity > 0 {
+			heap.Push(infoHeap, infoWithMinUsage)
 		}
 	}
+
 	// 这里 need 一定会为 0 出来，因为 volTotal 保证了一定大于 need
 	// 这里并不需要再次排序了，理论上的排序是基于资源使用率得到的 Deploy 最终方案
 	log.Debugf(ctx, "[GlobalPlan] strategyInfos: %v", strategyInfos)

strategy/global_test.go

@@ -2,6 +2,7 @@ package strategy
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -32,6 +33,32 @@ func TestGlobalPlan1(t *testing.T) {
 	assert.Equal(t, r[arg[0].Nodename], 1)
 }
 
+func TestGlobalPlan10086(t *testing.T) {
+	n1 := Info{
+		Nodename: "n1",
+		Usage:    0.8,
+		Rate:     0.02,
+		Capacity: 10,
+	}
+	n2 := Info{
+		Nodename: "n2",
+		Usage:    0.7,
+		Rate:     0.033,
+		Capacity: 9,
+	}
+	n3 := Info{
+		Nodename: "n3",
+		Usage:    0.6,
+		Rate:     0.05,
+		Capacity: 8,
+	}
+	arg := []Info{n1, n2, n3}
+	r, err := GlobalPlan(context.TODO(), arg, 3, 100, 0)
+	fmt.Println(r)
+	assert.NoError(t, err)
+	assert.Equal(t, r[arg[0].Nodename], 1)
+}
+
 func TestGlobalPlan2(t *testing.T) {
 	n1 := Info{
 		Nodename: "n1",