New Communisim Strategy (#313)

* new communisim: heap based

* eliminate strategy side effects

* add more tests for communism

Author: jschwinger233

scheduler/complex/potassium_test.go

@@ -3,6 +3,7 @@ package complexscheduler
 import (
 	"errors"
 	"fmt"
+	"sort"
 	"testing"
 
 	"math"
@@ -1011,8 +1012,10 @@ func TestSelectMemoryNodesSequence(t *testing.T) {
 	refreshPod(res, deployMap, mem, 0)
 	res, deployMap, err = SelectMemoryNodes(k, res, nil, cpu, mem, 5, false)
 	assert.NoError(t, err)
-	assert.Equal(t, deployMap[res[0].Name], 3)
-	assert.Equal(t, deployMap[res[1].Name], 2)
+	finalCounts := []int{deployMap[res[0].Name], deployMap[res[1].Name]}
+	sort.Ints(finalCounts)
+	assert.ElementsMatch(t, []int{2, 3}, finalCounts)
+
 }
 
 func TestSelectMemoryNodesGiven(t *testing.T) {
@@ -1024,15 +1027,15 @@ func TestSelectMemoryNodesGiven(t *testing.T) {
 	}
 
 	k, _ := newPotassium()
-	res, deployMap, err := SelectMemoryNodes(k, pod, countMap, 1.0, 512*int64(units.MiB), 2, false)
+	_, deployMap, err := SelectMemoryNodes(k, pod, countMap, 1.0, 512*int64(units.MiB), 2, false)
 	assert.NoError(t, err)
-	for _, node := range res {
-		if node.Name == "n3" {
-			assert.Equal(t, deployMap[node.Name], 2)
-			continue
-		}
-		assert.Equal(t, deployMap[node.Name], 0)
+	finalCounts := []int{}
+	for _, node := range pod {
+		finalCounts = append(finalCounts, countMap[node.Name]+deployMap[node.Name])
 	}
+	sort.Ints(finalCounts)
+	assert.ElementsMatch(t, []int{1, 1, 1, 2}, finalCounts)
+
 }
 
 func TestMaxIdleNode(t *testing.T) {
@@ -1219,10 +1222,11 @@ func TestSelectStorageNodesSequence(t *testing.T) {
 	res, deployMap, err := SelectStorageNodes(k, scheduleInfos, nil, stor, 1, false)
 	assert.NoError(t, err)
 	assert.Equal(t, 2, len(res))
-	assert.Equal(t, 1, deployMap[res[0].Name])
 	assert.Equal(t, 1, res[0].Capacity)
-	assert.Equal(t, 0, deployMap[res[1].Name])
 	assert.Equal(t, 2, res[1].Capacity)
+	counts := []int{deployMap[scheduleInfos[0].Name], deployMap[scheduleInfos[1].Name]}
+	sort.Ints(counts)
+	assert.ElementsMatch(t, []int{0, 1}, counts)
 
 	refreshPod(res, deployMap, mem, stor)
 	countMap := map[string]int{
@@ -1251,10 +1255,14 @@ func TestSelectStorageNodesSequence(t *testing.T) {
 		return
 	}
 	i, j := getLess(res)
-	assert.Equal(t, 0, deployMap[res[i].Name])
-	assert.Equal(t, 2, deployMap[res[j].Name])
-	assert.Equal(t, 1, res[i].Capacity)
-	assert.Equal(t, 0, res[j].Capacity)
+	getFinalCounts := func(deployMap, countMap map[string]int) (counts []int) {
+		for name, d := range deployMap {
+			counts = append(counts, d+countMap[name])
+		}
+		sort.Ints(counts)
+		return
+	}
+	assert.ElementsMatch(t, []int{1, 2}, getFinalCounts(deployMap, countMap))
 
 	refreshPod(res, deployMap, mem, stor)
 	countMap = map[string]int{
@@ -1268,13 +1276,6 @@ func TestSelectStorageNodesSequence(t *testing.T) {
 	assert.Equal(t, 2, len(res))
 	assert.Equal(t, 14, res[0].Capacity)
 	assert.Equal(t, 15, res[1].Capacity)
-
-	res, deployMap, err = SelectStorageNodes(k, res, countMap, int64(units.GiB), 1, false)
-	assert.NoError(t, err)
-	i, j = getLess(res)
-	assert.Equal(t, 2, len(res))
-	assert.Equal(t, 1, deployMap[res[i].Name])
-	assert.Equal(t, 0, res[i].Capacity)
 }
 
 func SelectStorageNodes(k *Potassium, scheduleInfos []resourcetypes.ScheduleInfo, countMap map[string]int, storage int64, need int, each bool) ([]resourcetypes.ScheduleInfo, map[string]int, error) {
@@ -1295,7 +1296,7 @@ func SelectStorageNodes(k *Potassium, scheduleInfos []resourcetypes.ScheduleInfo
 	for i, scheduleInfo := range scheduleInfos {
 		for _, si := range strategyInfos {
 			if si.Nodename == scheduleInfo.Name {
-				scheduleInfos[i].Capacity = si.Capacity
+				scheduleInfos[i].Capacity = si.Capacity - deployMap[si.Nodename]
 			}
 		}
 	}

strategy/average.go

@@ -12,13 +12,15 @@ import (
 // AveragePlan deploy workload each node
 // 容量够的机器每一台部署 N 个
 // need 是每台机器所需总量，limit 是限制节点数
-func AveragePlan(strategyInfos []Info, need, total, limit int, resourceType types.ResourceType) (map[string]int, error) {
+func AveragePlan(infos []Info, need, total, limit int, resourceType types.ResourceType) (map[string]int, error) {
 	log.Debugf("[AveragePlan] need %d limit %d", need, limit)
-	scheduleInfosLength := len(strategyInfos)
+	scheduleInfosLength := len(infos)
 	if scheduleInfosLength < limit {
 		return nil, errors.WithStack(types.NewDetailedErr(types.ErrInsufficientRes,
 			fmt.Sprintf("node len %d < limit, cannot alloc an average node plan", scheduleInfosLength)))
 	}
+	strategyInfos := make([]Info, scheduleInfosLength)
+	copy(strategyInfos, infos)
 	sort.Slice(strategyInfos, func(i, j int) bool { return strategyInfos[i].Capacity < strategyInfos[j].Capacity })
 	p := sort.Search(scheduleInfosLength, func(i int) bool { return strategyInfos[i].Capacity >= need })
 	if p == scheduleInfosLength {

strategy/average_test.go

@@ -1,6 +1,7 @@
 package strategy
 
 import (
+	"sort"
 	"testing"
 
 	"github.com/projecteru2/core/types"
@@ -10,22 +11,15 @@ import (
 func TestAveragePlan(t *testing.T) {
 	// 正常的
 	nodes := deployedNodes()
-	originCap := map[string]int{}
-	for _, v := range nodes {
-		originCap[v.Nodename] = v.Capacity
-	}
 	r, err := AveragePlan(nodes, 1, 0, 0, types.ResourceAll)
 	assert.NoError(t, err)
-	for i := range r {
-		assert.Equal(t, r[i], 1)
-		capacity := 0
-		for _, si := range nodes {
-			if si.Nodename == i {
-				capacity = si.Capacity
-			}
-		}
-		assert.Equal(t, capacity, originCap[i]-1)
+	finalCounts := []int{}
+	for _, node := range nodes {
+		finalCounts = append(finalCounts, node.Count+r[node.Nodename])
 	}
+	sort.Ints(finalCounts)
+	assert.ElementsMatch(t, []int{3, 4, 6, 8}, finalCounts)
+
 	// nodes len < limit
 	nodes = deployedNodes()
 	_, err = AveragePlan(nodes, 100, 0, 5, types.ResourceAll)

strategy/communism.go

@@ -1,57 +1,82 @@
 package strategy
 
 import (
+	"container/heap"
 	"fmt"
-	"sort"
 
 	"github.com/pkg/errors"
-	"github.com/projecteru2/core/log"
 	"github.com/projecteru2/core/types"
 )
 
+type infoHeap struct {
+	infos []Info
+	limit int
+}
+
+func (h infoHeap) Len() int {
+	return len(h.infos)
+}
+
+func (h infoHeap) Less(i, j int) bool {
+	return h.infos[i].Count < h.infos[j].Count || (h.infos[i].Count == h.infos[j].Count && h.infos[i].Capacity > h.infos[j].Capacity)
+}
+
+func (h infoHeap) Swap(i, j int) {
+	h.infos[i], h.infos[j] = h.infos[j], h.infos[i]
+}
+
+func (h *infoHeap) Push(x interface{}) {
+	info := x.(Info)
+	if info.Capacity == 0 || (h.limit > 0 && info.Count >= h.limit) {
+		return
+	}
+	h.infos = append(h.infos, info)
+}
+
+func (h *infoHeap) Pop() interface{} {
+	length := len(h.infos)
+	x := h.infos[length-1]
+	h.infos = h.infos[0 : length-1]
+	return x
+}
+
+func newInfoHeap(infos []Info, limit int) heap.Interface {
+	dup := infoHeap{
+		infos: []Info{},
+		limit: limit,
+	}
+	for _, info := range infos {
+		if info.Capacity == 0 || (limit > 0 && info.Count >= limit) {
+			continue
+		}
+		dup.infos = append(dup.infos, info)
+	}
+	return &dup
+}
+
 // CommunismPlan 吃我一记共产主义大锅饭
 // 部署完 N 个后全局尽可能平均
-func CommunismPlan(arg []Info, need, total, limit int, resourceType types.ResourceType) (map[string]int, error) {
+func CommunismPlan(infos []Info, need, total, limit int, resourceType types.ResourceType) (map[string]int, error) {
 	if total < need {
 		return nil, errors.WithStack(types.NewDetailedErr(types.ErrInsufficientRes,
 			fmt.Sprintf("need: %d, vol: %d", need, total)))
 	}
-	sort.Slice(arg, func(i, j int) bool { return arg[i].Count < arg[j].Count })
-	length := len(arg)
 
-	deployMap := map[string]int{}
-	for i := 0; i < length; i++ {
-		if need <= 0 {
-			break
-		}
-		req := need
-		if i < length-1 {
-			req = (arg[i+1].Count - arg[i].Count) * (i + 1)
-		}
-		if req > need {
-			req = need
+	deploy := map[string]int{}
+	iHeap := newInfoHeap(infos, limit)
+	heap.Init(iHeap)
+	for {
+		if iHeap.Len() == 0 {
+			return nil, errors.WithStack(types.ErrInsufficientRes)
 		}
-		for j := 0; j < i+1; j++ {
-			deploy := req / (i + 1 - j)
-			tail := req % (i + 1 - j)
-			d := deploy
-			if tail > 0 {
-				d++
-			}
-			if d > arg[j].Capacity {
-				d = arg[j].Capacity
-			}
-			if d == 0 {
-				continue
-			}
-			arg[j].Capacity -= d
-			deployMap[arg[j].Nodename] += d
-			need -= d
-			req -= d
+		info := heap.Pop(iHeap).(Info)
+		deploy[info.Nodename]++
+		need--
+		if need == 0 {
+			return deploy, nil
 		}
+		info.Count++
+		info.Capacity--
+		heap.Push(iHeap, info)
 	}
-	// 这里 need 一定会为 0 出来，因为 volTotal 保证了一定大于 need
-	// 这里并不需要再次排序了，理论上的排序是基于 Count 得到的 Deploy 最终方案
-	log.Debugf("[CommunismPlan] strategyInfo: %+v", arg)
-	return deployMap, nil
 }