Repository: skiptests/astraea
Pull request #604 (2022-08-20 19:33:06)

Avoid creating too many partition lists
<img width="1311" alt="截圖 2022-08-21 上午1 45 47" src="https://user-images.githubusercontent.com/6234750/185763400-097e9752-a4cb-450a-b466-0346e8af85f5.png">

主要原因是最後產生 <broker, topic> partitions 的時候會產生新的陣列，當 partitions 數量很多時會有明顯的成本，因此在建立ClusterInfo的時候我們就先分類好以避免後續的成本


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 7a910388..7b2e577c 100644
--- a/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
+++ b/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
@@ -16,7 +16,6 @@
  */
 package org.astraea.app.admin;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
@@ -68,21 +67,28 @@ public interface 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 =
+    var allReplicas =
         topics.stream()
-            .flatMap(t -> cluster.availablePartitionsForTopic(t).stream())
+            .flatMap(t -> cluster.partitionsForTopic(t).stream())
             .flatMap(p -> ReplicaInfo.of(p).stream())
             .collect(Collectors.toUnmodifiableList());
-    var availableReplicas = replicas.stream().collect(Collectors.groupingBy(ReplicaInfo::topic));
-    var availableReplicaLeaders =
-        availableReplicas.entrySet().stream()
-            .collect(
-                Collectors.toMap(
-                    Map.Entry::getKey,
-                    e ->
-                        e.getValue().stream()
-                            .filter(ReplicaInfo::isLeader)
-                            .collect(Collectors.toCollection(ArrayList::new))));
+    var replicasForTopic = allReplicas.stream().collect(Collectors.groupingBy(ReplicaInfo::topic));
+    var availableReplicasForTopic =
+        allReplicas.stream()
+            .filter(ReplicaInfo::isOnlineReplica)
+            .collect(Collectors.groupingBy(ReplicaInfo::topic));
+    var availableReplicaLeadersForTopics =
+        allReplicas.stream()
+            .filter(ReplicaInfo::isOnlineReplica)
+            .filter(ReplicaInfo::isLeader)
+            .collect(Collectors.groupingBy(ReplicaInfo::topic));
+    // This group is used commonly, so we cache it.
+    var availableLeaderReplicasForBrokersTopics =
+        allReplicas.stream()
+            .filter(ReplicaInfo::isOnlineReplica)
+            .filter(ReplicaInfo::isLeader)
+            .collect(Collectors.groupingBy(r -> Map.entry(r.nodeInfo().id(), r.topic())));
+
     return new ClusterInfo() {
       @Override
       public List<NodeInfo> nodes() {
@@ -101,37 +107,27 @@ public interface ClusterInfo {
 
       @Override
       public List<ReplicaInfo> availableReplicaLeaders(String topic) {
-        return availableReplicaLeaders.getOrDefault(topic, new ArrayList<>(0));
+        return availableReplicaLeadersForTopics.getOrDefault(topic, List.of());
+      }
+
+      @Override
+      public List<ReplicaInfo> availableReplicaLeaders(int broker, String topic) {
+        return availableLeaderReplicasForBrokersTopics.getOrDefault(
+            Map.entry(broker, topic), List.of());
       }
 
       @Override
       public List<ReplicaInfo> availableReplicas(String topic) {
-        return availableReplicas.getOrDefault(topic, List.of());
+        return availableReplicasForTopic.getOrDefault(topic, List.of());
       }
 
       @Override
       public List<ReplicaInfo> replicas(String topic) {
-        return replicas;
+        return replicasForTopic.getOrDefault(topic, List.of());
       }
     };
   }
 
-  /**
-   * find the node associated to specify node and port. Normally, the node + port should be unique
-   * in cluster.
-   *
-   * @param host hostname
-   * @param port client port
-   * @return the node information. It throws NoSuchElementException if specify node and port is not
-   *     associated to any node
-   */
-  default NodeInfo node(String host, int port) {
-    return nodes().stream()
-        .filter(n -> n.host().equals(host) && n.port() == port)
-        .findAny()
-        .orElseThrow(() -> new NoSuchElementException(host + ":" + port + " is nonexistent"));
-  }
-
   /**
    * find the node associated to node id.
    *
@@ -159,18 +155,30 @@ public interface ClusterInfo {
    * Get the list of replica leader information of each available partition for the given topic
    *
    * @param topic The Topic name
-   * @return A list of {@link ReplicaInfo}. It throws NoSuchElementException if the replica info of
-   *     the given topic is unknown to this ClusterInfo
+   * @return A list of {@link ReplicaInfo}.
    */
   List<ReplicaInfo> availableReplicaLeaders(String topic);
 
+  /**
+   * Get the list of replica leader information of each available partition for the given
+   * broker/topic
+   *
+   * @param broker the broker id
+   * @param topic The Topic name
+   * @return A list of {@link ReplicaInfo}.
+   */
+  default List<ReplicaInfo> availableReplicaLeaders(int broker, String topic) {
+    return availableReplicaLeaders(topic).stream()
+        .filter(r -> r.nodeInfo().id() == broker)
+        .collect(Collectors.toUnmodifiableList());
+  }
+
   /**
    * Get the list of replica information of each available partition/replica pair for the given
    * topic
    *
    * @param topic The topic name
-   * @return A list of {@link ReplicaInfo}. It throws NoSuchElementException if the replica info of
-   *     the given topic is unknown to this ClusterInfo
+   * @return A list of {@link ReplicaInfo}.
    */
   List<ReplicaInfo> availableReplicas(String topic);
 
@@ -185,8 +193,7 @@ public interface ClusterInfo {
    * Get the list of replica information of each partition/replica pair for the given topic
    *
    * @param topic The topic name
-   * @return A list of {@link ReplicaInfo}. It throws NoSuchElementException if the replica info of
-   *     the given topic is unknown to this ClusterInfo
+   * @return A list of {@link ReplicaInfo}.
    */
   List<ReplicaInfo> replicas(String topic);
 }
diff --git a/app/src/main/java/org/astraea/app/balancer/executor/RebalanceAdminImpl.java b/app/src/main/java/org/astraea/app/balancer/executor/RebalanceAdminImpl.java
index 742dd12f..1f9d32ac 100644
--- a/app/src/main/java/org/astraea/app/balancer/executor/RebalanceAdminImpl.java
+++ b/app/src/main/java/org/astraea/app/balancer/executor/RebalanceAdminImpl.java
@@ -121,6 +121,10 @@ class RebalanceAdminImpl implements RebalanceAdmin {
             expectedPlacement.stream()
                 .map(LogPlacement::broker)
                 .collect(Collectors.toUnmodifiableList()));
+
+    // wait until the whole cluster knows the replica list just changed
+    Utils.sleep(Duration.ofMillis(500));
+
     // do inter-data-directories migration
     var forCrossDirMigration =
         expectedPlacement.stream()
diff --git a/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java b/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
index db111ba3..aa66b66e 100644
--- a/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
@@ -85,6 +85,22 @@ public class StrictCostDispatcher implements Dispatcher {
     return beanCollector.register().host(host).port(port).fetcher(fetcher).build();
   }
 
+  void tryToUpdateFetcher(ClusterInfo clusterInfo) {
+    // add new receivers for new brokers
+    fetcher.ifPresent(
+        fetcher ->
+            clusterInfo.nodes().stream()
+                .filter(node -> !receivers.containsKey(node.id()))
+                .forEach(
+                    node ->
+                        jmxPortGetter
+                            .apply(node.id())
+                            .ifPresent(
+                                port ->
+                                    receivers.put(
+                                        node.id(), receiver(node.host(), port, fetcher)))));
+  }
+
   @Override
   public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
     var partitionLeaders = clusterInfo.availableReplicaLeaders(topic);
@@ -94,24 +110,7 @@ public class StrictCostDispatcher implements Dispatcher {
     // just return the only one available partition
     if (partitionLeaders.size() == 1) return partitionLeaders.get(0).partition();
 
-    // add new receivers for new brokers
-    receivers.putAll(
-        fetcher
-            .map(
-                fetcher ->
-                    clusterInfo.nodes().stream()
-                        .filter(nodeInfo -> !receivers.containsKey(nodeInfo.id()))
-                        .distinct()
-                        .filter(nodeInfo -> jmxPortGetter.apply(nodeInfo.id()).isPresent())
-                        .collect(
-                            Collectors.toMap(
-                                NodeInfo::id,
-                                nodeInfo ->
-                                    receiver(
-                                        nodeInfo.host(),
-                                        jmxPortGetter.apply(nodeInfo.id()).get(),
-                                        fetcher))))
-            .orElse(Map.of()));
+    tryToUpdateFetcher(clusterInfo);
 
     tryToUpdateRoundRobin(clusterInfo);
 
@@ -121,11 +120,7 @@ public class StrictCostDispatcher implements Dispatcher {
 
     // TODO: if the topic partitions are existent in fewer brokers, the target gets -1 in most cases
     var candidate =
-        target < 0
-            ? partitionLeaders
-            : partitionLeaders.stream()
-                .filter(r -> r.nodeInfo().id() == target)
-                .collect(Collectors.toUnmodifiableList());
+        target < 0 ? partitionLeaders : clusterInfo.availableReplicaLeaders(target, topic);
     candidate = candidate.isEmpty() ? partitionLeaders : candidate;
     return candidate.get((int) (Math.random() * candidate.size())).partition();
   }


Test patch:
diff --git a/app/src/test/java/org/astraea/app/admin/AdminTest.java b/app/src/test/java/org/astraea/app/admin/AdminTest.java
index 55af3992..f923076e 100644
--- a/app/src/test/java/org/astraea/app/admin/AdminTest.java
+++ b/app/src/test/java/org/astraea/app/admin/AdminTest.java
@@ -737,8 +737,6 @@ public class AdminTest extends RequireBrokerCluster {
           NoSuchElementException.class, () -> clusterInfo.availableReplicaLeaders("Unknown Topic"));
       Assertions.assertThrows(NoSuchElementException.class, () -> clusterInfo.dataDirectories(-1));
       Assertions.assertThrows(NoSuchElementException.class, () -> clusterInfo.node(-1));
-      Assertions.assertThrows(
-          NoSuchElementException.class, () -> clusterInfo.node("unknown", 1024));
     }
   }
 
@@ -1304,9 +1302,9 @@ public class AdminTest extends RequireBrokerCluster {
           .numberOfPartitions(partitions)
           .numberOfReplicas((short) 3)
           .create();
-      Utils.sleep(Duration.ofMillis(100));
+      Utils.sleep(Duration.ofSeconds(1));
       admin.replicationThrottler().throttle(topic).apply();
-      Utils.sleep(Duration.ofMillis(100));
+      Utils.sleep(Duration.ofSeconds(1));
 
       Assertions.assertEquals(currentLeaderLogs(admin, topic), fetchLeaderThrottle(topic));
       Assertions.assertEquals(currentFollowerLogs(admin, topic), fetchFollowerThrottle(topic));
diff --git a/app/src/test/java/org/astraea/app/balancer/executor/RebalanceAdminTest.java b/app/src/test/java/org/astraea/app/balancer/executor/RebalanceAdminTest.java
index 27cfea0f..07211113 100644
--- a/app/src/test/java/org/astraea/app/balancer/executor/RebalanceAdminTest.java
+++ b/app/src/test/java/org/astraea/app/balancer/executor/RebalanceAdminTest.java
@@ -47,7 +47,7 @@ import org.junit.jupiter.api.Test;
 
 class RebalanceAdminTest extends RequireBrokerCluster {
 
-  @Test
+  @RepeatedTest(value = 3)
   void alterReplicaPlacementByList() {
     // arrange
     try (var admin = Admin.of(bootstrapServers())) {
@@ -271,6 +271,7 @@ class RebalanceAdminTest extends RequireBrokerCluster {
       var topicPartition3 = TopicPartition.of(topic3, 0);
       Stream.of(topic1, topic2, topic3)
           .forEach(i -> admin.creator().topic(i).numberOfPartitions(1).create());
+      Utils.sleep(Duration.ofSeconds(1));
       var allowed = List.of(topic1, topic2);
       Predicate<String> filter = allowed::contains;
       var rebalanceAdmin = RebalanceAdmin.of(admin, filter);
diff --git a/app/src/test/java/org/astraea/app/cost/ClusterInfoTest.java b/app/src/test/java/org/astraea/app/cost/ClusterInfoTest.java
index 4e325291..6ca81a64 100644
--- a/app/src/test/java/org/astraea/app/cost/ClusterInfoTest.java
+++ b/app/src/test/java/org/astraea/app/cost/ClusterInfoTest.java
@@ -43,7 +43,6 @@ public class ClusterInfoTest {
 
     Assertions.assertEquals(1, clusterInfo.nodes().size());
     Assertions.assertEquals(NodeInfo.of(node), clusterInfo.nodes().get(0));
-    Assertions.assertEquals(clusterInfo.nodes().get(0), clusterInfo.node(node.host(), node.port()));
     Assertions.assertEquals(1, clusterInfo.availableReplicas(partition.topic()).size());
     Assertions.assertEquals(1, clusterInfo.replicas(partition.topic()).size());
     Assertions.assertEquals(
@@ -60,7 +59,6 @@ public class ClusterInfoTest {
     var clusterInfo = ClusterInfo.of(Cluster.empty());
     Assertions.assertEquals(0, clusterInfo.replicas("unknown").size());
     Assertions.assertThrows(NoSuchElementException.class, () -> clusterInfo.node(0));
-    Assertions.assertThrows(NoSuchElementException.class, () -> clusterInfo.node("", -1));
     Assertions.assertEquals(0, clusterInfo.availableReplicas("unknown").size());
     Assertions.assertEquals(0, clusterInfo.availableReplicaLeaders("unknown").size());
     Assertions.assertThrows(
diff --git a/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java b/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
index bede667f..a6887e39 100644
--- a/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
@@ -283,6 +283,8 @@ public class StrictCostDispatcherTest {
               Stream.of(replicaInfo0, replicaInfo1)
                   .map(ReplicaInfo::nodeInfo)
                   .collect(Collectors.toUnmodifiableList()));
+      Mockito.when(clusterInfo.availableReplicaLeaders(Mockito.anyInt(), Mockito.anyString()))
+          .thenReturn(List.of(replicaInfo0));
       Assertions.assertEquals(
           partitionId, dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo));
     }
@@ -328,4 +330,32 @@ public class StrictCostDispatcherTest {
       Assertions.assertNotEquals(t, dispatcher.timeToUpdateRoundRobin);
     }
   }
+
+  @Test
+  void testTryToUpdateFetcher() {
+    var receiverCount = new AtomicInteger(0);
+    try (var dispatcher =
+        new StrictCostDispatcher() {
+          @Override
+          Receiver receiver(String host, int port, Fetcher fetcher) {
+            receiverCount.incrementAndGet();
+            return Mockito.mock(Receiver.class);
+          }
+        }) {
+      var nodeInfo = NodeInfo.of(10, "host", 2222);
+      var clusterInfo = Mockito.mock(ClusterInfo.class);
+      Mockito.when(clusterInfo.nodes()).thenReturn(List.of(nodeInfo));
+
+      Assertions.assertEquals(0, dispatcher.receivers.size());
+      dispatcher.fetcher = Optional.of(Mockito.mock(Fetcher.class));
+      dispatcher.jmxPortGetter = id -> Optional.of(1111);
+      dispatcher.tryToUpdateFetcher(clusterInfo);
+      Assertions.assertEquals(1, dispatcher.receivers.size());
+      Assertions.assertEquals(1, receiverCount.get());
+
+      dispatcher.tryToUpdateFetcher(clusterInfo);
+      Assertions.assertEquals(1, dispatcher.receivers.size());
+      Assertions.assertEquals(1, receiverCount.get());
+    }
+  }
 }