Implement poisson partitioner

app/build.gradle

@@ -19,10 +19,12 @@ repositories {
 def versions = [
     kafka: project.properties['kafka.version'] ?: "2.8.0",
     junit: project.properties['junit.version'] ?: "5.7.2",
+    mockito: project.properties['mockito.version'] ?: "3.12.4",
 ]
 
 dependencies {
     testImplementation "org.junit.jupiter:junit-jupiter:${versions["junit"]}"
+    testImplementation "org.mockito:mockito-core:${versions["mockito"]}"
 
     implementation "org.apache.kafka:kafka-clients:${versions["kafka"]}"
 }

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

@@ -0,0 +1,53 @@
+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<Integer, int[]> brokerHashMap = new HashMap<>();
+    private LoadPoisson loadPoisson;
+
+    BrokersWeight(LoadPoisson loadPoisson){
+        this.loadPoisson = loadPoisson;
+    }
+
+    /**
+     * Change the weight of the node according to the current Poisson.
+     */
+    public void setBrokerHashMap() {
+        HashMap<Integer, Double> poissonMap = loadPoisson.setAllPoisson();
+
+        for (Map.Entry<Integer, 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 int getAllWeight() {
+        int allWeight = 0;
+        for (Map.Entry<Integer, int[]> entry:brokerHashMap.entrySet()) {
+            allWeight += entry.getValue()[0];
+        }
+        return allWeight;
+    }
+
+    public HashMap<Integer, int[]> getBrokerHashMap() {
+        return brokerHashMap;
+    }
+
+    public void setCurrentBrokerHashMap(HashMap<Integer, int[]> currentBrokerHashMap) {
+        brokerHashMap = currentBrokerHashMap;
+    }
+
+    //Only for test
+    public void setBrokerHashMapValue(Integer x, int y) {
+        brokerHashMap.put(x,new int[]{0,y});
+    }
+}

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

@@ -0,0 +1,40 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class LoadPoisson {
+
+    public HashMap<Integer, Double> setAllPoisson(){
+        HashMap<Integer, Double> poissonMap = new HashMap<>();
+        NodeLoadClient nodeLoadClient = new NodeLoadClient();
+        int lambda = nodeLoadClient.getAvgLoadCount();
+        for (Map.Entry<Integer,Integer> entry :nodeLoadClient.getOverLoadCount().entrySet()){
+            int x = nodeLoadClient.getBinOneCount(entry.getKey());
+            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,65 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import java.util.*;
+
+public class NodeLoadClient {
+
+    /**
+     *This value records the number of times each node has been overloaded within ten seconds.
+     */
+    private static HashMap<Integer, Integer> overLoadCount = new HashMap<Integer, Integer>();
+
+    public static void setOverLoadCount() {
+        Timer timer = new Timer();
+        timer.schedule(new TimerTask() {
+            public OverLoadNode overLoadNode = new OverLoadNode();
+
+            public void run() {
+                overLoadNode.monitorOverLoad(overLoadCount);
+            }
+        }, 1 , 1000);
+    }
+
+
+    public HashMap<Integer, Integer> getOverLoadCount() {
+        return this.overLoadCount;
+    }
+
+    public int getAvgLoadCount() {
+        double avgLoadCount = 0;
+        for (Map.Entry<Integer,Integer> entry : overLoadCount.entrySet()) {
+            avgLoadCount += getBinOneCount(entry.getValue());
+        }
+        return (int)avgLoadCount/overLoadCount.size();
+    }
+
+    /**
+     * Get the number of times a node is overloaded.
+     */
+    public 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;
+    }
+
+    //TODO
+    private Collection<Integer> getNodeID() {
+        return null;
+    }
+    //TODO
+    private Integer getNodeOverLoadCount(Integer nodeID){
+        return 0;
+    }
+    //TODO
+    private static int[] getNodesID() {
+        return null;
+    }
+}

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

@@ -0,0 +1,105 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class OverLoadNode {
+    private double standardDeviation = 0;
+    private double avgBrokersMsgPerSec = 0;
+    private int[] nodesID;
+    private int nodeNum;
+    private int mountCount;
+    private HashMap<Integer, Double> eachBrokerMsgPerSec = new HashMap();
+
+    OverLoadNode() {
+        this.nodesID = getNodesID();
+        this.nodeNum = getNodesID().length;
+    }
+
+    /**
+     * Monitor and update the number of overloads of each node.
+     */
+    public void monitorOverLoad(HashMap<Integer, Integer> overLoadCount) {
+        int ifOverLoad = 0;
+        setBrokersMsgPerSec();
+        setAvgBrokersMsgPerSec();
+        standardDeviationImperative();
+        for (Map.Entry<Integer,Double> entry : eachBrokerMsgPerSec.entrySet()){
+            if (entry.getValue() > (avgBrokersMsgPerSec + standardDeviation)) {
+                ifOverLoad = 1;
+            }
+            overLoadCount.put(entry.getKey(),setOverLoadCount(overLoadCount.get(entry.getKey()), mountCount%10, ifOverLoad));
+        }
+        this.mountCount = mountCount++;
+    }
+
+    /**
+     *Use bit operations to record whether the node exceeds the load per second,the position of the number represents the recorded time.
+     */
+    public int setOverLoadCount(int overLoadCount, int roundCount,int ifOverLoad){
+        int x = overLoadCount&1<<roundCount;
+        if(x == ifOverLoad<<roundCount){
+            return overLoadCount;
+        }else {
+            if (ifOverLoad!=0){
+                return overLoadCount | 1<<roundCount;
+            }
+            else {
+                return overLoadCount - (int)Math.pow(2, roundCount);
+            }
+        }
+    }
+
+    public void setBrokersMsgPerSec() {
+        for (int nodeID : nodesID){
+            eachBrokerMsgPerSec.put(nodeID, getEachBrokerMsgPerSec(nodeID));
+        }
+    }
+
+    public void setAvgBrokersMsgPerSec() {
+        double avg = 0;
+        for (Map.Entry<Integer,Double> entry : eachBrokerMsgPerSec.entrySet()) {
+            avg += entry.getValue();
+        }
+        this.avgBrokersMsgPerSec = avg/nodeNum;
+    }
+
+    public void standardDeviationImperative() {
+        double variance = 0;
+        for (Map.Entry<Integer,Double> entry : eachBrokerMsgPerSec.entrySet()) {
+            variance += (entry.getValue() - avgBrokersMsgPerSec) * (entry.getValue() - avgBrokersMsgPerSec);
+        }
+
+        this.standardDeviation = Math.sqrt(variance / nodeNum);
+    }
+
+    //Only for test
+    public void setEachBrokerMsgPerSec(HashMap<Integer, Double> hashMap) {
+        this.eachBrokerMsgPerSec = hashMap;
+    }
+
+    //Only for test
+    public double getStandardDeviation() {
+        return this.standardDeviation;
+    }
+
+    //Only for test
+    public double getAvgBrokersMsgPerSec() {
+        return this.avgBrokersMsgPerSec;
+    }
+
+    //Only for test
+    public void setMountCount(int i) {
+        this.mountCount = i;
+    }
+
+    //TODO
+    private double getEachBrokerMsgPerSec(int nodeID){
+        return 0;
+    }
+
+    //TODO
+    private int[] getNodesID() {
+        return new int[] {0, 1, 2, 3};
+    }
+}

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

@@ -0,0 +1,60 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import org.apache.kafka.clients.producer.Partitioner;
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.PartitionInfo;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+
+import static org.astraea.partitioner.nodeLoadMetric.NodeLoadClient.setOverLoadCount;
+
+public class SmoothPartitioner implements Partitioner {
+
+    /**
+     * Implement Smooth Weight Round Robin.
+     */
+    @Override
+    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
+
+        //I think I need to find a suitable place to put it.
+        setOverLoadCount();
+
+        LoadPoisson loadPoisson = new LoadPoisson();
+        BrokersWeight brokersWeight = new BrokersWeight(loadPoisson);
+        brokersWeight.setBrokerHashMap();
+        Map.Entry<Integer, int[]> maxWeightServer = null;
+
+        int allWeight = brokersWeight.getAllWeight();
+        HashMap<Integer, int[]> currentBrokerHashMap = brokersWeight.getBrokerHashMap();
+
+        for (Map.Entry<Integer, int[]> item : currentBrokerHashMap.entrySet()) {
+            Map.Entry<Integer, int[]> currentServer = item;
+            if (maxWeightServer == null || currentServer.getValue()[1] > maxWeightServer.getValue()[1]) {
+                maxWeightServer = currentServer;
+            }
+        }
+        assert maxWeightServer != null;
+        currentBrokerHashMap.put(maxWeightServer.getKey(), new int[]{maxWeightServer.getValue()[0], maxWeightServer.getValue()[1]-allWeight});
+        brokersWeight.setCurrentBrokerHashMap(currentBrokerHashMap);
+
+        ArrayList<Integer> partitionList = new ArrayList<>();
+        for (PartitionInfo partitionInfo : cluster.partitionsForNode(maxWeightServer.getKey())){
+            partitionList.add(partitionInfo.partition());
+        }
+        Random rand = new Random();
+        return partitionList.get(rand.nextInt(partitionList.size()));
+    }
+
+    @Override
+    public void close() {
+
+    }
+
+    @Override
+    public void configure(Map<String, ?> configs) {
+
+    }
+}

app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java

@@ -0,0 +1,13 @@
+package org.astraea.partitioner.nodeLoadMetric;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class NodeLoadClientTest {
+    @Test
+    public void testGetBinOneCount() {
+        NodeLoadClient nodeLoadClient = new NodeLoadClient();
+        assertEquals(nodeLoadClient.getBinOneCount(7), 3);
+    }
+}