feat(auto_balancer): allow change broker rack on register (AutoMQ#1771) (AutoMQ#1776)

Signed-off-by: Shichao Nie <niesc@automq.com>

Author: SCNieh

core/src/main/java/kafka/autobalancer/model/BrokerUpdater.java

@@ -24,7 +24,7 @@
 
 public class BrokerUpdater extends AbstractInstanceUpdater {
     private final int brokerId;
-    private final String rack;
+    private String rack;
     private boolean active;
 
     public BrokerUpdater(int brokerId, String rack, boolean active) {
@@ -41,6 +41,10 @@ public String rack() {
         return this.rack;
     }
 
+    public void setRack(String rack) {
+        this.rack = rack;
+    }
+
     public boolean isActive() {
         lock.lock();
         try {

core/src/main/java/kafka/autobalancer/model/ClusterModel.java

@@ -237,15 +237,19 @@ public boolean updateTopicPartitionMetrics(int brokerId, TopicPartition tp, Iter
     public void registerBroker(int brokerId, String rackId, boolean active) {
         clusterLock.lock();
         try {
-            if (brokerMap.containsKey(brokerId)) {
-                return;
-            }
             if (Utils.isBlank(rackId)) {
                 rackId = DEFAULT_RACK_ID;
             }
-            BrokerUpdater brokerUpdater = createBrokerUpdater(brokerId, rackId, active);
-            brokerMap.putIfAbsent(brokerId, brokerUpdater);
-            brokerReplicaMap.put(brokerId, new HashMap<>());
+            String finalRackId = rackId;
+            brokerMap.compute(brokerId, (id, brokerUpdater) -> {
+                if (brokerUpdater == null) {
+                    brokerReplicaMap.put(brokerId, new HashMap<>());
+                    return createBrokerUpdater(brokerId, finalRackId, active);
+                }
+                brokerUpdater.setRack(finalRackId);
+                brokerUpdater.setActive(active);
+                return brokerUpdater;
+            });
         } finally {
             clusterLock.unlock();
         }