forked from JointCloud/pcm-coordinator
94 lines
2.2 KiB
Go
94 lines
2.2 KiB
Go
package strategy
|
||
|
||
import (
|
||
"github.com/pkg/errors"
|
||
"math/rand"
|
||
)
|
||
|
||
type RandomStrategy struct {
|
||
clusterIds []string
|
||
replicas int32
|
||
}
|
||
|
||
func NewRandomStrategy(clusterIds []string, replicas int32) *RandomStrategy {
|
||
return &RandomStrategy{clusterIds: clusterIds,
|
||
replicas: replicas,
|
||
}
|
||
}
|
||
|
||
func (s *RandomStrategy) Schedule() ([]*AssignedCluster, error) {
|
||
var results []*AssignedCluster
|
||
if s.replicas < 1 {
|
||
return nil, errors.New("replicas must be greater than 0")
|
||
}
|
||
|
||
if len(s.clusterIds) < 1 {
|
||
return nil, errors.New("cluster must be greater than 0")
|
||
}
|
||
|
||
if len(s.clusterIds) == 0 || s.clusterIds == nil {
|
||
return nil, errors.New("weight must be set")
|
||
}
|
||
|
||
if s.replicas == 1 {
|
||
|
||
// 创建一个切片来保存每个部分的数量
|
||
parts := make([]int32, len(s.clusterIds))
|
||
// 剩余要分配的副本数
|
||
remaining := s.replicas
|
||
|
||
// 随机分配剩余的副本
|
||
for remaining > 0 {
|
||
// 随机选择一个部分（索引从0到numParts-1）
|
||
partIndex := rand.Intn(len(s.clusterIds))
|
||
|
||
// 如果该部分加上一个副本后不会超过总数，则分配一个副本
|
||
//if parts[partIndex]+1 <= s.replicas {
|
||
parts[partIndex]++
|
||
remaining--
|
||
//}
|
||
}
|
||
|
||
if len(s.clusterIds) == len(parts) {
|
||
for i, key := range s.clusterIds {
|
||
cluster := &AssignedCluster{ClusterId: key, Replicas: parts[i]}
|
||
results = append(results, cluster)
|
||
}
|
||
}
|
||
|
||
} else {
|
||
// 创建一个切片来保存每个部分的数量
|
||
parts := make([]int32, len(s.clusterIds))
|
||
|
||
// 首先将每个部分都分配至少一个副本
|
||
for i := range parts {
|
||
parts[i] = 1
|
||
s.replicas--
|
||
}
|
||
|
||
// 剩余要分配的副本数
|
||
remaining := s.replicas
|
||
|
||
// 随机分配剩余的副本
|
||
for remaining > 0 {
|
||
// 随机选择一个部分（索引从0到numParts-1）
|
||
partIndex := rand.Intn(len(s.clusterIds))
|
||
|
||
// 如果该部分加上一个副本后不会超过总数，则分配一个副本
|
||
//if parts[partIndex]+1 <= s.replicas {
|
||
parts[partIndex]++
|
||
remaining--
|
||
//}
|
||
}
|
||
|
||
if len(s.clusterIds) == len(parts) {
|
||
for i, key := range s.clusterIds {
|
||
cluster := &AssignedCluster{ClusterId: key, Replicas: parts[i]}
|
||
results = append(results, cluster)
|
||
}
|
||
}
|
||
|
||
}
|
||
return results, nil
|
||
}
|