Repository: skiptests/astraea
Pull request #642 (2022-08-27 09:05:21)

Show the size of replica when running balancer plan
1. 每個產生的規劃中會在來源地顯示replica大小，幫助閱讀
2. 重構幾隻常用的資料結構轉換


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 68dcf010..4289df4d 100644
--- a/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
+++ b/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
@@ -19,6 +19,7 @@ package org.astraea.app.admin;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -187,6 +188,14 @@ public interface ClusterInfo {
         .collect(Collectors.toUnmodifiableList());
   }
 
+  /**
+   * @param replica to search
+   * @return the replica matched to input replica
+   */
+  default Optional<ReplicaInfo> replica(TopicPartitionReplica replica) {
+    return replicas().stream().filter(r -> r.topicPartitionReplica().equals(replica)).findFirst();
+  }
+
   /** @return all replicas cached by this cluster info. */
   List<ReplicaInfo> replicas();
 }
diff --git a/app/src/main/java/org/astraea/app/admin/ReplicaInfo.java b/app/src/main/java/org/astraea/app/admin/ReplicaInfo.java
index f74440a7..02536e5d 100644
--- a/app/src/main/java/org/astraea/app/admin/ReplicaInfo.java
+++ b/app/src/main/java/org/astraea/app/admin/ReplicaInfo.java
@@ -113,6 +113,15 @@ public interface ReplicaInfo {
     };
   }
 
+  /**
+   * a helper to build TopicPartitionReplica quickly
+   *
+   * @return TopicPartitionReplica
+   */
+  default TopicPartitionReplica topicPartitionReplica() {
+    return TopicPartitionReplica.of(topic(), partition(), nodeInfo().id());
+  }
+
   /** @return topic name */
   String topic();
 
diff --git a/app/src/main/java/org/astraea/app/balancer/log/ClusterLogAllocation.java b/app/src/main/java/org/astraea/app/balancer/log/ClusterLogAllocation.java
index 9e9b1a95..e4c649a1 100644
--- a/app/src/main/java/org/astraea/app/balancer/log/ClusterLogAllocation.java
+++ b/app/src/main/java/org/astraea/app/balancer/log/ClusterLogAllocation.java
@@ -42,31 +42,24 @@ public interface ClusterLogAllocation {
         clusterInfo.replicas().stream()
             .filter(r -> r instanceof Replica)
             .map(r -> (Replica) r)
-            .collect(
-                Collectors.groupingBy(
-                    replica ->
-                        TopicPartition.of(replica.topic(), Integer.toString(replica.partition()))))
+            .collect(Collectors.groupingBy(r -> TopicPartition.of(r.topic(), r.partition())))
             .entrySet()
             .stream()
-            .map(
-                entry -> {
-                  // validate if the given log placements are valid
-                  if (entry.getValue().stream().filter(ReplicaInfo::isLeader).count() != 1)
-                    throw new IllegalArgumentException(
-                        "The " + entry.getKey() + " leader count mismatch 1.");
-
-                  final var topicPartition = entry.getKey();
-                  final var logPlacements =
-                      entry.getValue().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry -> {
+                      // validate if the given log placements are valid
+                      if (entry.getValue().stream().filter(ReplicaInfo::isLeader).count() != 1)
+                        throw new IllegalArgumentException(
+                            "The " + entry.getKey() + " leader count mismatch 1.");
+                      return entry.getValue().stream()
                           .sorted(Comparator.comparingInt(replica -> replica.isLeader() ? 0 : 1))
                           .map(
                               replica ->
                                   LogPlacement.of(replica.nodeInfo().id(), replica.dataFolder()))
                           .collect(Collectors.toList());
-
-                  return Map.entry(topicPartition, logPlacements);
-                })
-            .collect(Collectors.toConcurrentMap(Map.Entry::getKey, Map.Entry::getValue)));
+                    })));
   }
 
   static ClusterLogAllocation of(Map<TopicPartition, List<LogPlacement>> allocation) {
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 c19dfcbf..5a621092 100644
--- a/app/src/main/java/org/astraea/app/web/BalancerHandler.java
+++ b/app/src/main/java/org/astraea/app/web/BalancerHandler.java
@@ -19,9 +19,12 @@ package org.astraea.app.web;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.astraea.app.admin.Admin;
 import org.astraea.app.admin.ClusterBean;
+import org.astraea.app.admin.Replica;
+import org.astraea.app.admin.TopicPartitionReplica;
 import org.astraea.app.balancer.BalancerUtils;
 import org.astraea.app.balancer.RebalancePlanProposal;
 import org.astraea.app.balancer.generator.RebalancePlanGenerator;
@@ -86,14 +89,25 @@ class BalancerHandler implements Handler {
                                 new Change(
                                     tp.topic(),
                                     tp.partition(),
-                                    placements(clusterAllocation.logPlacements(tp)),
-                                    placements(entry.getKey().logPlacements(tp))))
+                                    // only log the size from source replicas
+                                    placements(
+                                        clusterAllocation.logPlacements(tp),
+                                        l ->
+                                            clusterInfo
+                                                .replica(
+                                                    TopicPartitionReplica.of(
+                                                        tp.topic(), tp.partition(), l.broker()))
+                                                .map(r -> ((Replica) r).size())
+                                                .orElse(null)),
+                                    placements(entry.getKey().logPlacements(tp), ignored -> null)))
                         .collect(Collectors.toUnmodifiableList()))
             .orElse(List.of()));
   }
 
-  static List<Placement> placements(List<LogPlacement> lps) {
-    return lps.stream().map(Placement::new).collect(Collectors.toUnmodifiableList());
+  static List<Placement> placements(List<LogPlacement> lps, Function<LogPlacement, Long> size) {
+    return lps.stream()
+        .map(p -> new Placement(p, size.apply(p)))
+        .collect(Collectors.toUnmodifiableList());
   }
 
   static class Placement {
@@ -101,9 +115,12 @@ class BalancerHandler implements Handler {
     final int brokerId;
     final String directory;
 
-    Placement(LogPlacement lp) {
+    final Long size;
+
+    Placement(LogPlacement lp, Long size) {
       this.brokerId = lp.broker();
       this.directory = lp.dataFolder();
+      this.size = size;
     }
   }
 
diff --git a/docs/web_server/web_api_balancer_chinese.md b/docs/web_server/web_api_balancer_chinese.md
index 7913d6eb..7bfef529 100644
--- a/docs/web_server/web_api_balancer_chinese.md
+++ b/docs/web_server/web_api_balancer_chinese.md
@@ -28,6 +28,9 @@ JSON Response 範例
   - `topic`: topic 名稱
   - `partition`: partition id
   - `before`: 原本的配置
+    - `brokerId`: 有掌管 replica 的節點 id
+    - `directory`: replica 存在資料的路徑
+    - `size`: replica 在硬碟上的資料大小
   - `after`: 比較好的配置
 ```json
 {
@@ -42,7 +45,8 @@ JSON Response 範例
       "before": [
         {
           "brokerId": 1006,
-          "directory": "/tmp/log-folder-0"
+          "directory": "/tmp/log-folder-0",
+          "size": 1234
         }
       ],
       "after": [
@@ -58,7 +62,8 @@ JSON Response 範例
       "before": [
         {
           "brokerId": 1003,
-          "directory": "/tmp/log-folder-0"
+          "directory": "/tmp/log-folder-0",
+          "size": 12355
         }
       ],
       "after": [


Test patch:
diff --git a/app/src/test/java/org/astraea/app/cost/ClusterInfoIntegratedTest.java b/app/src/test/java/org/astraea/app/cost/ClusterInfoIntegratedTest.java
new file mode 100644
index 00000000..4dbd3556
--- /dev/null
+++ b/app/src/test/java/org/astraea/app/cost/ClusterInfoIntegratedTest.java
@@ -0,0 +1,58 @@
+/*
+ * 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.cost;
+
+import java.time.Duration;
+import java.util.Optional;
+import org.astraea.app.admin.Admin;
+import org.astraea.app.common.Utils;
+import org.astraea.app.service.RequireBrokerCluster;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class ClusterInfoIntegratedTest extends RequireBrokerCluster {
+
+  @Test
+  void testQuery() {
+    try (var admin = Admin.of(bootstrapServers())) {
+      admin.creator().topic(Utils.randomString()).numberOfPartitions(10).create();
+      Utils.sleep(Duration.ofSeconds(2));
+
+      var clusterInfo = admin.clusterInfo();
+
+      // search by replica
+      clusterInfo
+          .replicas()
+          .forEach(
+              r ->
+                  Assertions.assertNotEquals(
+                      Optional.empty(), clusterInfo.replica(r.topicPartitionReplica())));
+
+      // search by topic
+      clusterInfo
+          .topics()
+          .forEach(t -> Assertions.assertNotEquals(0, clusterInfo.replicas(t).size()));
+      clusterInfo
+          .topics()
+          .forEach(t -> Assertions.assertNotEquals(0, clusterInfo.availableReplicas(t).size()));
+      clusterInfo
+          .topics()
+          .forEach(
+              t -> Assertions.assertNotEquals(0, clusterInfo.availableReplicaLeaders(t).size()));
+    }
+  }
+}
diff --git a/app/src/test/java/org/astraea/app/web/BalancerHandlerTest.java b/app/src/test/java/org/astraea/app/web/BalancerHandlerTest.java
index 3ece0fa2..d1a90383 100644
--- a/app/src/test/java/org/astraea/app/web/BalancerHandlerTest.java
+++ b/app/src/test/java/org/astraea/app/web/BalancerHandlerTest.java
@@ -16,15 +16,18 @@
  */
 package org.astraea.app.web;
 
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.IntStream;
 import org.astraea.app.admin.Admin;
 import org.astraea.app.admin.ClusterBean;
 import org.astraea.app.admin.ClusterInfo;
 import org.astraea.app.common.Utils;
 import org.astraea.app.cost.ClusterCost;
 import org.astraea.app.cost.HasClusterCost;
+import org.astraea.app.producer.Producer;
 import org.astraea.app.service.RequireBrokerCluster;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -35,8 +38,18 @@ public class BalancerHandlerTest extends RequireBrokerCluster {
   void testReport() {
     var topicName = Utils.randomString(10);
     try (var admin = Admin.of(bootstrapServers())) {
-      admin.creator().topic(topicName).numberOfPartitions(10).numberOfReplicas((short) 3).create();
+      admin.creator().topic(topicName).numberOfPartitions(3).numberOfReplicas((short) 1).create();
       Utils.sleep(Duration.ofSeconds(3));
+      try (var producer = Producer.of(bootstrapServers())) {
+        IntStream.range(0, 30)
+            .forEach(
+                index ->
+                    producer
+                        .sender()
+                        .topic(topicName)
+                        .key(String.valueOf(index).getBytes(StandardCharsets.UTF_8))
+                        .run());
+      }
       var handler = new BalancerHandler(admin, new MyCost());
       var report =
           Assertions.assertInstanceOf(
@@ -46,6 +59,14 @@ public class BalancerHandlerTest extends RequireBrokerCluster {
       Assertions.assertNotEquals(0, report.changes.size());
       Assertions.assertTrue(report.cost >= report.newCost);
       Assertions.assertEquals(MyCost.class.getSimpleName(), report.function);
+      // "before" should record size
+      report.changes.stream()
+          .flatMap(c -> c.before.stream())
+          .forEach(p -> Assertions.assertNotEquals(0, p.size));
+      // "after" should NOT record size
+      report.changes.stream()
+          .flatMap(c -> c.after.stream())
+          .forEach(p -> Assertions.assertNull(p.size));
     }
   }