Repository: skiptests/astraea
Pull request #252 (2022-02-24 08:09:20)

Flexible Round Robin partitoner v2
微調#251

1. 新增`Dispatcher`，該介面實作`Partitioner`並提供更適合astraea的方法
2. 封裝大部分的kafka cluster物件
3. 新增`Configuration`方便索取參數


Patch:
diff --git a/app/build.gradle b/app/build.gradle
index b37af5f3..19a3ba0c 100644
--- a/app/build.gradle
+++ b/app/build.gradle
@@ -10,12 +10,12 @@ repositories {
 }
 
 def versions = [
-    kafka: project.properties['kafka.version'] ?: "2.8.1",
+    kafka: project.properties['kafka.version'] ?: "3.1.0",
     junit: project.properties['junit.version'] ?: "5.8.2",
     mockito: project.properties['mockito.version'] ?: "4.2.0",
     jcommander: project.properties['jcommander.version'] ?: "1.81",
     slf4j: project.properties['slf4j.version'] ?: "1.7.32",
-    zookeeper: project.properties['zookeeper.version'] ?: "3.5.9",
+    zookeeper: project.properties['zookeeper.version'] ?: "3.6.3",
 ]
 
 dependencies {
diff --git a/app/src/main/java/org/astraea/partitioner/ClusterInfo.java b/app/src/main/java/org/astraea/partitioner/ClusterInfo.java
new file mode 100644
index 00000000..8b6767e5
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/ClusterInfo.java
@@ -0,0 +1,49 @@
+package org.astraea.partitioner;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+public interface ClusterInfo {
+
+  static ClusterInfo of(org.apache.kafka.common.Cluster cluster) {
+    return new ClusterInfo() {
+      @Override
+      public List<NodeInfo> nodes() {
+        return cluster.nodes().stream().map(NodeInfo::of).collect(Collectors.toUnmodifiableList());
+      }
+
+      @Override
+      public List<PartitionInfo> availablePartitions(String topic) {
+        return cluster.availablePartitionsForTopic(topic).stream()
+            .map(PartitionInfo::of)
+            .collect(Collectors.toUnmodifiableList());
+      }
+
+      @Override
+      public List<PartitionInfo> partitions(String topic) {
+        return cluster.partitionsForTopic(topic).stream()
+            .map(PartitionInfo::of)
+            .collect(Collectors.toUnmodifiableList());
+      }
+    };
+  }
+
+  /** @return The known set of nodes */
+  List<NodeInfo> nodes();
+
+  /**
+   * Get the list of available partitions for this topic
+   *
+   * @param topic The topic name
+   * @return A list of partitions
+   */
+  List<PartitionInfo> availablePartitions(String topic);
+
+  /**
+   * Get the list of partitions for this topic
+   *
+   * @param topic The topic name
+   * @return A list of partitions
+   */
+  List<PartitionInfo> partitions(String topic);
+}
diff --git a/app/src/main/java/org/astraea/partitioner/Configuration.java b/app/src/main/java/org/astraea/partitioner/Configuration.java
new file mode 100644
index 00000000..74000010
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/Configuration.java
@@ -0,0 +1,77 @@
+package org.astraea.partitioner;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+public interface Configuration {
+
+  static Configuration of(Map<String, ?> configs) {
+    return new Configuration() {
+      @Override
+      public String string(String key) {
+        return Objects.requireNonNull(configs.get(key), key + " is nonexistent").toString();
+      }
+
+      @Override
+      public List<String> list(String key, String separator) {
+        return Arrays.asList(string(key).split(separator));
+      }
+
+      @Override
+      public <T> Map<String, T> map(
+          String key,
+          String listSeparator,
+          String mapSeparator,
+          Function<String, T> valueConverter) {
+        Function<String, Map.Entry<String, T>> split =
+            s -> {
+              var items = s.split(mapSeparator);
+              if (items.length != 2)
+                throw new IllegalArgumentException(
+                    "the value: " + s + " is using incorrect separator");
+              return Map.entry(items[0], valueConverter.apply(items[1]));
+            };
+        return list(key, listSeparator).stream()
+            .map(split)
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+      }
+    };
+  }
+
+  /**
+   * @param key the key whose associated value is to be returned
+   * @return string value. never null
+   */
+  String string(String key);
+
+  /**
+   * @param key the key whose associated value is to be returned
+   * @param separator to split string to multiple strings
+   * @return string list. never null
+   */
+  List<String> list(String key, String separator);
+
+  /**
+   * @param key the key whose associated value is to be returned
+   * @param listSeparator to split string to multiple strings
+   * @param mapSeparator to split multiple strings to map
+   * @return string map. never null
+   */
+  default Map<String, String> map(String key, String listSeparator, String mapSeparator) {
+    return map(key, listSeparator, mapSeparator, s -> s);
+  }
+
+  /**
+   * @param key the key whose associated value is to be returned
+   * @param listSeparator to split string to multiple strings
+   * @param mapSeparator to split multiple strings to map
+   * @param valueConverter used to convert string to specify type
+   * @return string map. never null
+   */
+  <T> Map<String, T> map(
+      String key, String listSeparator, String mapSeparator, Function<String, T> valueConverter);
+}
diff --git a/app/src/main/java/org/astraea/partitioner/Dispatcher.java b/app/src/main/java/org/astraea/partitioner/Dispatcher.java
new file mode 100644
index 00000000..cfdddac8
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/Dispatcher.java
@@ -0,0 +1,37 @@
+package org.astraea.partitioner;
+
+import java.util.Map;
+import org.apache.kafka.clients.producer.Partitioner;
+import org.apache.kafka.common.Cluster;
+
+public interface Dispatcher extends Partitioner {
+  /**
+   * Compute the partition for the given record.
+   *
+   * @param topic The topic name
+   * @param key The key to partition on
+   * @param value The value to partition
+   * @param clusterInfo The current cluster metadata
+   */
+  int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo);
+
+  void configure(Configuration config);
+
+  @Override
+  default void configure(Map<String, ?> configs) {
+    configure(Configuration.of(configs));
+  }
+
+  @Override
+  default int partition(
+      String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
+    return partition(
+        topic,
+        keyBytes == null ? new byte[0] : keyBytes,
+        valueBytes == null ? new byte[0] : valueBytes,
+        ClusterInfo.of(cluster));
+  }
+
+  @Override
+  default void onNewBatch(String topic, Cluster cluster, int prevPartition) {}
+}
diff --git a/app/src/main/java/org/astraea/partitioner/FlexibleRoundRobinDispatcher.java b/app/src/main/java/org/astraea/partitioner/FlexibleRoundRobinDispatcher.java
new file mode 100644
index 00000000..8ffb1595
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/FlexibleRoundRobinDispatcher.java
@@ -0,0 +1,79 @@
+package org.astraea.partitioner;
+
+import static org.astraea.Utils.overSecond;
+import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.allPoisson;
+import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.toPositive;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.astraea.partitioner.nodeLoadMetric.NodeLoadClient;
+
+public class FlexibleRoundRobinDispatcher implements Dispatcher {
+  public static final String JMX_SERVERS_KEY = "jmx_servers";
+  private static final double SLOP = 0.2;
+  private NodeLoadClient nodeLoadClient;
+  private final ConcurrentMap<String, AtomicInteger> topicCounterMap = new ConcurrentHashMap<>();
+  private long lastTime;
+  private Map<Integer, Double> allPoisson;
+  private double avgPoisson;
+
+  @Override
+  public int partition(String topic, byte[] key, byte[] value, ClusterInfo clusterInfo) {
+    var loadCount =
+        Objects.requireNonNull(
+            nodeLoadClient.loadSituation(clusterInfo), "OverLoadCount should not be null.");
+    updatePoissonIfNeed(loadCount);
+
+    var partitions = clusterInfo.partitions(topic);
+    var numPartitions = partitions.size();
+    var availablePartitions = clusterInfo.availablePartitions(topic);
+    if (!availablePartitions.isEmpty()) {
+      return choosePartition(availablePartitions, topic);
+    } else {
+      var nextValue = nextValue(topic);
+      // no partitions are available, give a non-available partition
+      return toPositive(nextValue) % numPartitions;
+    }
+  }
+
+  @Override
+  public void configure(Configuration config) {
+    nodeLoadClient = new NodeLoadClient(config.map(JMX_SERVERS_KEY, ",", ":", Integer::valueOf));
+  }
+
+  @Override
+  public void close() {}
+
+  private int nextValue(String topic) {
+    var counter = topicCounterMap.computeIfAbsent(topic, k -> new AtomicInteger(0));
+    return counter.getAndIncrement();
+  }
+
+  private void updatePoissonIfNeed(Map<Integer, Integer> loadCount) {
+    if (overSecond(lastTime, 1)) {
+      lastTime = System.currentTimeMillis();
+      allPoisson = allPoisson(loadCount);
+      avgPoisson = allPoisson.values().stream().reduce(0.0, Double::sum) / allPoisson.size();
+    }
+  }
+
+  private int choosePartition(List<PartitionInfo> availablePartitions, String topic) {
+    var nextValue = nextValue(topic);
+    var part = toPositive(nextValue) % availablePartitions.size();
+    var partitionInfo = availablePartitions.get(part);
+    if (chooseOrNot(allPoisson.get(partitionInfo.leader().id()))) {
+      return partitionInfo.partition();
+    } else {
+      // no partitions are available, give a non-available partition
+      return toPositive(nextValue) % availablePartitions.size();
+    }
+  }
+
+  private boolean chooseOrNot(double brokerPoisson) {
+    return !(brokerPoisson >= (1 + SLOP) * avgPoisson);
+  }
+}
diff --git a/app/src/main/java/org/astraea/partitioner/NodeInfo.java b/app/src/main/java/org/astraea/partitioner/NodeInfo.java
new file mode 100644
index 00000000..17cb880e
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/NodeInfo.java
@@ -0,0 +1,36 @@
+package org.astraea.partitioner;
+
+public interface NodeInfo {
+
+  static NodeInfo of(org.apache.kafka.common.Node node) {
+    return of(node.id(), node.host(), node.port());
+  }
+
+  static NodeInfo of(int id, String host, int port) {
+    return new NodeInfo() {
+      @Override
+      public String host() {
+        return host;
+      }
+
+      @Override
+      public int id() {
+        return id;
+      }
+
+      @Override
+      public int port() {
+        return port;
+      }
+    };
+  }
+
+  /** @return The host name for this node */
+  String host();
+
+  /** @return The id for this node */
+  int id();
+
+  /** @return The client port for this node */
+  int port();
+}
diff --git a/app/src/main/java/org/astraea/partitioner/PartitionInfo.java b/app/src/main/java/org/astraea/partitioner/PartitionInfo.java
new file mode 100644
index 00000000..7e1b997d
--- /dev/null
+++ b/app/src/main/java/org/astraea/partitioner/PartitionInfo.java
@@ -0,0 +1,32 @@
+package org.astraea.partitioner;
+
+public interface PartitionInfo {
+
+  static PartitionInfo of(org.apache.kafka.common.PartitionInfo pf) {
+    return new PartitionInfo() {
+      @Override
+      public String topic() {
+        return pf.topic();
+      }
+
+      @Override
+      public int partition() {
+        return pf.partition();
+      }
+
+      @Override
+      public NodeInfo leader() {
+        return NodeInfo.of(pf.leader());
+      }
+    };
+  }
+
+  /** @return topic name */
+  String topic();
+
+  /** @return partition id */
+  int partition();
+
+  /** @return information of leader node */
+  NodeInfo leader();
+}
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
index ba3ae5be..f1269bd5 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
@@ -2,7 +2,6 @@ package org.astraea.partitioner.nodeLoadMetric;
 
 import static org.astraea.Utils.overSecond;
 
-import java.io.IOException;
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.ArrayList;
@@ -12,11 +11,11 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
-import org.apache.kafka.common.Cluster;
 import org.astraea.Utils;
 import org.astraea.metrics.collector.Receiver;
 import org.astraea.metrics.java.HasJvmMemory;
 import org.astraea.metrics.kafka.KafkaMetrics;
+import org.astraea.partitioner.ClusterInfo;
 
 /**
  * this clas is responsible for obtaining jmx metrics from BeanCollector and calculating the
@@ -39,7 +38,7 @@ public class NodeLoadClient {
   private static final String regex =
       "((25[0-5]|2[0-4]\\d|1\\d{2}|[1-9]?\\d)\\.){3}" + "(25[0-5]|2[0-4]\\d|1\\d{2}|[1-9]?\\d)$";
 
-  public NodeLoadClient(Map<String, Integer> jmxAddresses) throws IOException {
+  public NodeLoadClient(Map<String, Integer> jmxAddresses) {
     this.receiverList = RECEIVER_FACTORY.receiversList(jmxAddresses);
     receiverList.forEach(receiver -> Utils.waitFor(() -> receiver.current().size() > 0));
     currentJmxAddresses = jmxAddresses;
@@ -52,7 +51,7 @@ public class NodeLoadClient {
    * @param cluster from partitioner
    * @return each node load count in preset time
    */
-  public Map<Integer, Integer> loadSituation(Cluster cluster) throws UnknownHostException {
+  public Map<Integer, Integer> loadSituation(ClusterInfo cluster) {
     if (overSecond(lastTime, 1) && notInMethod) {
       notInMethod = false;
       nodesOverLoad(cluster);
@@ -63,7 +62,7 @@ public class NodeLoadClient {
   }
 
   /** @param cluster the cluster from the partitioner */
-  private synchronized void nodesOverLoad(Cluster cluster) {
+  private synchronized void nodesOverLoad(ClusterInfo cluster) {
     if (lastTime == -1) {
       var nodes = cluster.nodes();
       brokers =
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
index 38729d0d..81411eab 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
@@ -39,4 +39,21 @@ public class PartitionerUtils {
         overLoadCount.values().stream().mapToDouble(Integer::doubleValue).average().orElse(0);
     return (int) avgLoadCount;
   }
+
+  /**
+   * A cheap way to deterministically convert a number to a positive value. When the input is
+   * positive, the original value is returned. When the input number is negative, the returned
+   * positive value is the original value bit AND against 0x7fffffff which is not its absolute
+   * value.
+   *
+   * <p>Note: changing this method in the future will possibly cause partition selection not to be
+   * compatible with the existing messages already placed on a partition since it is used in
+   * producer's {@link org.apache.kafka.clients.producer.internals.DefaultPartitioner}
+   *
+   * @param number a given number
+   * @return a positive number.
+   */
+  public static int toPositive(int number) {
+    return number & 0x7fffffff;
+  }
 }
diff --git a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
index a2e41fe4..5d2a53d6 100644
--- a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
@@ -3,8 +3,6 @@ package org.astraea.partitioner.smoothPartitioner;
 import static org.astraea.Utils.overSecond;
 import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.allPoisson;
 
-import java.io.IOException;
-import java.net.UnknownHostException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -14,6 +12,7 @@ import java.util.Random;
 import java.util.concurrent.atomic.AtomicReference;
 import org.apache.kafka.clients.producer.Partitioner;
 import org.apache.kafka.common.Cluster;
+import org.astraea.partitioner.ClusterInfo;
 import org.astraea.partitioner.nodeLoadMetric.NodeLoadClient;
 
 /**
@@ -38,7 +37,7 @@ public class SmoothWeightPartitioner implements Partitioner {
   public int partition(
       String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
     Map<Integer, Integer> loadCount;
-    loadCount = loadCount(cluster);
+    loadCount = nodeLoadClient.loadSituation(ClusterInfo.of(cluster));
     Objects.requireNonNull(loadCount, "OverLoadCount should not be null.");
     updateWeightIfNeed(loadCount);
     AtomicReference<Map.Entry<Integer, int[]>> maxWeightServer = new AtomicReference<>();
@@ -79,23 +78,19 @@ public class SmoothWeightPartitioner implements Partitioner {
 
   @Override
   public void configure(Map<String, ?> configs) {
-    try {
-      var jmxAddresses =
-          Objects.requireNonNull(
-              configs.get("jmx_servers").toString(), "You must configure jmx_servers correctly");
-      var list = Arrays.asList((jmxAddresses).split(","));
-      HashMap<String, Integer> mapAddress = new HashMap<>();
-      list.forEach(
-          str ->
-              mapAddress.put(
-                  Arrays.asList(str.split(":")).get(0),
-                  Integer.parseInt(Arrays.asList(str.split(":")).get(1))));
-      Objects.requireNonNull(mapAddress, "You must configure jmx_servers correctly.");
-
-      nodeLoadClient = new NodeLoadClient(mapAddress);
-    } catch (IOException e) {
-      throw new RuntimeException();
-    }
+    var jmxAddresses =
+        Objects.requireNonNull(
+            configs.get("jmx_servers").toString(), "You must configure jmx_servers correctly");
+    var list = Arrays.asList((jmxAddresses).split(","));
+    HashMap<String, Integer> mapAddress = new HashMap<>();
+    list.forEach(
+        str ->
+            mapAddress.put(
+                Arrays.asList(str.split(":")).get(0),
+                Integer.parseInt(Arrays.asList(str.split(":")).get(1))));
+    Objects.requireNonNull(mapAddress, "You must configure jmx_servers correctly.");
+
+    nodeLoadClient = new NodeLoadClient(mapAddress);
   }
 
   /** Change the weight of the node according to the current Poisson. */
@@ -138,12 +133,4 @@ public class SmoothWeightPartitioner implements Partitioner {
   private boolean memoryWarning(int brokerID) {
     return nodeLoadClient.memoryUsage(brokerID) >= 0.8;
   }
-
-  private Map<Integer, Integer> loadCount(Cluster cluster) {
-    try {
-      return nodeLoadClient.loadSituation(cluster);
-    } catch (UnknownHostException e) {
-      throw new IllegalArgumentException(e);
-    }
-  }
 }
diff --git a/docker/start_broker.sh b/docker/start_broker.sh
index f63594f4..187cd986 100755
--- a/docker/start_broker.sh
+++ b/docker/start_broker.sh
@@ -4,7 +4,7 @@ declare -r DOCKER_FOLDER=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null
 source $DOCKER_FOLDER/docker_build_common.sh
 
 # ===============================[global variables]===============================
-declare -r VERSION=${REVISION:-${VERSION:-2.8.1}}
+declare -r VERSION=${REVISION:-${VERSION:-3.1.0}}
 declare -r REPO=${REPO:-ghcr.io/skiptests/astraea/broker}
 declare -r IMAGE_NAME="$REPO:$VERSION"
 declare -r DOCKERFILE=$DOCKER_FOLDER/broker.dockerfile
@@ -46,7 +46,7 @@ function showHelp() {
   echo "    REPO=astraea/broker                      set the docker repo"
   echo "    HEAP_OPTS=\"-Xmx2G -Xms2G\"                set broker JVM memory"
   echo "    REVISION=trunk                           set revision of kafka source code to build container"
-  echo "    VERSION=2.8.1                            set version of kafka distribution"
+  echo "    VERSION=3.1.0                            set version of kafka distribution"
   echo "    BUILD=false                              set true if you want to build image locally"
   echo "    RUN=false                                set false if you want to build/pull image only"
   echo "    DATA_FOLDERS=/tmp/folder1,/tmp/folder2   set host folders used by broker"


Test patch:
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
index 8fc7e8b2..5daa7582 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
@@ -10,8 +10,8 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
-import org.apache.kafka.common.Cluster;
-import org.apache.kafka.common.Node;
+import org.astraea.partitioner.ClusterInfo;
+import org.astraea.partitioner.NodeInfo;
 import org.astraea.service.RequireBrokerCluster;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
@@ -82,16 +82,16 @@ public class NodeLoadClientTest extends RequireBrokerCluster {
   @Test
   void testLoadSituation() throws UnknownHostException {
     var bootstrapServers = List.of(bootstrapServers().split(","));
-    List<Node> nodes = new ArrayList<>();
+    List<NodeInfo> nodes = new ArrayList<>();
     AtomicInteger count = new AtomicInteger(0);
     bootstrapServers.forEach(
         str -> {
           var hostPort = str.split(":");
-          nodes.add(new Node(count.get(), hostPort[0], Integer.parseInt(hostPort[1])));
+          nodes.add(NodeInfo.of(count.get(), hostPort[0], Integer.parseInt(hostPort[1])));
           count.getAndIncrement();
         });
 
-    Cluster cluster = Mockito.mock(Cluster.class);
+    var cluster = Mockito.mock(ClusterInfo.class);
     when(cluster.nodes()).thenReturn(nodes);
     var load = nodeLoadClient.loadSituation(cluster);
     Assertions.assertEquals(load.get(0), 5);