Repository: skiptests/astraea
Pull request #640 (2022-08-26 20:06:05)

Change balancer to Use NodeTopicSizeCost to evaluate cost
1. `NodeTopicSizeCost` 提供 `ClusterCost`
2. 撰寫新的`ClusterInfo`產生方式，盡量保持`Replica`資訊的完整性
3. `Balancer APIs`現在改用`topic size`作為評估手段


Patch:
diff --git a/app/src/main/java/org/astraea/app/admin/ClusterInfo.java b/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
index f6adbff7..68dcf010 100644
--- a/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
+++ b/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
@@ -45,13 +45,16 @@ public interface ClusterInfo {
    * @return ClusterInfo
    */
   static ClusterInfo of(org.apache.kafka.common.Cluster cluster) {
-    var nodes = cluster.nodes().stream().map(NodeInfo::of).collect(Collectors.toUnmodifiableList());
-    var topics = cluster.topics();
-    var replicas =
-        topics.stream()
+    return of(
+        cluster.nodes().stream().map(NodeInfo::of).collect(Collectors.toUnmodifiableList()),
+        cluster.topics().stream()
             .flatMap(t -> cluster.partitionsForTopic(t).stream())
             .flatMap(p -> ReplicaInfo.of(p).stream())
-            .collect(Collectors.toUnmodifiableList());
+            .collect(Collectors.toUnmodifiableList()));
+  }
+
+  static ClusterInfo of(List<NodeInfo> nodes, List<ReplicaInfo> replicas) {
+    var topics = replicas.stream().map(ReplicaInfo::topic).collect(Collectors.toUnmodifiableSet());
     var replicasForTopic = replicas.stream().collect(Collectors.groupingBy(ReplicaInfo::topic));
     var availableReplicasForTopic =
         replicas.stream()
diff --git a/app/src/main/java/org/astraea/app/balancer/BalancerUtils.java b/app/src/main/java/org/astraea/app/balancer/BalancerUtils.java
index 95281fed..7f8ecdd6 100644
--- a/app/src/main/java/org/astraea/app/balancer/BalancerUtils.java
+++ b/app/src/main/java/org/astraea/app/balancer/BalancerUtils.java
@@ -31,12 +31,62 @@ import org.astraea.app.admin.ClusterInfo;
 import org.astraea.app.admin.NodeInfo;
 import org.astraea.app.admin.Replica;
 import org.astraea.app.admin.ReplicaInfo;
+import org.astraea.app.admin.TopicPartition;
 import org.astraea.app.balancer.log.ClusterLogAllocation;
 import org.astraea.app.cost.HasClusterCost;
 import org.astraea.app.metrics.HasBeanObject;
 
 public class BalancerUtils {
 
+  /**
+   * Update the replicas of ClusterInfo according to ClusterLogAllocation. Noted that only "broker"
+   * and "data folder" get updated. The replicas matched to nothing from ClusterLogAllocation won't
+   * get any update.
+   *
+   * @param clusterInfo to get updated
+   * @param allocation offers new host and data folder
+   * @return new cluster info
+   */
+  public static ClusterInfo update(ClusterInfo clusterInfo, ClusterLogAllocation allocation) {
+    var newReplicas =
+        clusterInfo.replicas().stream()
+            .collect(Collectors.groupingBy(r -> TopicPartition.of(r.topic(), r.partition())))
+            .entrySet()
+            .stream()
+            .flatMap(
+                entry -> {
+                  var lps = allocation.logPlacements(entry.getKey());
+                  var replicas =
+                      entry.getValue().stream()
+                          .map(r -> (Replica) r)
+                          .collect(Collectors.toUnmodifiableList());
+                  return IntStream.range(0, replicas.size())
+                      .mapToObj(
+                          index -> {
+                            var previous = replicas.get(index);
+                            // return previous replica due to no new information
+                            if (index >= lps.size()) return previous;
+                            var lp = lps.get(index);
+                            return (ReplicaInfo)
+                                Replica.of(
+                                    previous.topic(),
+                                    previous.partition(),
+                                    clusterInfo.node(lp.broker()),
+                                    previous.lag(),
+                                    previous.size(),
+                                    index == 0,
+                                    previous.inSync(),
+                                    previous.isFuture(),
+                                    previous.isOffline(),
+                                    previous.isPreferredLeader(),
+                                    lp.dataFolder());
+                          });
+                })
+            .collect(Collectors.toUnmodifiableList());
+
+    return ClusterInfo.of(clusterInfo.nodes(), newReplicas);
+  }
+
   /**
    * Create a {@link ClusterInfo} with its log placement replaced by {@link ClusterLogAllocation}.
    * Every log will be marked as online & synced. Based on the given content in {@link
diff --git a/app/src/main/java/org/astraea/app/cost/NodeTopicSizeCost.java b/app/src/main/java/org/astraea/app/cost/NodeTopicSizeCost.java
index 07686a86..0db27e09 100644
--- a/app/src/main/java/org/astraea/app/cost/NodeTopicSizeCost.java
+++ b/app/src/main/java/org/astraea/app/cost/NodeTopicSizeCost.java
@@ -21,10 +21,14 @@ import java.util.Optional;
 import java.util.stream.Collectors;
 import org.astraea.app.admin.ClusterBean;
 import org.astraea.app.admin.ClusterInfo;
+import org.astraea.app.admin.NodeInfo;
+import org.astraea.app.admin.Replica;
 import org.astraea.app.metrics.broker.LogMetrics;
 import org.astraea.app.metrics.collector.Fetcher;
 
-public class NodeTopicSizeCost implements HasBrokerCost {
+public class NodeTopicSizeCost implements HasBrokerCost, HasClusterCost {
+  private final Dispersion dispersion = Dispersion.correlationCoefficient();
+
   @Override
   public Optional<Fetcher> fetcher() {
     return Optional.of(LogMetrics.Log.SIZE::fetch);
@@ -47,4 +51,21 @@ public class NodeTopicSizeCost implements HasBrokerCost {
                             .sum()));
     return () -> result;
   }
+
+  @Override
+  public ClusterCost clusterCost(ClusterInfo clusterInfo, ClusterBean clusterBean) {
+    var brokerCost =
+        clusterInfo.nodes().stream()
+            .collect(
+                Collectors.toMap(
+                    NodeInfo::id,
+                    nodeInfo ->
+                        clusterInfo.replicas().stream()
+                            .filter(r -> r.nodeInfo().id() == nodeInfo.id())
+                            .mapToLong(r -> r instanceof Replica ? ((Replica) r).size() : 0)
+                            .sum()));
+    return () ->
+        dispersion.calculate(
+            brokerCost.values().stream().map(v -> (double) v).collect(Collectors.toList()));
+  }
 }
diff --git a/app/src/main/java/org/astraea/app/web/BalancerHandler.java b/app/src/main/java/org/astraea/app/web/BalancerHandler.java
index 115e196b..c19dfcbf 100644
--- a/app/src/main/java/org/astraea/app/web/BalancerHandler.java
+++ b/app/src/main/java/org/astraea/app/web/BalancerHandler.java
@@ -28,7 +28,7 @@ import org.astraea.app.balancer.generator.RebalancePlanGenerator;
 import org.astraea.app.balancer.log.ClusterLogAllocation;
 import org.astraea.app.balancer.log.LogPlacement;
 import org.astraea.app.cost.HasClusterCost;
-import org.astraea.app.cost.ReplicaLeaderCost;
+import org.astraea.app.cost.NodeTopicSizeCost;
 
 class BalancerHandler implements Handler {
 
@@ -40,7 +40,7 @@ class BalancerHandler implements Handler {
   private final HasClusterCost costFunction;
 
   BalancerHandler(Admin admin) {
-    this(admin, new ReplicaLeaderCost());
+    this(admin, new NodeTopicSizeCost());
   }
 
   BalancerHandler(Admin admin, HasClusterCost costFunction) {
@@ -50,7 +50,7 @@ class BalancerHandler implements Handler {
 
   @Override
   public Response get(Channel channel) {
-    var clusterInfo = admin.clusterInfo();
+    var clusterInfo = admin.clusterInfo(admin.topicNames(false));
     var clusterAllocation = ClusterLogAllocation.of(clusterInfo);
     var cost = costFunction.clusterCost(clusterInfo, ClusterBean.EMPTY).value();
     var limit =
@@ -65,7 +65,7 @@ class BalancerHandler implements Handler {
                     Map.entry(
                         cla,
                         costFunction
-                            .clusterCost(BalancerUtils.merge(clusterInfo, cla), ClusterBean.EMPTY)
+                            .clusterCost(BalancerUtils.update(clusterInfo, cla), ClusterBean.EMPTY)
                             .value()))
             .filter(e -> e.getValue() <= cost)
             .min(Comparator.comparingDouble(Map.Entry::getValue));


Test patch:
diff --git a/app/src/test/java/org/astraea/app/balancer/BalancerUtilsIntegratedTest.java b/app/src/test/java/org/astraea/app/balancer/BalancerUtilsIntegratedTest.java
new file mode 100644
index 00000000..a67384e6
--- /dev/null
+++ b/app/src/test/java/org/astraea/app/balancer/BalancerUtilsIntegratedTest.java
@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.astraea.app.balancer;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.IntStream;
+import org.astraea.app.admin.Admin;
+import org.astraea.app.admin.Replica;
+import org.astraea.app.admin.TopicPartition;
+import org.astraea.app.balancer.log.ClusterLogAllocation;
+import org.astraea.app.balancer.log.LogPlacement;
+import org.astraea.app.common.Utils;
+import org.astraea.app.producer.Producer;
+import org.astraea.app.service.RequireBrokerCluster;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class BalancerUtilsIntegratedTest extends RequireBrokerCluster {
+
+  @Test
+  void testUpdate() {
+    var topicName = Utils.randomString(5);
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(topicName).create();
+      Utils.sleep(Duration.ofSeconds(3));
+
+      try (var producer = Producer.of(bootstrapServers())) {
+        IntStream.range(0, 100)
+            .forEach(ignored -> producer.sender().topic(topicName).key(new byte[10]).run());
+      }
+
+      var clusterInfo = admin.clusterInfo(Set.of(topicName));
+      clusterInfo.replicas().forEach(r -> Assertions.assertTrue(((Replica) r).size() > 0));
+
+      var replica = clusterInfo.replicas().iterator().next();
+      var newBrokerId =
+          brokerIds().stream().filter(id -> id != replica.nodeInfo().id()).findFirst().get();
+
+      var merged =
+          BalancerUtils.update(
+              clusterInfo,
+              ClusterLogAllocation.of(
+                  Map.of(
+                      TopicPartition.of(topicName, 0),
+                      // change the broker
+                      List.of(LogPlacement.of(newBrokerId, ((Replica) replica).dataFolder())))));
+
+      Assertions.assertEquals(clusterInfo.replicas().size(), merged.replicas().size());
+      Assertions.assertEquals(clusterInfo.topics().size(), merged.topics().size());
+      merged.replicas().forEach(r -> Assertions.assertTrue(((Replica) r).size() > 0));
+      Assertions.assertEquals(1, merged.replicas(topicName).size());
+      Assertions.assertEquals(newBrokerId, merged.replicas(topicName).get(0).nodeInfo().id());
+    }
+  }
+}