mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2024-12-14 10:39:24 +08:00
144 lines
4.2 KiB
Go
144 lines
4.2 KiB
Go
package pub_balancer
|
|
|
|
import (
|
|
cmap "github.com/orcaman/concurrent-map/v2"
|
|
"github.com/seaweedfs/seaweedfs/weed/glog"
|
|
"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
|
|
"math/rand"
|
|
"time"
|
|
)
|
|
|
|
func AllocateTopicPartitions(brokers cmap.ConcurrentMap[string, *BrokerStats], partitionCount int32) (assignments []*mq_pb.BrokerPartitionAssignment) {
|
|
// divide the ring into partitions
|
|
now := time.Now().UnixNano()
|
|
rangeSize := MaxPartitionCount / partitionCount
|
|
for i := int32(0); i < partitionCount; i++ {
|
|
assignment := &mq_pb.BrokerPartitionAssignment{
|
|
Partition: &mq_pb.Partition{
|
|
RingSize: MaxPartitionCount,
|
|
RangeStart: int32(i * rangeSize),
|
|
RangeStop: int32((i + 1) * rangeSize),
|
|
UnixTimeNs: now,
|
|
},
|
|
}
|
|
if i == partitionCount-1 {
|
|
assignment.Partition.RangeStop = MaxPartitionCount
|
|
}
|
|
assignments = append(assignments, assignment)
|
|
}
|
|
|
|
// pick the brokers
|
|
pickedBrokers := pickBrokers(brokers, partitionCount)
|
|
|
|
// assign the partitions to brokers
|
|
for i, assignment := range assignments {
|
|
assignment.LeaderBroker = pickedBrokers[i]
|
|
}
|
|
glog.V(0).Infof("allocate topic partitions %d: %v", len(assignments), assignments)
|
|
return
|
|
}
|
|
|
|
// randomly pick n brokers, which may contain duplicates
|
|
// TODO pick brokers based on the broker stats
|
|
func pickBrokers(brokers cmap.ConcurrentMap[string, *BrokerStats], count int32) []string {
|
|
candidates := make([]string, 0, brokers.Count())
|
|
for brokerStatsItem := range brokers.IterBuffered() {
|
|
candidates = append(candidates, brokerStatsItem.Key)
|
|
}
|
|
pickedBrokers := make([]string, 0, count)
|
|
for i := int32(0); i < count; i++ {
|
|
p := rand.Intn(len(candidates))
|
|
pickedBrokers = append(pickedBrokers, candidates[p])
|
|
}
|
|
return pickedBrokers
|
|
}
|
|
|
|
// reservoir sampling select N brokers from the active brokers, with exclusion of the excluded brokers
|
|
func pickBrokersExcluded(brokers []string, count int, excludedLeadBroker string, excludedBrokers []string) []string {
|
|
// convert the excluded brokers to a map
|
|
excludedBrokerMap := make(map[string]bool)
|
|
for _, broker := range excludedBrokers {
|
|
excludedBrokerMap[broker] = true
|
|
}
|
|
if excludedLeadBroker != "" {
|
|
excludedBrokerMap[excludedLeadBroker] = true
|
|
}
|
|
|
|
pickedBrokers := make([]string, 0, count)
|
|
for i, broker := range brokers {
|
|
if _, found := excludedBrokerMap[broker]; found {
|
|
continue
|
|
}
|
|
if len(pickedBrokers) < count {
|
|
pickedBrokers = append(pickedBrokers, broker)
|
|
} else {
|
|
j := rand.Intn(i + 1)
|
|
if j < count {
|
|
pickedBrokers[j] = broker
|
|
}
|
|
}
|
|
}
|
|
|
|
// shuffle the picked brokers
|
|
count = len(pickedBrokers)
|
|
for i := 0; i < count; i++ {
|
|
j := rand.Intn(count)
|
|
pickedBrokers[i], pickedBrokers[j] = pickedBrokers[j], pickedBrokers[i]
|
|
}
|
|
|
|
return pickedBrokers
|
|
}
|
|
|
|
func EnsureAssignmentsToActiveBrokers(activeBrokers cmap.ConcurrentMap[string, *BrokerStats], followerCount int, assignments []*mq_pb.BrokerPartitionAssignment) (hasChanges bool) {
|
|
candidates := make([]string, 0, activeBrokers.Count())
|
|
for brokerStatsItem := range activeBrokers.IterBuffered() {
|
|
candidates = append(candidates, brokerStatsItem.Key)
|
|
}
|
|
|
|
for _, assignment := range assignments {
|
|
// count how many brokers are needed
|
|
count := 0
|
|
if assignment.LeaderBroker == "" {
|
|
count++
|
|
} else if _, found := activeBrokers.Get(assignment.LeaderBroker); !found {
|
|
assignment.LeaderBroker = ""
|
|
count++
|
|
}
|
|
for i:=0; i<followerCount; i++ {
|
|
if i >= len(assignment.FollowerBrokers) {
|
|
count++
|
|
continue
|
|
}
|
|
if assignment.FollowerBrokers[i] == "" {
|
|
count++
|
|
} else if _, found := activeBrokers.Get(assignment.FollowerBrokers[i]); !found {
|
|
assignment.FollowerBrokers[i] = ""
|
|
count++
|
|
}
|
|
}
|
|
|
|
if count > 0 {
|
|
pickedBrokers := pickBrokersExcluded(candidates, count, assignment.LeaderBroker, assignment.FollowerBrokers)
|
|
i := 0
|
|
if assignment.LeaderBroker == "" {
|
|
assignment.LeaderBroker = pickedBrokers[i]
|
|
i++
|
|
}
|
|
j := 0
|
|
for ; j<len(assignment.FollowerBrokers); j++ {
|
|
if assignment.FollowerBrokers[j] == "" {
|
|
assignment.FollowerBrokers[j] = pickedBrokers[i]
|
|
i++
|
|
}
|
|
}
|
|
if i < len(pickedBrokers) {
|
|
assignment.FollowerBrokers = append(assignment.FollowerBrokers, pickedBrokers[i:]...)
|
|
}
|
|
}
|
|
|
|
hasChanges = hasChanges || count > 0
|
|
}
|
|
|
|
return
|
|
}
|