Repository: skiptests/astraea
Pull request #471 (2022-07-09 18:41:41)

Cache the cluster info to reduce the cost of converting kafka cluster…
兩個重點：
1. Producer並不會一直更新`Cluster`，因此我們可以透過暫存來避免過多的轉換吃掉頻寬
2. 讓RR的更新頻率可以被參數化
3. 將RR的更新頻率預設值調整成與更新beans的時間一致

chia7712: @chinghongfang 這隻PR合併後，你可否麻煩在實驗室環境中做測試？我這邊有做一個簡單的測試，在叢集中塞入一台狀況不好的機器，`StrictCostDispatcher`的表現比預設的partitioner還好，但希望你能用實驗室的設備再幫忙驗證一次

BTW，預設的`StrictCostDispatcher`並不需要任何額外的參數就可以運作的很好，因為他現在仰賴的是`Producer metrics`

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 37830767..764c1f72 100644
--- a/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
+++ b/app/src/main/java/org/astraea/app/admin/ClusterInfo.java
@@ -27,11 +27,35 @@ import org.astraea.app.metrics.HasBeanObject;
 
 public interface ClusterInfo {
 
+  /**
+   * convert the kafka Cluster to our ClusterInfo. All data structure are converted immediately, so
+   * you should cache the result if the performance is critical
+   *
+   * @param cluster kafka ClusterInfo
+   * @return astraea 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()
+            .flatMap(t -> cluster.availablePartitionsForTopic(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.toUnmodifiableList())));
     return new ClusterInfo() {
       @Override
       public List<NodeInfo> nodes() {
-        return cluster.nodes().stream().map(NodeInfo::of).collect(Collectors.toUnmodifiableList());
+        return nodes;
       }
 
       @Override
@@ -41,48 +65,22 @@ public interface ClusterInfo {
       }
 
       public Set<String> topics() {
-        return cluster.topics();
+        return topics;
       }
 
       @Override
       public List<ReplicaInfo> availableReplicaLeaders(String topic) {
-        var info =
-            cluster.availablePartitionsForTopic(topic).stream()
-                .map(ReplicaInfo::of)
-                .map(
-                    replicas ->
-                        replicas.stream().filter(ReplicaInfo::isLeader).findFirst().orElseThrow())
-                .collect(Collectors.toUnmodifiableList());
-        if (info.isEmpty())
-          throw new NoSuchElementException(
-              "This ClusterInfo have no information about topic \"" + topic + "\"");
-        return info;
+        return availableReplicaLeaders.getOrDefault(topic, List.of());
       }
 
       @Override
       public List<ReplicaInfo> availableReplicas(String topic) {
-        var info =
-            cluster.availablePartitionsForTopic(topic).stream()
-                .map(ReplicaInfo::of)
-                .flatMap(Collection::stream)
-                .collect(Collectors.toUnmodifiableList());
-        if (info.isEmpty())
-          throw new NoSuchElementException(
-              "This ClusterInfo have no information about topic \"" + topic + "\"");
-        return info;
+        return availableReplicas.getOrDefault(topic, List.of());
       }
 
       @Override
       public List<ReplicaInfo> replicas(String topic) {
-        var info =
-            cluster.partitionsForTopic(topic).stream()
-                .map(ReplicaInfo::of)
-                .flatMap(Collection::stream)
-                .collect(Collectors.toUnmodifiableList());
-        if (info.isEmpty())
-          throw new NoSuchElementException(
-              "This ClusterInfo have no information about topic \"" + topic + "\"");
-        return info;
+        return replicas;
       }
 
       @Override
diff --git a/app/src/main/java/org/astraea/app/partitioner/Dispatcher.java b/app/src/main/java/org/astraea/app/partitioner/Dispatcher.java
index a28fcf9a..85310bab 100644
--- a/app/src/main/java/org/astraea/app/partitioner/Dispatcher.java
+++ b/app/src/main/java/org/astraea/app/partitioner/Dispatcher.java
@@ -17,12 +17,19 @@
 package org.astraea.app.partitioner;
 
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 import org.apache.kafka.clients.producer.Partitioner;
 import org.apache.kafka.common.Cluster;
 import org.astraea.app.admin.ClusterInfo;
 
 public interface Dispatcher extends Partitioner {
+  /**
+   * cache the cluster info to reduce the cost of converting cluster. Producer does not update
+   * Cluster frequently, so it is ok to cache it.
+   */
+  ConcurrentHashMap<Cluster, ClusterInfo> CLUSTER_CACHE = new ConcurrentHashMap<>();
+
   /**
    * Compute the partition for the given record.
    *
@@ -59,7 +66,7 @@ public interface Dispatcher extends Partitioner {
         topic,
         keyBytes == null ? new byte[0] : keyBytes,
         valueBytes == null ? new byte[0] : valueBytes,
-        ClusterInfo.of(cluster));
+        CLUSTER_CACHE.computeIfAbsent(cluster, ignored -> ClusterInfo.of(cluster)));
   }
 
   @Override
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 f9fb5530..79fbdd1b 100644
--- a/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
+++ b/app/src/main/java/org/astraea/app/partitioner/StrictCostDispatcher.java
@@ -27,6 +27,7 @@ import java.util.stream.Collectors;
 import org.astraea.app.admin.ClusterInfo;
 import org.astraea.app.admin.NodeInfo;
 import org.astraea.app.admin.ReplicaInfo;
+import org.astraea.app.argument.DurationField;
 import org.astraea.app.common.Utils;
 import org.astraea.app.cost.CostFunction;
 import org.astraea.app.cost.HasBrokerCost;
@@ -52,11 +53,13 @@ import org.astraea.app.metrics.collector.Receiver;
  */
 public class StrictCostDispatcher implements Dispatcher {
   public static final String JMX_PORT = "jmx.port";
-  private static final Duration ROUND_ROBIN_LEASE = Duration.ofSeconds(30);
+  public static final String ROUND_ROBIN_LEASE_KEY = "round.robin.lease";
 
   private final BeanCollector beanCollector =
       BeanCollector.builder().interval(Duration.ofSeconds(4)).build();
 
+  Duration roundRobinLease;
+
   // The cost-functions we consider and the weight of them. It is visible for test
   Map<CostFunction, Double> functions = Map.of();
 
@@ -104,15 +107,8 @@ public class StrictCostDispatcher implements Dispatcher {
                                         fetcher))))
             .orElse(Map.of()));
 
-    // get latest beans for each node
-    var beans =
-        receivers.entrySet().stream()
-            .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().current()));
+    tryToUpdateRoundRobin(clusterInfo);
 
-    if (roundRobin == null || System.currentTimeMillis() >= timeToUpdateRoundRobin) {
-      roundRobin = newRoundRobin(functions, ClusterInfo.of(clusterInfo, beans));
-      timeToUpdateRoundRobin = System.currentTimeMillis() + ROUND_ROBIN_LEASE.toMillis();
-    }
     return roundRobin
         .next(partitionLeaders.stream().map(r -> r.nodeInfo().id()).collect(Collectors.toSet()))
         .flatMap(
@@ -125,6 +121,19 @@ public class StrictCostDispatcher implements Dispatcher {
         .orElse(0);
   }
 
+  void tryToUpdateRoundRobin(ClusterInfo clusterInfo) {
+    if (roundRobin == null || System.currentTimeMillis() >= timeToUpdateRoundRobin) {
+      roundRobin =
+          newRoundRobin(
+              functions,
+              ClusterInfo.of(
+                  clusterInfo,
+                  receivers.entrySet().stream()
+                      .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().current()))));
+      timeToUpdateRoundRobin = System.currentTimeMillis() + roundRobinLease.toMillis();
+    }
+  }
+
   /**
    * The value of cost returned from cost function is conflict to score, since the higher cost
    * represents lower score. This helper reverses the cost by subtracting the cost from "max cost".
@@ -175,7 +184,13 @@ public class StrictCostDispatcher implements Dispatcher {
     configure(
         configuredFunctions.isEmpty() ? Map.of(new NodeLatencyCost(), 1D) : configuredFunctions,
         config.integer(JMX_PORT),
-        PartitionerUtils.parseIdJMXPort(config));
+        PartitionerUtils.parseIdJMXPort(config),
+        config
+            .string(ROUND_ROBIN_LEASE_KEY)
+            .map(DurationField::toDuration)
+            // The duration of updating beans is 4 seconds, so
+            // the default duration of updating RR is 4 seconds.
+            .orElse(Duration.ofSeconds(4)));
   }
 
   /**
@@ -188,7 +203,8 @@ public class StrictCostDispatcher implements Dispatcher {
   void configure(
       Map<CostFunction, Double> functions,
       Optional<Integer> jmxPortDefault,
-      Map<Integer, Integer> customJmxPort) {
+      Map<Integer, Integer> customJmxPort,
+      Duration roundRobinLease) {
     this.functions = functions;
     this.fetcher = Fetcher.of(this.functions.keySet());
     this.jmxPortGetter = id -> Optional.ofNullable(customJmxPort.get(id)).or(() -> jmxPortDefault);
@@ -196,6 +212,7 @@ public class StrictCostDispatcher implements Dispatcher {
     // put local mbean client first
     this.fetcher.ifPresent(
         f -> receivers.put(-1, beanCollector.register().local().fetcher(f).build()));
+    this.roundRobinLease = roundRobinLease;
   }
 
   /**


Test patch:
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 dfa54b1f..df94f547 100644
--- a/app/src/test/java/org/astraea/app/cost/ClusterInfoTest.java
+++ b/app/src/test/java/org/astraea/app/cost/ClusterInfoTest.java
@@ -19,6 +19,7 @@ package org.astraea.app.cost;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Set;
 import org.apache.kafka.common.Cluster;
 import org.astraea.app.admin.ClusterBean;
 import org.astraea.app.admin.ClusterInfo;
@@ -35,6 +36,7 @@ public class ClusterInfoTest {
     var node = NodeInfoTest.node();
     var partition = ReplicaInfoTest.partitionInfo();
     var kafkaCluster = Mockito.mock(Cluster.class);
+    Mockito.when(kafkaCluster.topics()).thenReturn(Set.of(partition.topic()));
     Mockito.when(kafkaCluster.availablePartitionsForTopic(partition.topic()))
         .thenReturn(List.of(partition));
     Mockito.when(kafkaCluster.partitionsForTopic(partition.topic())).thenReturn(List.of(partition));
@@ -59,14 +61,11 @@ public class ClusterInfoTest {
   @Test
   void testIllegalQuery() {
     var clusterInfo = ClusterInfo.of(Cluster.empty());
-
-    Assertions.assertThrows(NoSuchElementException.class, () -> clusterInfo.replicas("unknown"));
+    Assertions.assertEquals(0, clusterInfo.replicas("unknown").size());
     Assertions.assertThrows(NoSuchElementException.class, () -> clusterInfo.node(0));
     Assertions.assertThrows(NoSuchElementException.class, () -> clusterInfo.node("", -1));
-    Assertions.assertThrows(
-        NoSuchElementException.class, () -> clusterInfo.availableReplicas("unknown"));
-    Assertions.assertThrows(
-        NoSuchElementException.class, () -> clusterInfo.availableReplicaLeaders("unknown"));
+    Assertions.assertEquals(0, clusterInfo.availableReplicas("unknown").size());
+    Assertions.assertEquals(0, clusterInfo.availableReplicaLeaders("unknown").size());
     Assertions.assertThrows(
         UnsupportedOperationException.class, () -> clusterInfo.dataDirectories(0));
   }
diff --git a/app/src/test/java/org/astraea/app/partitioner/DispatcherTest.java b/app/src/test/java/org/astraea/app/partitioner/DispatcherTest.java
index aa66ca34..bf5629d5 100644
--- a/app/src/test/java/org/astraea/app/partitioner/DispatcherTest.java
+++ b/app/src/test/java/org/astraea/app/partitioner/DispatcherTest.java
@@ -16,9 +16,12 @@
  */
 package org.astraea.app.partitioner;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.kafka.common.Cluster;
 import org.astraea.app.admin.ClusterInfo;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -44,10 +47,27 @@ public class DispatcherTest {
           }
         };
     Assertions.assertEquals(0, count.get());
-    // it should not throw NPE
-    dispatcher.partition("t", null, null, null, null, null);
+    dispatcher.partition(
+        "t", null, null, null, null, new Cluster("aa", List.of(), List.of(), Set.of(), Set.of()));
     Assertions.assertEquals(1, count.get());
     dispatcher.configure(Map.of("a", "b"));
     Assertions.assertEquals(2, count.get());
   }
+
+  @Test
+  void testClusterCache() {
+    var dispatcher =
+        new Dispatcher() {
+          @Override
+          public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
+            return 0;
+          }
+        };
+    var initialCount = Dispatcher.CLUSTER_CACHE.size();
+    var cluster = new Cluster("aa", List.of(), List.of(), Set.of(), Set.of());
+    dispatcher.partition("topic", "a", new byte[0], "v", new byte[0], cluster);
+    Assertions.assertEquals(initialCount + 1, Dispatcher.CLUSTER_CACHE.size());
+    dispatcher.partition("topic", "a", new byte[0], "v", new byte[0], cluster);
+    Assertions.assertEquals(initialCount + 1, Dispatcher.CLUSTER_CACHE.size());
+  }
 }
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 4a238f2a..23bb8b69 100644
--- a/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
+++ b/app/src/test/java/org/astraea/app/partitioner/StrictCostDispatcherTest.java
@@ -16,6 +16,7 @@
  */
 package org.astraea.app.partitioner;
 
+import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -24,6 +25,7 @@ import org.astraea.app.admin.ClusterBean;
 import org.astraea.app.admin.ClusterInfo;
 import org.astraea.app.admin.NodeInfo;
 import org.astraea.app.admin.ReplicaInfo;
+import org.astraea.app.common.Utils;
 import org.astraea.app.cost.BrokerCost;
 import org.astraea.app.cost.BrokerInputCost;
 import org.astraea.app.cost.CostFunction;
@@ -102,7 +104,7 @@ public class StrictCostDispatcherTest {
     var clusterInfo = Mockito.mock(ClusterInfo.class);
     Mockito.when(clusterInfo.availableReplicaLeaders(Mockito.anyString())).thenReturn(List.of());
     try (var dispatcher = new StrictCostDispatcher()) {
-      dispatcher.configure(Map.of(), Optional.empty(), Map.of());
+      dispatcher.configure(Map.of(), Optional.empty(), Map.of(), Duration.ofSeconds(10));
       Assertions.assertEquals(
           0, dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo));
     }
@@ -116,7 +118,7 @@ public class StrictCostDispatcherTest {
     Mockito.when(clusterInfo.availableReplicaLeaders(Mockito.anyString()))
         .thenReturn(List.of(replicaInfo));
     try (var dispatcher = new StrictCostDispatcher()) {
-      dispatcher.configure(Map.of(), Optional.empty(), Map.of());
+      dispatcher.configure(Map.of(), Optional.empty(), Map.of(), Duration.ofSeconds(10));
       Assertions.assertEquals(
           10, dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo));
     }
@@ -171,7 +173,8 @@ public class StrictCostDispatcherTest {
         .thenReturn(List.of(replicaInfo0, replicaInfo1));
     Mockito.when(clusterInfo.clusterBean()).thenReturn(ClusterBean.of(Map.of()));
     try (var dispatcher = new StrictCostDispatcher()) {
-      dispatcher.configure(Map.of(costFunction, 1D), Optional.empty(), Map.of());
+      dispatcher.configure(
+          Map.of(costFunction, 1D), Optional.empty(), Map.of(), Duration.ofSeconds(10));
       dispatcher.partition("topic", new byte[0], new byte[0], clusterInfo);
       Assertions.assertNotNull(dispatcher.roundRobin);
       Assertions.assertEquals(0, dispatcher.receivers.size());
@@ -198,7 +201,8 @@ public class StrictCostDispatcherTest {
             return Mockito.mock(Receiver.class);
           }
         };
-    dispatcher.configure(Map.of(costFunction, 1D), Optional.of(jmxPort), Map.of());
+    dispatcher.configure(
+        Map.of(costFunction, 1D), Optional.of(jmxPort), Map.of(), Duration.ofSeconds(10));
     var replicaInfo0 = ReplicaInfo.of("topic", 0, NodeInfo.of(10, "host", 11111), true, true, true);
     var replicaInfo1 = ReplicaInfo.of("topic", 1, NodeInfo.of(10, "host", 11111), true, true, true);
     var replicaInfo2 =
@@ -227,13 +231,19 @@ public class StrictCostDispatcherTest {
     var dispatcher = new StrictCostDispatcher();
 
     // pass due to local mbean
-    dispatcher.configure(Map.of(new ThroughputCost(), 1D), Optional.empty(), Map.of());
+    dispatcher.configure(
+        Map.of(new ThroughputCost(), 1D), Optional.empty(), Map.of(), Duration.ofSeconds(10));
 
     // pass due to default port
-    dispatcher.configure(Map.of(new ThroughputCost(), 1D), Optional.of(111), Map.of());
+    dispatcher.configure(
+        Map.of(new ThroughputCost(), 1D), Optional.of(111), Map.of(), Duration.ofSeconds(10));
 
     // pass due to specify port
-    dispatcher.configure(Map.of(new ThroughputCost(), 1D), Optional.empty(), Map.of(222, 111));
+    dispatcher.configure(
+        Map.of(new ThroughputCost(), 1D),
+        Optional.empty(),
+        Map.of(222, 111),
+        Duration.ofSeconds(10));
   }
 
   @Test
@@ -248,7 +258,8 @@ public class StrictCostDispatcherTest {
             return () -> Map.of(brokerId, 10D);
           }
         };
-    dispatcher.configure(Map.of(costFunction, 1D), Optional.empty(), Map.of());
+    dispatcher.configure(
+        Map.of(costFunction, 1D), Optional.empty(), Map.of(), Duration.ofSeconds(10));
 
     var replicaInfo0 =
         ReplicaInfo.of(
@@ -278,4 +289,26 @@ public class StrictCostDispatcherTest {
     var score = StrictCostDispatcher.costToScore(cost);
     Assertions.assertTrue(score.get(2) > score.get(1));
   }
+
+  @Test
+  void testRoundRobinLease() {
+    var dispatcher = new StrictCostDispatcher();
+    dispatcher.configure(
+        Configuration.of(Map.of(StrictCostDispatcher.ROUND_ROBIN_LEASE_KEY, "2s")));
+    Assertions.assertEquals(Duration.ofSeconds(2), dispatcher.roundRobinLease);
+
+    var clusterInfo = Mockito.mock(ClusterInfo.class);
+    Mockito.when(clusterInfo.clusterBean()).thenReturn(ClusterBean.of(Map.of()));
+    dispatcher.tryToUpdateRoundRobin(clusterInfo);
+    var rr = dispatcher.roundRobin;
+    Assertions.assertNotNull(rr);
+    // the rr is not updated yet
+    dispatcher.tryToUpdateRoundRobin(clusterInfo);
+    Assertions.assertEquals(rr, dispatcher.roundRobin);
+
+    Utils.sleep(Duration.ofSeconds(3));
+    dispatcher.tryToUpdateRoundRobin(clusterInfo);
+    // rr is updated already
+    Assertions.assertNotEquals(rr, dispatcher.roundRobin);
+  }
 }