broker report stats to balancer

weed/mq/balancer/balancer.go

@@ -1,16 +1,88 @@
 package balancer
 
 import (
+	"fmt"
 	cmap "github.com/orcaman/concurrent-map/v2"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+)
+
+const (
+	MaxPartitionCount = 1024
 )
 
 type Balancer struct {
 	Brokers cmap.ConcurrentMap[string, *BrokerStats]
 }
+
+func (b *Balancer) LookupOrAllocateTopicPartitions(topic *mq_pb.Topic, publish bool) ([]*mq_pb.BrokerPartitionAssignment, error) {
+	// TODO lock the topic
+
+	// find the topic partitions on the filer
+	// if the topic is not found
+	//   if the request is_for_publish
+	//     create the topic
+	//   if the request is_for_subscribe
+	//     return error not found
+	// t := topic.FromPbTopic(request.Topic)
+	return []*mq_pb.BrokerPartitionAssignment{
+		{
+			LeaderBroker:    "localhost:17777",
+			FollowerBrokers: []string{"localhost:17777"},
+			Partition: &mq_pb.Partition{
+				RingSize:   MaxPartitionCount,
+				RangeStart: 0,
+				RangeStop:  MaxPartitionCount,
+			},
+		},
+	}, nil
+}
+
 type BrokerStats struct {
 	TopicPartitionCount int32
 	ConsumerCount       int32
 	CpuUsagePercent     int32
+	Stats               cmap.ConcurrentMap[string, *TopicPartitionStats]
+}
+
+func (bs *BrokerStats) UpdateStats(stats *mq_pb.BrokerStats) {
+	bs.TopicPartitionCount = int32(len(stats.Stats))
+	bs.CpuUsagePercent = stats.CpuUsagePercent
+
+	var consumerCount int32
+	currentTopicPartitions := bs.Stats.Items()
+	for _, topicPartitionStats := range stats.Stats {
+		tps := &TopicPartitionStats{
+			TopicPartition: TopicPartition{
+				Namespace:  topicPartitionStats.Topic.Namespace,
+				Topic:      topicPartitionStats.Topic.Name,
+				RangeStart: topicPartitionStats.Partition.RangeStart,
+				RangeStop:  topicPartitionStats.Partition.RangeStop,
+			},
+			ConsumerCount: topicPartitionStats.ConsumerCount,
+		}
+		consumerCount += topicPartitionStats.ConsumerCount
+		key := tps.TopicPartition.String()
+		bs.Stats.Set(key, tps)
+		delete(currentTopicPartitions, key)
+	}
+	// remove the topic partitions that are not in the stats
+	for key := range currentTopicPartitions {
+		bs.Stats.Remove(key)
+	}
+	bs.ConsumerCount = consumerCount
+
+}
+
+type TopicPartition struct {
+	Namespace  string
+	Topic      string
+	RangeStart int32
+	RangeStop  int32
+}
+
+type TopicPartitionStats struct {
+	TopicPartition
+	ConsumerCount int32
 }
 
 func NewBalancer() *Balancer {
@@ -20,5 +92,11 @@ func NewBalancer() *Balancer {
 }
 
 func NewBrokerStats() *BrokerStats {
-	return &BrokerStats{}
+	return &BrokerStats{
+		Stats: cmap.New[*TopicPartitionStats](),
+	}
+}
+
+func (tp *TopicPartition) String() string {
+	return fmt.Sprintf("%v.%v-%04d-%04d", tp.Namespace, tp.Topic, tp.RangeStart, tp.RangeStop)
 }

weed/mq/broker/broker_grpc_admin.go

@@ -12,10 +12,6 @@ import (
 	"sync"
 )
 
-const (
-	MaxPartitionCount = 1024
-)
-
 func (broker *MessageQueueBroker) FindBrokerLeader(c context.Context, request *mq_pb.FindBrokerLeaderRequest) (*mq_pb.FindBrokerLeaderResponse, error) {
 	ret := &mq_pb.FindBrokerLeaderResponse{}
 	err := broker.withMasterClient(false, broker.MasterClient.GetMaster(), func(client master_pb.SeaweedClient) error {

weed/mq/broker/broker_grpc_balancer.go

@@ -40,9 +40,7 @@ func (broker *MessageQueueBroker) ConnectToBalancer(stream mq_pb.SeaweedMessagin
 			return status.Errorf(codes.Unavailable, "not current broker balancer")
 		}
 		if receivedStats := req.GetStats(); receivedStats != nil {
-			brokerStats.TopicPartitionCount = receivedStats.TopicPartitionCount
+			brokerStats.UpdateStats(receivedStats)
-			brokerStats.ConsumerCount = receivedStats.ConsumerCount
-			brokerStats.CpuUsagePercent = receivedStats.CpuUsagePercent
 
 			glog.V(3).Infof("broker %s stats: %+v", initMessage.Broker, brokerStats)
 		}

weed/mq/broker/broker_grpc_lookup.go

@@ -3,6 +3,8 @@ package broker
 import (
 	"context"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // FindTopicBrokers returns the brokers that are serving the topic
@@ -17,30 +19,25 @@ import (
 //     2.2 if the topic is found, return the brokers
 //
 //  3. unlock the topic
-func (broker *MessageQueueBroker) LookupTopicBrokers(ctx context.Context, request *mq_pb.LookupTopicBrokersRequest) (*mq_pb.LookupTopicBrokersResponse, error) {
+func (broker *MessageQueueBroker) LookupTopicBrokers(ctx context.Context, request *mq_pb.LookupTopicBrokersRequest) (resp *mq_pb.LookupTopicBrokersResponse, err error) {
-	ret := &mq_pb.LookupTopicBrokersResponse{}
+	if broker.currentBalancer == "" {
-	// TODO lock the topic
+		return nil, status.Errorf(codes.Unavailable, "no balancer")
-
-	// find the topic partitions on the filer
-	// if the topic is not found
-	//   if the request is_for_publish
-	//     create the topic
-	//   if the request is_for_subscribe
-	//     return error not found
-	// t := topic.FromPbTopic(request.Topic)
-	ret.Topic = request.Topic
-	ret.BrokerPartitionAssignments = []*mq_pb.BrokerPartitionAssignment{
-		{
-			LeaderBroker:    "localhost:17777",
-			FollowerBrokers: []string{"localhost:17777"},
-			Partition: &mq_pb.Partition{
-				RingSize:   MaxPartitionCount,
-				RangeStart: 0,
-				RangeStop:  MaxPartitionCount,
-			},
-		},
 	}
-	return ret, nil
+	if !broker.lockAsBalancer.IsLocked() {
+		proxyErr := broker.withBrokerClient(false, broker.currentBalancer, func(client mq_pb.SeaweedMessagingClient) error {
+			resp, err = client.LookupTopicBrokers(ctx, request)
+			return nil
+		})
+		if proxyErr != nil {
+			return nil, proxyErr
+		}
+		return resp, err
+	}
+
+	ret := &mq_pb.LookupTopicBrokersResponse{}
+	ret.Topic = request.Topic
+	ret.BrokerPartitionAssignments, err = broker.Balancer.LookupOrAllocateTopicPartitions(ret.Topic, request.IsForPublish)
+	return ret, err
 }
 
 // CheckTopicPartitionsStatus check the topic partitions on the broker

weed/mq/broker/broker_server.go

@@ -39,6 +39,7 @@ type MessageQueueBroker struct {
 	localTopicManager *topic.LocalTopicManager
 	Balancer          *balancer.Balancer
 	lockAsBalancer    *cluster.LiveLock
+	currentBalancer   pb.ServerAddress
 }
 
 func NewMessageBroker(option *MessageQueueBrokerOption, grpcDialOption grpc.DialOption) (mqBroker *MessageQueueBroker, err error) {

weed/mq/broker/broker_stats.go

@@ -34,6 +34,7 @@ func (broker *MessageQueueBroker) BrokerConnectToBalancer(self string) error {
 	}
 
 	glog.V(1).Infof("broker %s found balancer %s", self, brokerBalancer)
+	broker.currentBalancer = pb.ServerAddress(brokerBalancer)
 
 	// connect to the lock owner
 	err = pb.WithBrokerGrpcClient(false, brokerBalancer, broker.grpcDialOption, func(client mq_pb.SeaweedMessagingClient) error {

weed/mq/client/pub_client/publish.go

@@ -2,13 +2,13 @@ package pub_client
 
 import (
 	"fmt"
-	"github.com/seaweedfs/seaweedfs/weed/mq/broker"
+	"github.com/seaweedfs/seaweedfs/weed/mq/balancer"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"github.com/seaweedfs/seaweedfs/weed/util"
 )
 
 func (p *TopicPublisher) Publish(key, value []byte) error {
-	hashKey := util.HashToInt32(key) % broker.MaxPartitionCount
+	hashKey := util.HashToInt32(key) % balancer.MaxPartitionCount
 	if hashKey < 0 {
 		hashKey = -hashKey
 	}

weed/mq/client/pub_client/publisher.go

@@ -2,7 +2,7 @@ package pub_client
 
 import (
 	"github.com/rdleal/intervalst/interval"
-	"github.com/seaweedfs/seaweedfs/weed/mq/broker"
+	"github.com/seaweedfs/seaweedfs/weed/mq/balancer"
 	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -46,7 +46,7 @@ func (p *TopicPublisher) Connect(bootstrapBroker string) error {
 
 func (p *TopicPublisher) Shutdown() error {
 
-	if clients, found := p.partition2Broker.AllIntersections(0, broker.MaxPartitionCount); found {
+	if clients, found := p.partition2Broker.AllIntersections(0, balancer.MaxPartitionCount); found {
 		for _, client := range clients {
 			client.CloseSend()
 		}

weed/mq/topic/local_manager.go

@@ -58,11 +58,23 @@ func (manager *LocalTopicManager) RemoveTopicPartition(topic Topic, partition Pa
 }
 
 func (manager *LocalTopicManager) CollectStats(duration time.Duration) *mq_pb.BrokerStats {
-	stats := &mq_pb.BrokerStats{}
+	stats := &mq_pb.BrokerStats{
+		Stats: make(map[string]*mq_pb.TopicPartitionStats),
+	}
 	manager.topics.IterCb(func(topic string, localTopic *LocalTopic) {
 		for _, localPartition := range localTopic.Partitions {
-			stats.TopicPartitionCount++
+			stats.Stats[topic] = &mq_pb.TopicPartitionStats{
-			stats.ConsumerCount += localPartition.ConsumerCount
+				Topic: &mq_pb.Topic{
+					Namespace: string(localTopic.Namespace),
+					Name:      localTopic.Name,
+				},
+				Partition: &mq_pb.Partition{
+					RingSize:   localPartition.RingSize,
+					RangeStart: localPartition.RangeStart,
+					RangeStop:  localPartition.RangeStop,
+				},
+				ConsumerCount: localPartition.ConsumerCount,
+			}
 		}
 	})
 

weed/pb/mq.proto

@@ -105,10 +105,15 @@ message CheckBrokerLoadResponse {
 
 //////////////////////////////////////////////////
 message BrokerStats {
-    int32 topic_partition_count = 1;
+    int32 cpu_usage_percent = 1;
-    int32 consumer_count = 2;
+    map<string, TopicPartitionStats> stats = 2;
-    int32 cpu_usage_percent = 3;
 }
+message TopicPartitionStats {
+    Topic topic = 1;
+    Partition partition = 2;
+    int32 consumer_count = 3;
+}
+
 message ConnectToBalancerRequest {
     message InitMessage {
         string broker = 1;