add/remove assigned partitions

This commit is contained in:
chrislu
2024-01-22 10:47:39 -08:00
parent 861ad732ca
commit e8b05ecc91
4 changed files with 104 additions and 72 deletions

View File

@@ -63,7 +63,7 @@ func (bs *BrokerStats) UpdateStats(stats *mq_pb.BrokerStats) {
}
func (bs *BrokerStats) RegisterAssignment(t *mq_pb.Topic, partition *mq_pb.Partition) {
func (bs *BrokerStats) RegisterAssignment(t *mq_pb.Topic, partition *mq_pb.Partition, isAdd bool) {
tps := &TopicPartitionStats{
TopicPartition: topic.TopicPartition{
Topic: topic.Topic{Namespace: t.Namespace, Name: t.Name},
@@ -78,5 +78,9 @@ func (bs *BrokerStats) RegisterAssignment(t *mq_pb.Topic, partition *mq_pb.Parti
IsLeader: true,
}
key := tps.TopicPartition.String()
bs.TopicPartitionStats.Set(key, tps)
if isAdd {
bs.TopicPartitionStats.SetIfAbsent(key, tps)
} else {
bs.TopicPartitionStats.Remove(key)
}
}