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

     1  package strategy
     2  
     3  import (
     4  	"container/heap"
     5  	"context"
     6  
     7  	"github.com/cockroachdb/errors"
     8  	"github.com/projecteru2/core/log"
     9  	"github.com/projecteru2/core/types"
    10  )
    11  
    12  type infoHeapForGlobalStrategy []Info
    13  
    14  // Len .
    15  func (r infoHeapForGlobalStrategy) Len() int {
    16  	return len(r)
    17  }
    18  
    19  // Less .
    20  func (r infoHeapForGlobalStrategy) Less(i, j int) bool {
    21  	return (r[i].Usage + r[i].Rate) < (r[j].Usage + r[j].Rate)
    22  }
    23  
    24  // Swap .
    25  func (r infoHeapForGlobalStrategy) Swap(i, j int) {
    26  	r[i], r[j] = r[j], r[i]
    27  }
    28  
    29  // Push .
    30  func (r *infoHeapForGlobalStrategy) Push(x any) {
    31  	*r = append(*r, x.(Info))
    32  }
    33  
    34  // Pop .
    35  func (r *infoHeapForGlobalStrategy) Pop() any {
    36  	old := *r
    37  	n := len(old)
    38  	x := old[n-1]
    39  	*r = old[:n-1]
    40  	return x
    41  }
    42  
    43  // GlobalPlan 基于全局资源配额
    44  // 尽量使得资源消耗平均
    45  func GlobalPlan(ctx context.Context, infos []Info, need, total, _ int) (map[string]int, error) {
    46  	if total < need {
    47  		return nil, errors.Wrapf(types.ErrInsufficientResource, "need: %d, available: %d", need, total)
    48  	}
    49  	strategyInfos := make([]Info, len(infos))
    50  	copy(strategyInfos, infos)
    51  	deployMap := map[string]int{}
    52  
    53  	infoHeap := &infoHeapForGlobalStrategy{}
    54  	for _, info := range strategyInfos {
    55  		if info.Capacity > 0 {
    56  			infoHeap.Push(info)
    57  		}
    58  	}
    59  	heap.Init(infoHeap)
    60  
    61  	for i := 0; i < need; i++ {
    62  		if infoHeap.Len() == 0 {
    63  			return nil, errors.Wrapf(types.ErrInsufficientResource, "need: %d, available: %d", need, i)
    64  		}
    65  		infoWithMinUsage := heap.Pop(infoHeap).(Info)
    66  		deployMap[infoWithMinUsage.Nodename]++
    67  		infoWithMinUsage.Usage += infoWithMinUsage.Rate
    68  		infoWithMinUsage.Capacity--
    69  
    70  		if infoWithMinUsage.Capacity > 0 {
    71  			heap.Push(infoHeap, infoWithMinUsage)
    72  		}
    73  	}
    74  
    75  	// 这里 need 一定会为 0 出来，因为 volTotal 保证了一定大于 need
    76  	// 这里并不需要再次排序了，理论上的排序是基于资源使用率得到的 Deploy 最终方案
    77  	log.WithFunc("strategy.GlobalPlan").Debugf(ctx, "strategyInfos: %+v", strategyInfos)
    78  	return deployMap, nil
    79  }