receive broker stats

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
+}