github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/strategy/communism.go (about)

     1  package strategy
     2  
     3  import (
     4  	"container/heap"
     5  	"context"
     6  
     7  	"github.com/projecteru2/core/types"
     8  
     9  	"github.com/cockroachdb/errors"
    10  )
    11  
    12  type infoHeap struct {
    13  	infos []Info
    14  	limit int
    15  }
    16  
    17  func (h infoHeap) Len() int {
    18  	return len(h.infos)
    19  }
    20  
    21  func (h infoHeap) Less(i, j int) bool {
    22  	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)
    23  }
    24  
    25  func (h infoHeap) Swap(i, j int) {
    26  	h.infos[i], h.infos[j] = h.infos[j], h.infos[i]
    27  }
    28  
    29  func (h *infoHeap) Push(x any) {
    30  	info := x.(Info)
    31  	if info.Capacity == 0 || (h.limit > 0 && info.Count >= h.limit) {
    32  		return
    33  	}
    34  	h.infos = append(h.infos, info)
    35  }
    36  
    37  func (h *infoHeap) Pop() any {
    38  	length := len(h.infos)
    39  	x := h.infos[length-1]
    40  	h.infos = h.infos[0 : length-1]
    41  	return x
    42  }
    43  
    44  func newInfoHeap(infos []Info, limit int) heap.Interface {
    45  	dup := infoHeap{
    46  		infos: []Info{},
    47  		limit: limit,
    48  	}
    49  	for _, info := range infos {
    50  		if info.Capacity == 0 || (limit > 0 && info.Count >= limit) {
    51  			continue
    52  		}
    53  		dup.infos = append(dup.infos, info)
    54  	}
    55  	return &dup
    56  }
    57  
    58  // CommunismPlan 吃我一记共产主义大锅饭
    59  // 部署完 N 个后全局尽可能平均
    60  func CommunismPlan(_ context.Context, infos []Info, need, total, limit int) (map[string]int, error) {
    61  	if total < need {
    62  		return nil, errors.Wrapf(types.ErrInsufficientResource, "need: %d, available: %d", need, total)
    63  	}
    64  
    65  	deploy := map[string]int{}
    66  	iHeap := newInfoHeap(infos, limit)
    67  	heap.Init(iHeap)
    68  	for {
    69  		if iHeap.Len() == 0 {
    70  			return nil, errors.Wrapf(types.ErrInsufficientResource, "reached nodelimit, a node can host at most %d instances", limit)
    71  		}
    72  		info := heap.Pop(iHeap).(Info)
    73  		deploy[info.Nodename]++
    74  		need--
    75  		if need == 0 {
    76  			return deploy, nil
    77  		}
    78  		info.Count++
    79  		info.Capacity--
    80  		heap.Push(iHeap, info)
    81  	}
    82  }