Implement poisson partitioner

app/build.gradle

@@ -18,18 +18,22 @@ repositories {
 }
 
 def versions = [
-    kafka: project.properties['kafka.version'] ?: "2.8.1",
+    kafka: project.properties['kafka.version'] ?: "2.8.0",
     junit: project.properties['junit.version'] ?: "5.8.1",
+    mockito: project.properties['mockito.version'] ?: "3.12.4",
     jcommander: project.properties['jcommander.version'] ?: "1.81",
     slf4j: project.properties['slf4j.version'] ?: "1.7.32",
     zookeeper: project.properties['zookeeper.version'] ?: "3.5.9",
 ]
 
 dependencies {
     testImplementation "org.junit.jupiter:junit-jupiter:${versions["junit"]}"
+    testImplementation "org.mockito:mockito-core:${versions["mockito"]}"
+    testImplementation "org.mockito:mockito-inline:${versions["mockito"]}"
     testImplementation "org.apache.zookeeper:zookeeper:${versions["zookeeper"]}"
     testImplementation "org.apache.kafka:kafka_2.12:${versions["kafka"]}"
 
+
     implementation "org.apache.kafka:kafka-clients:${versions["kafka"]}"
     implementation "com.beust:jcommander:${versions["jcommander"]}"
     // we don't use slf4j actually, and it is used by kafka so we swallow the log.

app/src/main/java/org/astraea/partitioner/nodeLoadMetric/BrokersWeight.java

@@ -0,0 +1,55 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class BrokersWeight {
+
+  /**
+   * Record the current weight of each node according to Poisson calculation and the weight after
+   * partitioner calculation.
+   */
+  private static HashMap<String, int[]> brokerHashMap = new HashMap<>();
+
+  private LoadPoisson loadPoisson;
+
+  public BrokersWeight(LoadPoisson loadPoisson) {
+    this.loadPoisson = loadPoisson;
+  }
+
+  /** Change the weight of the node according to the current Poisson. */
+  public synchronized void setBrokerHashMap() {
+    HashMap<String, Double> poissonMap = loadPoisson.setAllPoisson();
+
+    for (Map.Entry<String, Double> entry : poissonMap.entrySet()) {
+      if (!brokerHashMap.containsKey(entry.getKey())) {
+        brokerHashMap.put(entry.getKey(), new int[] {(int) ((1 - entry.getValue()) * 20), 0});
+      } else {
+        brokerHashMap.put(
+            entry.getKey(),
+            new int[] {(int) ((1 - entry.getValue()) * 20), brokerHashMap.get(entry.getKey())[1]});
+      }
+    }
+  }
+
+  public synchronized int getAllWeight() {
+    int allWeight = 0;
+    for (Map.Entry<String, int[]> entry : brokerHashMap.entrySet()) {
+      allWeight += entry.getValue()[0];
+    }
+    return allWeight;
+  }
+
+  public synchronized HashMap<String, int[]> getBrokerHashMap() {
+    return brokerHashMap;
+  }
+
+  public synchronized void setCurrentBrokerHashMap(HashMap<String, int[]> currentBrokerHashMap) {
+    brokerHashMap = currentBrokerHashMap;
+  }
+
+  // Only for test
+  void setBrokerHashMapValue(String x, int y) {
+    brokerHashMap.put(x, new int[] {0, y});
+  }
+}

app/src/main/java/org/astraea/partitioner/nodeLoadMetric/LoadPoisson.java

@@ -0,0 +1,42 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class LoadPoisson {
+  private NodeLoadClient nodeLoadClient;
+
+  public LoadPoisson(NodeLoadClient nodeLoadClient) {
+    this.nodeLoadClient = nodeLoadClient;
+  }
+
+  public synchronized HashMap<String, Double> setAllPoisson() {
+    HashMap<String, Double> poissonMap = new HashMap<>();
+    int lambda = nodeLoadClient.getAvgLoadCount();
+    for (Map.Entry<String, Integer> entry : nodeLoadClient.getAllOverLoadCount().entrySet()) {
+      int x = nodeLoadClient.getBinOneCount(entry.getValue());
+      poissonMap.put(entry.getKey(), doPoisson(lambda, x));
+    }
+    return poissonMap;
+  }
+
+  public double doPoisson(int lambda, int x) {
+    double Probability = 0;
+    double ans = 0;
+
+    for (int i = 0; i <= x; i++) {
+      double j = Math.pow(lambda, i);
+      double e = Math.exp(-lambda);
+      long h = factorial(i);
+      Probability = (j * e) / h;
+      ans += Probability;
+    }
+
+    return ans;
+  }
+
+  public long factorial(long number) {
+    if (number <= 1) return 1;
+    else return number * factorial(number - 1);
+  }
+}

app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java

@@ -0,0 +1,94 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.concurrent.TimeUnit;
+import org.astraea.concurrent.ThreadPool;
+import org.astraea.metrics.jmx.MBeanClient;
+
+public class NodeLoadClient implements ThreadPool.Executor {
+
+  private final OverLoadNode overLoadNode;
+  private final Collection<NodeMetadata> nodeMetadataCollection = new ArrayList<>();
+
+  public NodeLoadClient(HashMap<String, String> jmxAddresses) throws IOException {
+    for (HashMap.Entry<String, String> entry : jmxAddresses.entrySet()) {
+      this.nodeMetadataCollection.add(
+          new NodeMetadata(entry.getKey(), createNodeMetrics(entry.getKey(), entry.getValue())));
+    }
+    this.overLoadNode = new OverLoadNode(this.nodeMetadataCollection);
+  }
+
+  public NodeMetrics createNodeMetrics(String key, String value) throws IOException {
+    return new NodeMetrics(key, value);
+  }
+
+  /** A thread that continuously updates metricsfor NodeLoadClient. */
+  @Override
+  public State execute() throws InterruptedException {
+    try {
+      refreshNodesMetrics();
+      overLoadNode.monitorOverLoad();
+      TimeUnit.SECONDS.sleep(1);
+    } catch (InterruptedException e) {
+      e.printStackTrace();
+    }
+    return null;
+  }
+
+  @Override
+  public void close() {
+    for (NodeMetadata nodeMetadata : nodeMetadataCollection) {
+      NodeMetrics nodeMetrics = nodeMetadata.getNodeMetrics();
+      MBeanClient mBeanClient = nodeMetrics.getKafkaMetricClient();
+      try {
+        mBeanClient.close();
+      } catch (Exception e) {
+        e.printStackTrace();
+      }
+    }
+  }
+
+  public synchronized HashMap<String, Integer> getAllOverLoadCount() {
+    HashMap<String, Integer> overLoadCount = new HashMap<>();
+    for (NodeMetadata nodeMetadata : nodeMetadataCollection) {
+      overLoadCount.put(nodeMetadata.getNodeID(), nodeMetadata.getOverLoadCount());
+    }
+    return overLoadCount;
+  }
+
+  public synchronized int getAvgLoadCount() {
+    double avgLoadCount = 0;
+    for (NodeMetadata nodeMetadata : nodeMetadataCollection) {
+      avgLoadCount += getBinOneCount(nodeMetadata.getOverLoadCount());
+    }
+    return nodeMetadataCollection.size() > 0
+        ? (int) avgLoadCount / nodeMetadataCollection.size()
+        : 0;
+  }
+
+  /** Get the number of times a node is overloaded. */
+  public static int getBinOneCount(int n) {
+    int index = 0;
+    int count = 0;
+    while (n > 0) {
+      int x = n & 1 << index;
+      if (x != 0) {
+        count++;
+        n = n - (1 << index);
+      }
+      index++;
+    }
+    return count;
+  }
+
+  public void refreshNodesMetrics() {
+    for (NodeMetadata nodeMetadata : nodeMetadataCollection) {
+      NodeMetrics nodeMetrics = nodeMetadata.getNodeMetrics();
+      nodeMetrics.refreshMetrics();
+      nodeMetadata.setTotalBytes(nodeMetrics.totalBytesPerSec());
+    }
+  }
+}

app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeMetadata.java

@@ -0,0 +1,40 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+/** Store information about each node */
+public class NodeMetadata {
+  private String nodeID;
+  private NodeMetrics nodeMetrics;
+  private double totalBytes;
+  private int overLoadCount;
+
+  NodeMetadata(String nodeID, NodeMetrics nodeMetrics) {
+    this.nodeID = nodeID;
+    this.nodeMetrics = nodeMetrics;
+    this.overLoadCount = 0;
+    this.totalBytes = 0.0;
+  }
+
+  public NodeMetrics getNodeMetrics() {
+    return this.nodeMetrics;
+  }
+
+  public void setOverLoadCount(int count) {
+    this.overLoadCount = count;
+  }
+
+  public void setTotalBytes(double bytes) {
+    this.totalBytes = bytes;
+  }
+
+  public double getTotalBytes() {
+    return this.totalBytes;
+  }
+
+  public String getNodeID() {
+    return this.nodeID;
+  }
+
+  public int getOverLoadCount() {
+    return this.overLoadCount;
+  }
+}

app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeMetrics.java

@@ -0,0 +1,62 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import javax.management.remote.JMXServiceURL;
+import org.astraea.metrics.jmx.MBeanClient;
+import org.astraea.metrics.kafka.BrokerTopicMetricsResult;
+import org.astraea.metrics.kafka.KafkaMetrics;
+
+/** Responsible for connecting jmx according to the received address */
+public class NodeMetrics {
+  private final String JMX_URI_FORMAT = "service:jmx:rmi:///jndi/rmi://" + "%s" + "/jmxrmi";
+  private final JMXServiceURL serviceURL;
+  private final MBeanClient mBeanClient;
+  private final String nodeID;
+  private HashMap<String, Double> metricsValues;
+  private Collection<String> argumentTargetMetrics = new ArrayList<>();
+
+  NodeMetrics(String ID, String address) throws IOException {
+    argumentTargetMetrics.add("BytesInPerSec");
+    argumentTargetMetrics.add("BytesOutPerSec");
+    nodeID = ID;
+    if (Pattern.compile("^service:").matcher(address).find())
+      serviceURL = new JMXServiceURL(address);
+    else serviceURL = new JMXServiceURL(createJmxUrl(address));
+    mBeanClient = new MBeanClient(serviceURL);
+    metricsValues = new HashMap();
+  }
+
+  public String createJmxUrl(String address) {
+    return String.format(JMX_URI_FORMAT, address);
+  }
+
+  public void refreshMetrics() {
+    List<KafkaMetrics.BrokerTopic> metrics =
+        argumentTargetMetrics.stream()
+            .map(KafkaMetrics.BrokerTopic::of)
+            .collect(Collectors.toUnmodifiableList());
+
+    List<BrokerTopicMetricsResult> collect =
+        metrics.stream().map(x -> x.fetch(mBeanClient)).collect(Collectors.toUnmodifiableList());
+
+    for (BrokerTopicMetricsResult result : collect) {
+      metricsValues.put(
+          result.beanObject().getProperties().get("name"),
+          (Double) result.beanObject().getAttributes().get("MeanRate"));
+    }
+  }
+
+  public double totalBytesPerSec() {
+    return metricsValues.get("BytesInPerSec") + metricsValues.get("BytesOutPerSec");
+  }
+
+  public MBeanClient getKafkaMetricClient() {
+    return this.mBeanClient;
+  }
+}