snapshot of https://github.com/seaweedfs/seaweedfs.git (master, 2026-09-16; blobs > 131072B omitted)

weed/mq/pub_balancer/allocate.go

@@ -0,0 +1,154 @@
+package pub_balancer
+
+import (
+	"math/rand/v2"
+	"time"
+
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+// findActiveBroker looks up a broker in activeBrokers, tolerating address
+// encoding differences (plain host:port vs host:port.grpcPort) by falling
+// back to an Equals-based scan when the direct key lookup misses.
+func findActiveBroker(activeBrokers cmap.ConcurrentMap[string, *BrokerStats], addr string) (string, bool) {
+	if addr == "" {
+		return "", false
+	}
+	if _, found := activeBrokers.Get(addr); found {
+		return addr, true
+	}
+	target := pb.ServerAddress(addr)
+	for item := range activeBrokers.IterBuffered() {
+		if pb.ServerAddress(item.Key).Equals(target) {
+			return item.Key, true
+		}
+	}
+	return "", false
+}
+
+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: &schema_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)
+	}
+
+	EnsureAssignmentsToActiveBrokers(brokers, 1, assignments)
+
+	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 broker
+func pickBrokersExcluded(brokers []string, count int, excludedLeadBroker string, excludedBroker string) []string {
+	pickedBrokers := make([]string, 0, count)
+	for i, broker := range brokers {
+		if broker == excludedBroker {
+			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
+}
+
+// EnsureAssignmentsToActiveBrokers ensures the assignments are assigned to active brokers
+func EnsureAssignmentsToActiveBrokers(activeBrokers cmap.ConcurrentMap[string, *BrokerStats], followerCount int, assignments []*mq_pb.BrokerPartitionAssignment) (hasChanges bool) {
+	glog.V(4).Infof("EnsureAssignmentsToActiveBrokers: activeBrokers: %v, followerCount: %d, assignments: %v", activeBrokers.Count(), followerCount, assignments)
+
+	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 canonical, found := findActiveBroker(activeBrokers, assignment.LeaderBroker); !found {
+			assignment.LeaderBroker = ""
+			count++
+		} else if canonical != assignment.LeaderBroker {
+			assignment.LeaderBroker = canonical
+			hasChanges = true
+		}
+		if assignment.FollowerBroker == "" {
+			count++
+		} else if canonical, found := findActiveBroker(activeBrokers, assignment.FollowerBroker); !found {
+			assignment.FollowerBroker = ""
+			count++
+		} else if canonical != assignment.FollowerBroker {
+			assignment.FollowerBroker = canonical
+			hasChanges = true
+		}
+
+		if count > 0 {
+			pickedBrokers := pickBrokersExcluded(candidates, count, assignment.LeaderBroker, assignment.FollowerBroker)
+			i := 0
+			if assignment.LeaderBroker == "" {
+				if i < len(pickedBrokers) {
+					assignment.LeaderBroker = pickedBrokers[i]
+					i++
+					hasChanges = true
+				}
+			}
+			if assignment.FollowerBroker == "" {
+				if i < len(pickedBrokers) {
+					assignment.FollowerBroker = pickedBrokers[i]
+					i++
+					hasChanges = true
+				}
+			}
+		}
+
+	}
+
+	glog.V(4).Infof("EnsureAssignmentsToActiveBrokers: activeBrokers: %v, followerCount: %d, assignments: %v hasChanges: %v", activeBrokers.Count(), followerCount, assignments, hasChanges)
+	return
+}

weed/mq/pub_balancer/allocate_test.go

@@ -0,0 +1,206 @@
+package pub_balancer
+
+import (
+	"fmt"
+	"testing"
+
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_allocateOneBroker(t *testing.T) {
+	brokers := cmap.New[*BrokerStats]()
+	brokers.SetIfAbsent("localhost:17777", &BrokerStats{
+		TopicPartitionCount: 0,
+		CpuUsagePercent:     0,
+	})
+
+	tests := []struct {
+		name            string
+		args            args
+		wantAssignments []*mq_pb.BrokerPartitionAssignment
+	}{
+		{
+			name: "test only one broker",
+			args: args{
+				brokers:        brokers,
+				partitionCount: 1,
+			},
+			wantAssignments: []*mq_pb.BrokerPartitionAssignment{
+				{
+					LeaderBroker: "localhost:17777",
+					Partition: &schema_pb.Partition{
+						RingSize:   MaxPartitionCount,
+						RangeStart: 0,
+						RangeStop:  MaxPartitionCount,
+					},
+				},
+			},
+		},
+	}
+	testThem(t, tests)
+}
+
+type args struct {
+	brokers        cmap.ConcurrentMap[string, *BrokerStats]
+	partitionCount int32
+}
+
+func testThem(t *testing.T, tests []struct {
+	name            string
+	args            args
+	wantAssignments []*mq_pb.BrokerPartitionAssignment
+}) {
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gotAssignments := AllocateTopicPartitions(tt.args.brokers, tt.args.partitionCount)
+			assert.Equal(t, len(tt.wantAssignments), len(gotAssignments))
+			for i, gotAssignment := range gotAssignments {
+				assert.Equal(t, tt.wantAssignments[i].LeaderBroker, gotAssignment.LeaderBroker)
+				assert.Equal(t, tt.wantAssignments[i].Partition.RangeStart, gotAssignment.Partition.RangeStart)
+				assert.Equal(t, tt.wantAssignments[i].Partition.RangeStop, gotAssignment.Partition.RangeStop)
+				assert.Equal(t, tt.wantAssignments[i].Partition.RingSize, gotAssignment.Partition.RingSize)
+			}
+		})
+	}
+}
+
+func TestEnsureAssignmentsToActiveBrokersX(t *testing.T) {
+	type args struct {
+		activeBrokers cmap.ConcurrentMap[string, *BrokerStats]
+		followerCount int
+		assignments   []*mq_pb.BrokerPartitionAssignment
+	}
+	activeBrokers := cmap.New[*BrokerStats]()
+	activeBrokers.SetIfAbsent("localhost:1", &BrokerStats{})
+	activeBrokers.SetIfAbsent("localhost:2", &BrokerStats{})
+	activeBrokers.SetIfAbsent("localhost:3", &BrokerStats{})
+	activeBrokers.SetIfAbsent("localhost:4", &BrokerStats{})
+	activeBrokers.SetIfAbsent("localhost:5", &BrokerStats{})
+	activeBrokers.SetIfAbsent("localhost:6", &BrokerStats{})
+	lowActiveBrokers := cmap.New[*BrokerStats]()
+	lowActiveBrokers.SetIfAbsent("localhost:1", &BrokerStats{})
+	lowActiveBrokers.SetIfAbsent("localhost:2", &BrokerStats{})
+	singleActiveBroker := cmap.New[*BrokerStats]()
+	singleActiveBroker.SetIfAbsent("localhost:1", &BrokerStats{})
+	tests := []struct {
+		name       string
+		args       args
+		hasChanges bool
+	}{
+		{
+			name: "test empty leader",
+			args: args{
+				activeBrokers: activeBrokers,
+				followerCount: 1,
+				assignments: []*mq_pb.BrokerPartitionAssignment{
+					{
+						LeaderBroker:   "",
+						Partition:      &schema_pb.Partition{},
+						FollowerBroker: "localhost:2",
+					},
+				},
+			},
+			hasChanges: true,
+		},
+		{
+			name: "test empty follower",
+			args: args{
+				activeBrokers: activeBrokers,
+				followerCount: 1,
+				assignments: []*mq_pb.BrokerPartitionAssignment{
+					{
+						LeaderBroker:   "localhost:1",
+						Partition:      &schema_pb.Partition{},
+						FollowerBroker: "",
+					},
+				},
+			},
+			hasChanges: true,
+		},
+		{
+			name: "test dead follower",
+			args: args{
+				activeBrokers: activeBrokers,
+				followerCount: 1,
+				assignments: []*mq_pb.BrokerPartitionAssignment{
+					{
+						LeaderBroker:   "localhost:1",
+						Partition:      &schema_pb.Partition{},
+						FollowerBroker: "localhost:200",
+					},
+				},
+			},
+			hasChanges: true,
+		},
+		{
+			name: "test dead leader and follower",
+			args: args{
+				activeBrokers: activeBrokers,
+				followerCount: 1,
+				assignments: []*mq_pb.BrokerPartitionAssignment{
+					{
+						LeaderBroker:   "localhost:100",
+						Partition:      &schema_pb.Partition{},
+						FollowerBroker: "localhost:200",
+					},
+				},
+			},
+			hasChanges: true,
+		},
+		{
+			name: "test low active brokers",
+			args: args{
+				activeBrokers: lowActiveBrokers,
+				followerCount: 3,
+				assignments: []*mq_pb.BrokerPartitionAssignment{
+					{
+						LeaderBroker:   "localhost:1",
+						Partition:      &schema_pb.Partition{},
+						FollowerBroker: "localhost:2",
+					},
+				},
+			},
+			hasChanges: false,
+		},
+		{
+			name: "test low active brokers with one follower",
+			args: args{
+				activeBrokers: lowActiveBrokers,
+				followerCount: 1,
+				assignments: []*mq_pb.BrokerPartitionAssignment{
+					{
+						LeaderBroker: "localhost:1",
+						Partition:    &schema_pb.Partition{},
+					},
+				},
+			},
+			hasChanges: true,
+		},
+		{
+			name: "test single active broker",
+			args: args{
+				activeBrokers: singleActiveBroker,
+				followerCount: 3,
+				assignments: []*mq_pb.BrokerPartitionAssignment{
+					{
+						LeaderBroker:   "localhost:1",
+						Partition:      &schema_pb.Partition{},
+						FollowerBroker: "localhost:2",
+					},
+				},
+			},
+			hasChanges: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			fmt.Printf("%v before %v\n", tt.name, tt.args.assignments)
+			hasChanges := EnsureAssignmentsToActiveBrokers(tt.args.activeBrokers, tt.args.followerCount, tt.args.assignments)
+			assert.Equalf(t, tt.hasChanges, hasChanges, "EnsureAssignmentsToActiveBrokers(%v, %v, %v)", tt.args.activeBrokers, tt.args.followerCount, tt.args.assignments)
+			fmt.Printf("%v after %v\n", tt.name, tt.args.assignments)
+		})
+	}
+}

weed/mq/pub_balancer/balance.go

@@ -0,0 +1,73 @@
+package pub_balancer
+
+import (
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"google.golang.org/grpc"
+)
+
+/*
+* Assuming a topic has [x,y] number of partitions when publishing, and there are b number of brokers.
+* and p is the number of partitions per topic.
+* if the broker number b <= x, then p = x.
+* if the broker number x < b < y, then x <= p <= b.
+* if the broker number b >= y, x <= p <= y
+
+Balance topic partitions to brokers
+===================================
+
+When the goal is to make sure that low traffic partitions can be merged, (and p >= x, and after last rebalance interval):
+1. Calculate the average load(throughput) of partitions per topic.
+2. If any two neighboring partitions have a load that is less than the average load, merge them.
+3. If min(b, y) < p, then merge two neighboring partitions that have the least combined load.
+
+When the goal is to make sure that high traffic partitions can be split, (and p < y and p < b, and after last rebalance interval):
+1. Calculate the average number of partitions per broker.
+2. If any partition has a load that is more than the average load, split it into two partitions.
+
+When the goal is to make sure that each broker has the same number of partitions:
+1. Calculate the average number of partitions per broker.
+2. For the brokers that have more than the average number of partitions, move the partitions to the brokers that have less than the average number of partitions.
+
+*/
+
+type BalanceAction interface {
+}
+type BalanceActionMerge struct {
+	Before []topic.TopicPartition
+	After  topic.TopicPartition
+}
+type BalanceActionSplit struct {
+	Before topic.TopicPartition
+	After  []topic.TopicPartition
+}
+
+type BalanceActionMove struct {
+	TopicPartition topic.TopicPartition
+	SourceBroker   string
+	TargetBroker   string
+}
+
+type BalanceActionCreate struct {
+	TopicPartition topic.TopicPartition
+	TargetBroker   string
+}
+
+// BalancePublishers check the stats of all brokers,
+// and balance the publishers to the brokers.
+func (balancer *PubBalancer) BalancePublishers() []BalanceAction {
+	action := BalanceTopicPartitionOnBrokers(balancer.Brokers)
+	return []BalanceAction{action}
+}
+
+func (balancer *PubBalancer) ExecuteBalanceAction(actions []BalanceAction, grpcDialOption grpc.DialOption) (err error) {
+	for _, action := range actions {
+		switch action.(type) {
+		case *BalanceActionMove:
+			err = balancer.ExecuteBalanceActionMove(action.(*BalanceActionMove), grpcDialOption)
+		}
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

weed/mq/pub_balancer/balance_action.go

@@ -0,0 +1,59 @@
+package pub_balancer
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/seaweedfs/seaweedfs/weed/pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"google.golang.org/grpc"
+)
+
+// PubBalancer <= PublisherToPubBalancer() <= Broker <=> Publish()
+// ExecuteBalanceActionMove from PubBalancer => AssignTopicPartitions() => Broker => Publish()
+
+func (balancer *PubBalancer) ExecuteBalanceActionMove(move *BalanceActionMove, grpcDialOption grpc.DialOption) error {
+	if _, found := balancer.Brokers.Get(move.SourceBroker); !found {
+		return fmt.Errorf("source broker %s not found", move.SourceBroker)
+	}
+	if _, found := balancer.Brokers.Get(move.TargetBroker); !found {
+		return fmt.Errorf("target broker %s not found", move.TargetBroker)
+	}
+
+	err := pb.WithBrokerGrpcClient(false, move.TargetBroker, grpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+		_, err := client.AssignTopicPartitions(context.Background(), &mq_pb.AssignTopicPartitionsRequest{
+			Topic: move.TopicPartition.Topic.ToPbTopic(),
+			BrokerPartitionAssignments: []*mq_pb.BrokerPartitionAssignment{
+				{
+					Partition: move.TopicPartition.ToPbPartition(),
+				},
+			},
+			IsLeader:   true,
+			IsDraining: false,
+		})
+		return err
+	})
+	if err != nil {
+		return fmt.Errorf("assign topic partition %v to %s: %v", move.TopicPartition, move.TargetBroker, err)
+	}
+
+	err = pb.WithBrokerGrpcClient(false, move.SourceBroker, grpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {
+		_, err := client.AssignTopicPartitions(context.Background(), &mq_pb.AssignTopicPartitionsRequest{
+			Topic: move.TopicPartition.Topic.ToPbTopic(),
+			BrokerPartitionAssignments: []*mq_pb.BrokerPartitionAssignment{
+				{
+					Partition: move.TopicPartition.ToPbPartition(),
+				},
+			},
+			IsLeader:   true,
+			IsDraining: true,
+		})
+		return err
+	})
+	if err != nil {
+		return fmt.Errorf("assign topic partition %v to %s: %v", move.TopicPartition, move.SourceBroker, err)
+	}
+
+	return nil
+
+}

weed/mq/pub_balancer/balance_action_split.go

@@ -0,0 +1,43 @@
+package pub_balancer
+
+/*
+Sequence of operations to ensure ordering
+
+Assuming Publisher P10 is publishing to Topic Partition TP10, and Subscriber S10 is subscribing to Topic TP10.
+After splitting Topic TP10 into Topic Partition TP11 and Topic Partition TP21,
+Publisher P11 is publishing to Topic Partition TP11, and Publisher P21 is publishing to Topic Partition TP21.
+Subscriber S12 is subscribing to Topic Partition TP11, and Subscriber S21 is subscribing to Topic Partition TP21.
+
+(The last digit is ephoch generation number, which is increasing when the topic partitioning is changed.)
+
+The diagram is as follows:
+P10 -> TP10 -> S10
+   ||
+   \/
+P11 -> TP11 -> S11
+P21 -> TP21 -> S21
+
+The following is the sequence of events:
+1. Create Topic Partition TP11 and TP21
+2. Close Publisher(s) P10
+3. Close Subscriber(s) S10
+4. Close Topic Partition TP10
+5. Start Publisher P11, P21
+6. Start Subscriber S11, S21
+
+The dependency is as follows:
+          2  => 3  => 4
+          |     |
+          v     v
+    1 => (5  |  6)
+
+And also:
+2 => 5
+3 => 6
+
+For brokers:
+1. Close all publishers for a topic partition
+2. Close all subscribers for a topic partition
+3. Close the topic partition
+
+*/

weed/mq/pub_balancer/balance_brokers.go

@@ -0,0 +1,53 @@
+package pub_balancer
+
+import (
+	"math/rand/v2"
+
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+)
+
+func BalanceTopicPartitionOnBrokers(brokers cmap.ConcurrentMap[string, *BrokerStats]) BalanceAction {
+	// 1. calculate the average number of partitions per broker
+	var totalPartitionCount int32
+	var totalBrokerCount int32
+	for brokerStats := range brokers.IterBuffered() {
+		totalBrokerCount++
+		totalPartitionCount += brokerStats.Val.TopicPartitionCount
+	}
+	averagePartitionCountPerBroker := totalPartitionCount / totalBrokerCount
+	minPartitionCountPerBroker := averagePartitionCountPerBroker
+	maxPartitionCountPerBroker := averagePartitionCountPerBroker
+	var sourceBroker, targetBroker string
+	var candidatePartition *topic.TopicPartition
+	for brokerStats := range brokers.IterBuffered() {
+		if minPartitionCountPerBroker > brokerStats.Val.TopicPartitionCount {
+			minPartitionCountPerBroker = brokerStats.Val.TopicPartitionCount
+			targetBroker = brokerStats.Key
+		}
+		if maxPartitionCountPerBroker < brokerStats.Val.TopicPartitionCount {
+			maxPartitionCountPerBroker = brokerStats.Val.TopicPartitionCount
+			sourceBroker = brokerStats.Key
+			// select a random partition from the source broker
+			randomPartitionIndex := rand.IntN(int(brokerStats.Val.TopicPartitionCount))
+			index := 0
+			for topicPartitionStats := range brokerStats.Val.TopicPartitionStats.IterBuffered() {
+				if index == randomPartitionIndex {
+					candidatePartition = &topicPartitionStats.Val.TopicPartition
+					break
+				} else {
+					index++
+				}
+			}
+		}
+	}
+	if minPartitionCountPerBroker >= maxPartitionCountPerBroker-1 {
+		return nil
+	}
+	// 2. move the partitions from the source broker to the target broker
+	return &BalanceActionMove{
+		TopicPartition: *candidatePartition,
+		SourceBroker:   sourceBroker,
+		TargetBroker:   targetBroker,
+	}
+}

weed/mq/pub_balancer/balance_brokers_test.go

@@ -0,0 +1,68 @@
+package pub_balancer
+
+import (
+	"reflect"
+	"testing"
+
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+)
+
+func TestBalanceTopicPartitionOnBrokers(t *testing.T) {
+
+	brokers := cmap.New[*BrokerStats]()
+	broker1Stats := &BrokerStats{
+		TopicPartitionCount: 1,
+		CpuUsagePercent:     1,
+		TopicPartitionStats: cmap.New[*TopicPartitionStats](),
+	}
+	broker1Stats.TopicPartitionStats.Set("topic1:0", &TopicPartitionStats{
+		TopicPartition: topic.TopicPartition{
+			Topic:     topic.Topic{Namespace: "topic1", Name: "topic1"},
+			Partition: topic.Partition{RangeStart: 0, RangeStop: 512, RingSize: 1024},
+		},
+	})
+	broker2Stats := &BrokerStats{
+		TopicPartitionCount: 2,
+		CpuUsagePercent:     1,
+		TopicPartitionStats: cmap.New[*TopicPartitionStats](),
+	}
+	broker2Stats.TopicPartitionStats.Set("topic1:1", &TopicPartitionStats{
+		TopicPartition: topic.TopicPartition{
+			Topic:     topic.Topic{Namespace: "topic1", Name: "topic1"},
+			Partition: topic.Partition{RangeStart: 512, RangeStop: 1024, RingSize: 1024},
+		},
+	})
+	broker2Stats.TopicPartitionStats.Set("topic2:0", &TopicPartitionStats{
+		TopicPartition: topic.TopicPartition{
+			Topic:     topic.Topic{Namespace: "topic2", Name: "topic2"},
+			Partition: topic.Partition{RangeStart: 0, RangeStop: 1024, RingSize: 1024},
+		},
+	})
+	brokers.Set("broker1", broker1Stats)
+	brokers.Set("broker2", broker2Stats)
+
+	type args struct {
+		brokers cmap.ConcurrentMap[string, *BrokerStats]
+	}
+	tests := []struct {
+		name string
+		args args
+		want BalanceAction
+	}{
+		{
+			name: "test",
+			args: args{
+				brokers: brokers,
+			},
+			want: nil,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := BalanceTopicPartitionOnBrokers(tt.args.brokers); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("BalanceTopicPartitionOnBrokers() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

weed/mq/pub_balancer/broker_stats.go

@@ -0,0 +1,90 @@
+package pub_balancer
+
+import (
+	"fmt"
+
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+type BrokerStats struct {
+	TopicPartitionCount int32
+	PublisherCount      int32
+	SubscriberCount     int32
+	CpuUsagePercent     int32
+	TopicPartitionStats cmap.ConcurrentMap[string, *TopicPartitionStats] // key: topic_partition
+	Topics              []topic.Topic
+}
+type TopicPartitionStats struct {
+	topic.TopicPartition
+	PublisherCount  int32
+	SubscriberCount int32
+}
+
+func NewBrokerStats() *BrokerStats {
+	return &BrokerStats{
+		TopicPartitionStats: cmap.New[*TopicPartitionStats](),
+	}
+}
+func (bs *BrokerStats) String() string {
+	return fmt.Sprintf("BrokerStats{TopicPartitionCount:%d, Publishers:%d, Subscribers:%d CpuUsagePercent:%d, Stats:%+v}",
+		bs.TopicPartitionCount, bs.PublisherCount, bs.SubscriberCount, bs.CpuUsagePercent, bs.TopicPartitionStats.Items())
+}
+
+func (bs *BrokerStats) UpdateStats(stats *mq_pb.BrokerStats) {
+	bs.TopicPartitionCount = int32(len(stats.Stats))
+	bs.CpuUsagePercent = stats.CpuUsagePercent
+
+	var publisherCount, subscriberCount int32
+	currentTopicPartitions := bs.TopicPartitionStats.Items()
+	for _, topicPartitionStats := range stats.Stats {
+		tps := &TopicPartitionStats{
+			TopicPartition: topic.TopicPartition{
+				Topic: topic.Topic{Namespace: topicPartitionStats.Topic.Namespace, Name: topicPartitionStats.Topic.Name},
+				Partition: topic.Partition{
+					RangeStart: topicPartitionStats.Partition.RangeStart,
+					RangeStop:  topicPartitionStats.Partition.RangeStop,
+					RingSize:   topicPartitionStats.Partition.RingSize,
+					UnixTimeNs: topicPartitionStats.Partition.UnixTimeNs,
+				},
+			},
+			PublisherCount:  topicPartitionStats.PublisherCount,
+			SubscriberCount: topicPartitionStats.SubscriberCount,
+		}
+		publisherCount += topicPartitionStats.PublisherCount
+		subscriberCount += topicPartitionStats.SubscriberCount
+		key := tps.TopicPartition.TopicPartitionId()
+		bs.TopicPartitionStats.Set(key, tps)
+		delete(currentTopicPartitions, key)
+	}
+	// remove the topic partitions that are not in the stats
+	for key := range currentTopicPartitions {
+		bs.TopicPartitionStats.Remove(key)
+	}
+	bs.PublisherCount = publisherCount
+	bs.SubscriberCount = subscriberCount
+}
+
+func (bs *BrokerStats) RegisterAssignment(t *schema_pb.Topic, partition *schema_pb.Partition, isAdd bool) {
+	tps := &TopicPartitionStats{
+		TopicPartition: topic.TopicPartition{
+			Topic: topic.Topic{Namespace: t.Namespace, Name: t.Name},
+			Partition: topic.Partition{
+				RangeStart: partition.RangeStart,
+				RangeStop:  partition.RangeStop,
+				RingSize:   partition.RingSize,
+				UnixTimeNs: partition.UnixTimeNs,
+			},
+		},
+		PublisherCount:  0,
+		SubscriberCount: 0,
+	}
+	key := tps.TopicPartition.TopicPartitionId()
+	if isAdd {
+		bs.TopicPartitionStats.SetIfAbsent(key, tps)
+	} else {
+		bs.TopicPartitionStats.Remove(key)
+	}
+}

weed/mq/pub_balancer/lookup.go

@@ -0,0 +1,37 @@
+package pub_balancer
+
+import (
+	"errors"
+
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+var (
+	ErrNoBroker = errors.New("no broker")
+)
+
+func (balancer *PubBalancer) LookupTopicPartitions(topic *schema_pb.Topic) (assignments []*mq_pb.BrokerPartitionAssignment) {
+	// find existing topic partition assignments
+	for brokerStatsItem := range balancer.Brokers.IterBuffered() {
+		broker, brokerStats := brokerStatsItem.Key, brokerStatsItem.Val
+		for topicPartitionStatsItem := range brokerStats.TopicPartitionStats.IterBuffered() {
+			topicPartitionStat := topicPartitionStatsItem.Val
+			if topicPartitionStat.TopicPartition.Namespace == topic.Namespace &&
+				topicPartitionStat.TopicPartition.Name == topic.Name {
+				assignment := &mq_pb.BrokerPartitionAssignment{
+					Partition: &schema_pb.Partition{
+						RingSize:   MaxPartitionCount,
+						RangeStart: topicPartitionStat.RangeStart,
+						RangeStop:  topicPartitionStat.RangeStop,
+						UnixTimeNs: topicPartitionStat.UnixTimeNs,
+					},
+				}
+				// TODO fix follower setting
+				assignment.LeaderBroker = broker
+				assignments = append(assignments, assignment)
+			}
+		}
+	}
+	return
+}

weed/mq/pub_balancer/partition_list_broker.go

@@ -0,0 +1,60 @@
+package pub_balancer
+
+import (
+	"github.com/seaweedfs/seaweedfs/weed/glog"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+type PartitionSlotToBroker struct {
+	RangeStart     int32
+	RangeStop      int32
+	UnixTimeNs     int64
+	AssignedBroker string
+	FollowerBroker string
+}
+
+type PartitionSlotToBrokerList struct {
+	PartitionSlots []*PartitionSlotToBroker
+	RingSize       int32
+}
+
+func NewPartitionSlotToBrokerList(ringSize int32) *PartitionSlotToBrokerList {
+	return &PartitionSlotToBrokerList{
+		RingSize: ringSize,
+	}
+}
+
+func (ps *PartitionSlotToBrokerList) AddBroker(partition *schema_pb.Partition, broker string, follower string) {
+	for _, partitionSlot := range ps.PartitionSlots {
+		if partitionSlot.RangeStart == partition.RangeStart && partitionSlot.RangeStop == partition.RangeStop {
+			if partitionSlot.AssignedBroker != "" && partitionSlot.AssignedBroker != broker {
+				glog.V(0).Infof("partition %s broker change: %s => %s", partition, partitionSlot.AssignedBroker, broker)
+				partitionSlot.AssignedBroker = broker
+			}
+			if partitionSlot.FollowerBroker != "" && partitionSlot.FollowerBroker != follower {
+				glog.V(0).Infof("partition %s follower change: %s => %s", partition, partitionSlot.FollowerBroker, follower)
+				partitionSlot.FollowerBroker = follower
+			}
+
+			return
+		}
+	}
+	ps.PartitionSlots = append(ps.PartitionSlots, &PartitionSlotToBroker{
+		RangeStart:     partition.RangeStart,
+		RangeStop:      partition.RangeStop,
+		UnixTimeNs:     partition.UnixTimeNs,
+		AssignedBroker: broker,
+		FollowerBroker: follower,
+	})
+}
+func (ps *PartitionSlotToBrokerList) RemoveBroker(broker string) {
+	ps.ReplaceBroker(broker, "")
+}
+
+func (ps *PartitionSlotToBrokerList) ReplaceBroker(oldBroker string, newBroker string) {
+	for _, partitionSlot := range ps.PartitionSlots {
+		if partitionSlot.AssignedBroker == oldBroker {
+			partitionSlot.AssignedBroker = newBroker
+		}
+	}
+}

weed/mq/pub_balancer/pub_balancer.go

@@ -0,0 +1,102 @@
+package pub_balancer
+
+import (
+	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
+)
+
+const (
+	MaxPartitionCount  = 8 * 9 * 5 * 7 //2520
+	LockBrokerBalancer = "broker_balancer"
+)
+
+// PubBalancer collects stats from all brokers.
+//
+//	When publishers wants to create topics, it picks brokers to assign the topic partitions.
+//	When consumers wants to subscribe topics, it tells which brokers are serving the topic partitions.
+//
+// When a partition needs to be split or merged, or a partition needs to be moved to another broker,
+// the balancer will let the broker tell the consumer instance to stop processing the partition.
+// The existing consumer instance will flush the internal state, and then stop processing.
+// Then the balancer will tell the brokers to start sending new messages in the new/moved partition to the consumer instances.
+//
+// Failover to standby consumer instances:
+//
+//	A consumer group can have min and max number of consumer instances.
+//	For consumer instances joined after the max number, they will be in standby mode.
+//
+//	When a consumer instance is down, the broker will notice this and inform the balancer.
+//	The balancer will then tell the broker to send the partition to another standby consumer instance.
+type PubBalancer struct {
+	Brokers cmap.ConcurrentMap[string, *BrokerStats] // key: broker address
+	// Collected from all brokers when they connect to the broker leader
+	TopicToBrokers    cmap.ConcurrentMap[string, *PartitionSlotToBrokerList] // key: topic name
+	OnPartitionChange func(topic *schema_pb.Topic, assignments []*mq_pb.BrokerPartitionAssignment)
+}
+
+func NewPubBalancer() *PubBalancer {
+	return &PubBalancer{
+		Brokers:        cmap.New[*BrokerStats](),
+		TopicToBrokers: cmap.New[*PartitionSlotToBrokerList](),
+	}
+}
+
+func (balancer *PubBalancer) AddBroker(broker string) (brokerStats *BrokerStats) {
+	var found bool
+	brokerStats, found = balancer.Brokers.Get(broker)
+	if !found {
+		brokerStats = NewBrokerStats()
+		if !balancer.Brokers.SetIfAbsent(broker, brokerStats) {
+			brokerStats, _ = balancer.Brokers.Get(broker)
+		}
+	}
+	balancer.onPubAddBroker(broker, brokerStats)
+	return brokerStats
+}
+
+func (balancer *PubBalancer) RemoveBroker(broker string, stats *BrokerStats) {
+	balancer.Brokers.Remove(broker)
+
+	// update TopicToBrokers
+	for _, topic := range stats.Topics {
+		partitionSlotToBrokerList, found := balancer.TopicToBrokers.Get(topic.String())
+		if !found {
+			continue
+		}
+		pickedBroker := pickBrokers(balancer.Brokers, 1)
+		if len(pickedBroker) == 0 {
+			partitionSlotToBrokerList.RemoveBroker(broker)
+		} else {
+			partitionSlotToBrokerList.ReplaceBroker(broker, pickedBroker[0])
+		}
+	}
+	balancer.onPubRemoveBroker(broker, stats)
+}
+
+func (balancer *PubBalancer) OnBrokerStatsUpdated(broker string, brokerStats *BrokerStats, receivedStats *mq_pb.BrokerStats) {
+	brokerStats.UpdateStats(receivedStats)
+
+	// update TopicToBrokers
+	for _, topicPartitionStats := range receivedStats.Stats {
+		topicKey := topic.FromPbTopic(topicPartitionStats.Topic).String()
+		partition := topicPartitionStats.Partition
+		partitionSlotToBrokerList, found := balancer.TopicToBrokers.Get(topicKey)
+		if !found {
+			partitionSlotToBrokerList = NewPartitionSlotToBrokerList(MaxPartitionCount)
+			if !balancer.TopicToBrokers.SetIfAbsent(topicKey, partitionSlotToBrokerList) {
+				partitionSlotToBrokerList, _ = balancer.TopicToBrokers.Get(topicKey)
+			}
+		}
+		partitionSlotToBrokerList.AddBroker(partition, broker, topicPartitionStats.Follower)
+	}
+}
+
+// OnPubAddBroker is called when a broker is added for a publisher coordinator
+func (balancer *PubBalancer) onPubAddBroker(broker string, brokerStats *BrokerStats) {
+}
+
+// OnPubRemoveBroker is called when a broker is removed for a publisher coordinator
+func (balancer *PubBalancer) onPubRemoveBroker(broker string, brokerStats *BrokerStats) {
+}

weed/mq/pub_balancer/repair.go

@@ -0,0 +1,12 @@
+package pub_balancer
+
+import ()
+
+func (balancer *PubBalancer) RepairTopics() []BalanceAction {
+	action := BalanceTopicPartitionOnBrokers(balancer.Brokers)
+	return []BalanceAction{action}
+}
+
+type TopicPartitionInfo struct {
+	Broker string
+}