Repository: skiptests/astraea
Pull request #299 (2022-04-13 17:16:30)

Revise `ClusterInfo` and `PartitionInfo` implementation
這個 PR 對 `ClusterInfo` 和 `PartitionInfo` 新增一些 Balancer 需要的欄位

* add `ClusterInfo#topics`
  有 `ClusterInfo#partitoins(String topic)` 和 `ClusterInfo#availablePartitions(String topic)`, 但是沒有辦法單從 `ClusterInfo` 得知有哪些 topic 存在，所以這些 method 沒辦法直接利用，如果 `ClusterInfo` 是 cluster 在某個時間點的 snapshot 的話 `topic` 資訊應該也要暴露。
* add `replicas`, `inSyncReplicas`, `offlineReplicas` fields to `PartitionInfo`
  `ClusterInfo` 如果是 cluster snapshot 的話，這些 replica 分佈的資訊應該也要包含在內，不過原先的 `PartitionInfo` 並沒有這些欄位，導致後面 CostFunction 沒辦法從 `ClusterInfo` 中得知 replica 分佈資訊。


Patch:
diff --git a/app/src/main/java/org/astraea/cost/ClusterInfo.java b/app/src/main/java/org/astraea/cost/ClusterInfo.java
index c6bccbe4..96fe70dd 100644
--- a/app/src/main/java/org/astraea/cost/ClusterInfo.java
+++ b/app/src/main/java/org/astraea/cost/ClusterInfo.java
@@ -7,6 +7,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Set;
 import java.util.stream.Collectors;
 import org.astraea.metrics.HasBeanObject;
 
@@ -19,6 +20,10 @@ public interface ClusterInfo {
         return cluster.nodes().stream().map(NodeInfo::of).collect(Collectors.toUnmodifiableList());
       }
 
+      public Set<String> topics() {
+        return cluster.topics();
+      }
+
       @Override
       public List<PartitionInfo> availablePartitions(String topic) {
         return cluster.availablePartitionsForTopic(topic).stream()
@@ -70,6 +75,11 @@ public interface ClusterInfo {
         return cluster.availablePartitions(topic);
       }
 
+      @Override
+      public Set<String> topics() {
+        return cluster.topics();
+      }
+
       @Override
       public List<PartitionInfo> partitions(String topic) {
         return cluster.partitions(topic);
@@ -128,6 +138,13 @@ public interface ClusterInfo {
    */
   List<PartitionInfo> availablePartitions(String topic);
 
+  /**
+   * All topic names
+   *
+   * @return return a set of topic names
+   */
+  Set<String> topics();
+
   /**
    * Get the list of partitions for this topic
    *
diff --git a/app/src/main/java/org/astraea/cost/PartitionInfo.java b/app/src/main/java/org/astraea/cost/PartitionInfo.java
index 50d3ffca..87c19003 100644
--- a/app/src/main/java/org/astraea/cost/PartitionInfo.java
+++ b/app/src/main/java/org/astraea/cost/PartitionInfo.java
@@ -1,12 +1,32 @@
 package org.astraea.cost;
 
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
 public interface PartitionInfo {
 
   static PartitionInfo of(org.apache.kafka.common.PartitionInfo pf) {
-    return of(pf.topic(), pf.partition(), NodeInfo.of(pf.leader()));
+    return of(
+        pf.topic(),
+        pf.partition(),
+        NodeInfo.of(pf.leader()),
+        Arrays.stream(pf.replicas()).map(NodeInfo::of).collect(Collectors.toUnmodifiableList()),
+        Arrays.stream(pf.inSyncReplicas())
+            .map(NodeInfo::of)
+            .collect(Collectors.toUnmodifiableList()),
+        Arrays.stream(pf.offlineReplicas())
+            .map(NodeInfo::of)
+            .collect(Collectors.toUnmodifiableList()));
   }
 
-  static PartitionInfo of(String topic, int partition, NodeInfo leader) {
+  static PartitionInfo of(
+      String topic,
+      int partition,
+      NodeInfo leader,
+      List<NodeInfo> replicas,
+      List<NodeInfo> inSyncReplicas,
+      List<NodeInfo> offlineReplicas) {
     return new PartitionInfo() {
       @Override
       public String topic() {
@@ -22,6 +42,21 @@ public interface PartitionInfo {
       public NodeInfo leader() {
         return leader;
       }
+
+      @Override
+      public List<NodeInfo> replicas() {
+        return List.copyOf(replicas);
+      }
+
+      @Override
+      public List<NodeInfo> inSyncReplica() {
+        return List.copyOf(inSyncReplicas);
+      }
+
+      @Override
+      public List<NodeInfo> offlineReplicas() {
+        return List.copyOf(offlineReplicas);
+      }
     };
   }
 
@@ -33,4 +68,13 @@ public interface PartitionInfo {
 
   /** @return information of leader node */
   NodeInfo leader();
+
+  /** @return a list of replicas */
+  List<NodeInfo> replicas();
+
+  /** @return a list of in-sync replicas */
+  List<NodeInfo> inSyncReplica();
+
+  /** @return a list of offline replicas */
+  List<NodeInfo> offlineReplicas();
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/cost/FakeClusterInfo.java b/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
index 7f6d13df..e2707feb 100644
--- a/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
+++ b/app/src/test/java/org/astraea/cost/FakeClusterInfo.java
@@ -3,6 +3,7 @@ package org.astraea.cost;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.astraea.metrics.HasBeanObject;
 
 public class FakeClusterInfo implements ClusterInfo {
@@ -16,6 +17,11 @@ public class FakeClusterInfo implements ClusterInfo {
     return List.of();
   }
 
+  @Override
+  public Set<String> topics() {
+    return Set.of();
+  }
+
   @Override
   public List<PartitionInfo> partitions(String topic) {
     return List.of();
diff --git a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
index f471e71f..b14a6d57 100644
--- a/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/partitioner/StrictCostDispatcherTest.java
@@ -17,6 +17,41 @@ import org.mockito.Mockito;
 
 public class StrictCostDispatcherTest {
 
+  private static PartitionInfo createFakePartitionInfo(
+      String topic, int partition, NodeInfo leaderNode) {
+    return new PartitionInfo() {
+      @Override
+      public String topic() {
+        return topic;
+      }
+
+      @Override
+      public int partition() {
+        return partition;
+      }
+
+      @Override
+      public NodeInfo leader() {
+        return leaderNode;
+      }
+
+      @Override
+      public List<NodeInfo> replicas() {
+        return null;
+      }
+
+      @Override
+      public List<NodeInfo> inSyncReplica() {
+        return null;
+      }
+
+      @Override
+      public List<NodeInfo> offlineReplicas() {
+        return null;
+      }
+    };
+  }
+
   @Test
   void testClose() {
     var dispatcher = new StrictCostDispatcher();
@@ -41,8 +76,8 @@ public class StrictCostDispatcherTest {
   void testBestPartition() {
     var partitions =
         List.of(
-            PartitionInfo.of("t", 0, NodeInfo.of(10, "h0", 1000)),
-            PartitionInfo.of("t", 1, NodeInfo.of(11, "h1", 1000)));
+            createFakePartitionInfo("t", 0, NodeInfo.of(10, "h0", 1000)),
+            createFakePartitionInfo("t", 1, NodeInfo.of(11, "h1", 1000)));
 
     var score = List.of(Map.of(10, 0.8D), Map.of(11, 0.7D));
 
@@ -56,10 +91,10 @@ public class StrictCostDispatcherTest {
   void testPartition() {
     // n0 is busy
     var n0 = NodeInfo.of(10, "host0", 12345);
-    var p0 = PartitionInfo.of("aa", 1, n0);
+    var p0 = createFakePartitionInfo("aa", 1, n0);
     // n1 is free
     var n1 = NodeInfo.of(11, "host1", 12345);
-    var p1 = PartitionInfo.of("aa", 2, n1);
+    var p1 = createFakePartitionInfo("aa", 2, n1);
 
     var receiver = Mockito.mock(Receiver.class);
     Mockito.when(receiver.current()).thenReturn(List.of());