pcm-coordinator/internal/scheduler/strategy/random.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
}