receive broker stats

weed/mq/balancer/balancer.go

@@ -1,12 +1,18 @@
 package balancer
 
-import cmap "github.com/orcaman/concurrent-map"
+import (
+	"fmt"
+	cmap "github.com/orcaman/concurrent-map/v2"
+)
 
 type Balancer struct {
-	brokers cmap.ConcurrentMap[string, *BrokerStats]
+	Brokers cmap.ConcurrentMap[string, *BrokerStats]
 }
 type BrokerStats struct {
-	stats map[TopicPartition]*TopicPartitionStats
+	TopicPartitionCount int32
+	MessageCount        int64
+	BytesCount          int64
+	CpuUsagePercent     int32
 }
 
 type TopicPartition struct {
@@ -16,5 +22,22 @@ type TopicPartition struct {
 }
 
 type TopicPartitionStats struct {
-	Throughput int64
+	TopicPartition
+	Throughput          int64
+	ConsumerCount       int64
+	TopicPartitionCount int64
+}
+
+func NewBalancer() *Balancer {
+	return &Balancer{
+		Brokers: cmap.New[*BrokerStats](),
+	}
+}
+
+func NewBrokerStats() *BrokerStats {
+	return &BrokerStats{}
+}
+
+func (tp *TopicPartition) String() string {
+	return fmt.Sprintf("%v-%04d-%04d", tp.Topic, tp.RangeStart, tp.RangeStop)
 }

weed/mq/broker/broker_grpc_balancer.go

@@ -0,0 +1,41 @@
+package broker
+
+import (
+	"github.com/seaweedfs/seaweedfs/weed/mq/balancer"
+	"github.com/seaweedfs/seaweedfs/weed/pb/mq_pb"
+)
+
+func (broker *MessageQueueBroker) ConnectToBalancer(stream mq_pb.SeaweedMessaging_ConnectToBalancerServer) error {
+	req, err := stream.Recv()
+	if err != nil {
+		return err
+	}
+	response := &mq_pb.ConnectToBalancerResponse{}
+	initMessage := req.GetInit()
+	brokerStats := balancer.NewBrokerStats()
+	if initMessage != nil {
+		broker.Balancer.Brokers.Set(initMessage.Broker, brokerStats)
+	} else {
+		response.Error = "balancer init message is empty"
+		return stream.Send(response)
+	}
+	defer func() {
+		broker.Balancer.Brokers.Remove(initMessage.Broker)
+	}()
+	stream.Send(response)
+
+	for {
+		req, err := stream.Recv()
+		if err != nil {
+			return err
+		}
+		if receivedStats := req.GetStats(); receivedStats != nil {
+			brokerStats.TopicPartitionCount = receivedStats.TopicPartitionCount
+			brokerStats.MessageCount = receivedStats.MessageCount
+			brokerStats.BytesCount = receivedStats.BytesCount
+			brokerStats.CpuUsagePercent = receivedStats.CpuUsagePercent
+		}
+	}
+
+	return nil
+}

weed/mq/broker/broker_server.go

@@ -1,6 +1,7 @@
 package broker
 
 import (
+	"github.com/seaweedfs/seaweedfs/weed/mq/balancer"
 	"github.com/seaweedfs/seaweedfs/weed/mq/topic"
 	"time"
 
@@ -34,6 +35,7 @@ type MessageQueueBroker struct {
 	filers            map[pb.ServerAddress]struct{}
 	currentFiler      pb.ServerAddress
 	localTopicManager *topic.LocalTopicManager
+	Balancer          *balancer.Balancer
 }
 
 func NewMessageBroker(option *MessageQueueBrokerOption, grpcDialOption grpc.DialOption) (mqBroker *MessageQueueBroker, err error) {
@@ -41,9 +43,10 @@ func NewMessageBroker(option *MessageQueueBrokerOption, grpcDialOption grpc.Dial
 	mqBroker = &MessageQueueBroker{
 		option:            option,
 		grpcDialOption:    grpcDialOption,
-		MasterClient:   wdclient.NewMasterClient(grpcDialOption, option.FilerGroup, cluster.BrokerType, pb.NewServerAddress(option.Ip, option.Port, 0), option.DataCenter, option.Rack, *pb.NewServiceDiscoveryFromMap(option.Masters)),
+		MasterClient:      wdclient.NewMasterClient(grpcDialOption, option.FilerGroup, cluster.BrokerType, pb.NewServerAddress(option.Ip, option.Port, 0), option.DataCenter, option.Rack, *pb.NewServiceDiscoveryFromMap(option.Masters)),
 		filers:            make(map[pb.ServerAddress]struct{}),
 		localTopicManager: topic.NewLocalTopicManager(),
+		Balancer:          balancer.NewBalancer(),
 	}
 	mqBroker.MasterClient.SetOnPeerUpdateFn(mqBroker.OnBrokerUpdate)