refactor

2024-01-21 01:24:12 -08:00
parent a5b906c0dd
commit 9ce7c482b3
2 changed files with 12 additions and 3 deletions
--- a/weed/mq/pub_balancer/partition_list_broker.go
+++ b/weed/mq/pub_balancer/partition_list_broker.go
@@ -44,9 +44,13 @@ func (ps *PartitionSlotToBrokerList) AddBroker(partition *mq_pb.Partition, broke
 	})
 }
 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 == broker {
-			partitionSlot.AssignedBroker = ""
+		if partitionSlot.AssignedBroker == oldBroker {
+			partitionSlot.AssignedBroker = newBroker
 		}
 	}
 }