Repository: skiptests/astraea
Pull request #241 (2022-02-18 12:57:43)

Optimization Partitioner
優化Partitioner效能，現在在10KiB的record.size下也能保持一定的效能。
右邊爲改善過後的astraea partitioner的數據，原先的數據結果大概爲30min avg=500MB/s
![Screenshot from 2022-02-18 20-25-17](https://user-images.githubusercontent.com/43372856/154685915-a6ab6538-bb36-4b45-a46d-f32edbc40c40.png)

主要的改善點爲減少不必要的該方法`brokersWeight`調用。JMC告訴我該方法佔了整個程式12%的cpu使用量。
而如果metrics不變這個方法的結果就不會改變。metrics最快也是1s更新一次，所以改爲超過1s再進行一次計算。




Patch:
diff --git a/app/src/main/java/org/astraea/Utils.java b/app/src/main/java/org/astraea/Utils.java
index b022c9ec..298ec86b 100644
--- a/app/src/main/java/org/astraea/Utils.java
+++ b/app/src/main/java/org/astraea/Utils.java
@@ -137,5 +137,9 @@ public final class Utils {
     return value;
   }
 
+  public static boolean overSecond(long lastTime, int second) {
+    return (lastTime + Duration.ofSeconds(second).toMillis()) < System.currentTimeMillis();
+  }
+
   private Utils() {}
 }
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 95909263..ba3ae5be 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
@@ -1,15 +1,15 @@
 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.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import org.apache.kafka.common.Cluster;
@@ -32,8 +32,9 @@ public class NodeLoadClient {
   List<Broker> brokers;
   private Map<Integer, Integer> currentLoadNode;
   private int referenceBrokerID;
-  private CountDownLatch countDownLatch;
   private boolean notInMethod = true;
+  private double totalInput;
+  private double totalOutput;
 
   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)$";
@@ -51,25 +52,20 @@ public class NodeLoadClient {
    * @param cluster from partitioner
    * @return each node load count in preset time
    */
-  public synchronized Map<Integer, Integer> loadSituation(Cluster cluster)
-      throws UnknownHostException {
-    if (overOneSecond() && notInMethod) {
+  public Map<Integer, Integer> loadSituation(Cluster cluster) throws UnknownHostException {
+    if (overSecond(lastTime, 1) && notInMethod) {
       notInMethod = false;
       nodesOverLoad(cluster);
+      notInMethod = true;
     }
-    try {
-      countDownLatch.await();
-    } catch (InterruptedException e) {
-      e.printStackTrace();
-    }
+    Utils.waitFor(() -> currentLoadNode != null);
     return currentLoadNode;
   }
 
   /** @param cluster the cluster from the partitioner */
   private synchronized void nodesOverLoad(Cluster cluster) {
-    var nodes = cluster.nodes();
     if (lastTime == -1) {
-      countDownLatch = new CountDownLatch(1);
+      var nodes = cluster.nodes();
       brokers =
           nodes.stream()
               .map(node -> new Broker(node.id(), node.host(), node.port()))
@@ -92,9 +88,8 @@ public class NodeLoadClient {
     }
 
     brokersMsgPerSec();
-
-    var totalInput = brokers.stream().map(broker -> broker.input).reduce(0.0, Double::sum);
-    var totalOutput = brokers.stream().map(broker -> broker.output).reduce(0.0, Double::sum);
+    totalInput = totalInput();
+    totalOutput = totalOutput();
 
     AtomicReference<Double> totalBrokerSituation = new AtomicReference<>(0.0);
     brokers.forEach(
@@ -107,21 +102,15 @@ public class NodeLoadClient {
           totalBrokerSituation.updateAndGet(v -> v + broker.brokerSituation);
         });
 
-    var totalSituation = totalBrokerSituation.get();
-
-    brokers.forEach(broker -> broker.brokerSituationNormalized(totalSituation));
+    brokers.forEach(broker -> broker.brokerSituationNormalized(totalBrokerSituation.get()));
 
-    var avg = totalSituation / brokers.size();
-
-    // TODO Regarding countermeasures against cluster conditions
-    var standardDeviation = standardDeviationImperative(avg);
-
-    brokers.forEach(broker -> broker.load(brokerLoad(broker.brokerSituation, avg)));
+    brokers.forEach(
+        broker ->
+            broker.load(
+                brokerLoad(broker.brokerSituation, totalBrokerSituation.get() / brokers.size())));
     currentLoadNode =
         brokers.stream().collect(Collectors.toMap(broker -> broker.brokerID, Broker::totalLoad));
     lastTime = System.currentTimeMillis();
-    countDownLatch.countDown();
-    notInMethod = true;
   }
 
   private void brokersMsgPerSec() {
@@ -149,9 +138,13 @@ public class NodeLoadClient {
                       .getAttributes()
                       .get("MeanRate")
                       .toString());
+          broker.jvm =
+              (HasJvmMemory)
+                  broker.metrics.get("Memory").current().stream().findAny().orElse(broker.jvm);
         });
   }
 
+  // TODO Regarding countermeasures against cluster conditions
   public double standardDeviationImperative(double avgBrokersSituation) {
     var variance = new AtomicReference<>(0.0);
     brokers.forEach(
@@ -225,10 +218,6 @@ public class NodeLoadClient {
     return !host.matches(regex);
   }
 
-  private boolean overOneSecond() {
-    return lastTime + Duration.ofSeconds(1).toMillis() <= System.currentTimeMillis();
-  }
-
   public void close() {
     RECEIVER_FACTORY.close(currentJmxAddresses);
   }
@@ -240,6 +229,14 @@ public class NodeLoadClient {
         .collect(Collectors.toList());
   }
 
+  private double totalInput() {
+    return brokers.stream().map(broker -> broker.input).reduce(0.0, Double::sum);
+  }
+
+  private double totalOutput() {
+    return brokers.stream().map(broker -> broker.output).reduce(0.0, Double::sum);
+  }
+
   // visible of test
   int brokerLoad(double brokerSituation, double avg) {
     var initialization = 5;
@@ -274,6 +271,7 @@ public class NodeLoadClient {
     private double inputNormalized;
     private double outputNormalized;
     private double brokerSituationNormalized;
+    private HasJvmMemory jvm;
     private static final double inputWeights = 0.5;
     private static final double outputWeights = 0.5;
 
@@ -314,13 +312,12 @@ public class NodeLoadClient {
       this.maxThoughPut = Math.max(maxThoughPut, input + output);
     }
 
+    // TODO Algorithms adapted to computers of different specifications
     private void thoughPutComparison(double referenceThoughPut) {
       if (count >= 20) {
-        this.thoughPutComparison =
-            (maxThoughPut - referenceThoughPut) / (referenceThoughPut + 1) + 1;
+        this.thoughPutComparison = 1;
       } else if (count >= 10) {
-        this.thoughPutComparison =
-            (maxThoughPut - referenceThoughPut) / (referenceThoughPut + 1) / 2 + 1;
+        this.thoughPutComparison = 1;
       } else if (count >= 0) {
         this.thoughPutComparison = 1;
       }
@@ -330,7 +327,6 @@ public class NodeLoadClient {
       if (count < 10) {
         return 0.1;
       } else {
-        var jvm = (HasJvmMemory) metrics.get("Memory").current().stream().findAny().get();
         return (jvm.heapMemoryUsage().getUsed() + 0.0) / (jvm.heapMemoryUsage().getMax() + 1);
       }
     }
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/LoadPoisson.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
similarity index 73%
rename from app/src/main/java/org/astraea/partitioner/nodeLoadMetric/LoadPoisson.java
rename to app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
index b94c6723..38729d0d 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/LoadPoisson.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
@@ -4,17 +4,17 @@ import java.util.HashMap;
 import java.util.Map;
 
 /** do poisson for node's load situation */
-public class LoadPoisson {
+public class PartitionerUtils {
+  private PartitionerUtils() {}
 
-  public HashMap<Integer, Double> allPoisson(Map<Integer, Integer> overLoadCount) {
+  public static HashMap<Integer, Double> allPoisson(Map<Integer, Integer> overLoadCount) {
     var poissonMap = new HashMap<Integer, Double>();
     var lambda = avgLoadCount(overLoadCount);
     overLoadCount.forEach((nodeID, count) -> poissonMap.put(nodeID, doPoisson(lambda, count)));
     return poissonMap;
   }
 
-  // visible for test
-  double doPoisson(int lambda, int x) {
+  static double doPoisson(int lambda, int x) {
     var Probability = 0.0;
     var ans = 0.0;
     var i = 0;
@@ -29,13 +29,12 @@ public class LoadPoisson {
     return ans;
   }
 
-  // visible for test
-  long factorial(long number) {
+  static long factorial(long number) {
     if (number <= 1) return 1;
     else return number * factorial(number - 1);
   }
 
-  private int avgLoadCount(Map<Integer, Integer> overLoadCount) {
+  private static int avgLoadCount(Map<Integer, Integer> overLoadCount) {
     var avgLoadCount =
         overLoadCount.values().stream().mapToDouble(Integer::doubleValue).average().orElse(0);
     return (int) avgLoadCount;
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 76585318..a2e41fe4 100644
--- a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
@@ -1,5 +1,8 @@
 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;
@@ -11,7 +14,6 @@ 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.nodeLoadMetric.LoadPoisson;
 import org.astraea.partitioner.nodeLoadMetric.NodeLoadClient;
 
 /**
@@ -29,18 +31,17 @@ public class SmoothWeightPartitioner implements Partitioner {
 
   private NodeLoadClient nodeLoadClient;
 
+  private long lastTime = -1;
+  private final Random rand = new Random();
+
   @Override
   public int partition(
       String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
     Map<Integer, Integer> loadCount;
-    var rand = new Random();
     loadCount = loadCount(cluster);
-    var loadPoisson = new LoadPoisson();
-
     Objects.requireNonNull(loadCount, "OverLoadCount should not be null.");
-    brokersWeight(loadPoisson.allPoisson(loadCount));
+    updateWeightIfNeed(loadCount);
     AtomicReference<Map.Entry<Integer, int[]>> maxWeightServer = new AtomicReference<>();
-    var allWeight = allNodesWeight();
     brokersWeight.forEach(
         (nodeID, weight) -> {
           if (maxWeightServer.get() == null || weight[1] > maxWeightServer.get().getValue()[1]) {
@@ -52,7 +53,8 @@ public class SmoothWeightPartitioner implements Partitioner {
     brokersWeight.put(
         maxWeightServer.get().getKey(),
         new int[] {
-          maxWeightServer.get().getValue()[0], maxWeightServer.get().getValue()[1] - allWeight
+          maxWeightServer.get().getValue()[0],
+          maxWeightServer.get().getValue()[1] - allNodesWeight()
         });
     brokersWeight
         .keySet()
@@ -112,6 +114,17 @@ public class SmoothWeightPartitioner implements Partitioner {
         });
   }
 
+  void updateWeightIfNeed(Map<Integer, Integer> loadCount) {
+    if (overSecond(lastTime, 1)) {
+      brokersWeight(allPoisson(loadCount));
+      lastTime = System.currentTimeMillis();
+    }
+  }
+
+  Map<Integer, int[]> brokersWeight() {
+    return brokersWeight;
+  }
+
   private void subBrokerWeight(int brokerID, int subNumber) {
     brokersWeight.put(
         brokerID,
diff --git a/docker/start_zookeeper.sh b/docker/start_zookeeper.sh
index aee8be0f..59612197 100755
--- a/docker/start_zookeeper.sh
+++ b/docker/start_zookeeper.sh
@@ -11,6 +11,7 @@ declare -r ZOOKEEPER_PORT=${ZOOKEEPER_PORT:-"$(getRandomPort)"}
 declare -r DOCKERFILE=$DOCKER_FOLDER/zookeeper.dockerfile
 declare -r DATA_FOLDER_IN_CONTAINER="/tmp/zookeeper-dir"
 declare -r CONTAINER_NAME="zookeeper-$ZOOKEEPER_PORT"
+declare -r HEAP_OPTS="${HEAP_OPTS:-"-Xmx1G -Xms256m"}"
 
 # ===================================[functions]===================================
 
@@ -85,11 +86,13 @@ if [[ -n "$DATA_FOLDER" ]]; then
   echo "mount $DATA_FOLDER to container"
   docker run -d --init \
     --name $CONTAINER_NAME \
+    -e JVMFLAGS="$HEAP_OPTS" \
     -p $ZOOKEEPER_PORT:2181 \
     -v "$DATA_FOLDER":$DATA_FOLDER_IN_CONTAINER \
     "$IMAGE_NAME" ./bin/zkServer.sh start-foreground
 else
   docker run -d --init \
+    -e JVMFLAGS="$HEAP_OPTS" \
     -p $ZOOKEEPER_PORT:2181 \
     "$IMAGE_NAME" ./bin/zkServer.sh start-foreground
 fi
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index 1a2cb8ee..976bbe0a 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,6 +1,6 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
 # please update the gradle version in gradlew too
-distributionUrl=https\://services.gradle.org/distributions/gradle-7.3.3-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-7.4-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
diff --git a/gradlew b/gradlew
index f4597373..017e6341 100755
--- a/gradlew
+++ b/gradlew
@@ -116,7 +116,7 @@ esac
 
 # gradle-wrapper.jar is not in git repo anymore, so we have to download it now.
 if [ ! -e "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" ]; then
-  if ! curl -s -S --retry 3 -L -o "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" "https://raw.githubusercontent.com/gradle/gradle/v7.3.3/gradle/wrapper/gradle-wrapper.jar"; then
+  if ! curl -s -S --retry 3 -L -o "$APP_HOME/gradle/wrapper/gradle-wrapper.jar" "https://raw.githubusercontent.com/gradle/gradle/v7.4.0/gradle/wrapper/gradle-wrapper.jar"; then
     rm -f "$APP_HOME/gradle/wrapper/gradle-wrapper.jar"
     echo "failed to download gradle-wrapper.jar"
     exit 1


Test patch:
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/LoadPoissonTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
similarity index 55%
rename from app/src/test/java/org/astraea/partitioner/nodeLoadMetric/LoadPoissonTest.java
rename to app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
index a87672e5..8729045e 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/LoadPoissonTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
@@ -1,28 +1,25 @@
 package org.astraea.partitioner.nodeLoadMetric;
 
+import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.*;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import java.util.HashMap;
 import org.junit.jupiter.api.Test;
 
-public class LoadPoissonTest {
+public class PartitionerUtilsTest {
 
   @Test
   void testFactorial() {
-    LoadPoisson loadPoisson = new LoadPoisson();
-
-    assertEquals(loadPoisson.factorial(3), 6);
-    assertEquals(loadPoisson.factorial(5), 120);
+    assertEquals(factorial(3), 6);
+    assertEquals(factorial(5), 120);
   }
 
   @Test
   void testDoPoisson() {
-    LoadPoisson loadPoisson = new LoadPoisson();
-
-    assertEquals(loadPoisson.doPoisson(10, 15), 0.9512595966960214);
-    assertEquals(loadPoisson.doPoisson(5, 5), 0.6159606548330632);
-    assertEquals(loadPoisson.doPoisson(5, 8), 0.9319063652781515);
-    assertEquals(loadPoisson.doPoisson(5, 2), 0.12465201948308113);
+    assertEquals(doPoisson(10, 15), 0.9512595966960214);
+    assertEquals(doPoisson(5, 5), 0.6159606548330632);
+    assertEquals(doPoisson(5, 8), 0.9319063652781515);
+    assertEquals(doPoisson(5, 2), 0.12465201948308113);
   }
 
   @Test
@@ -38,9 +35,7 @@ public class LoadPoissonTest {
     testPoissonMap.put(1, 0.6159606548330632);
     testPoissonMap.put(2, 0.9319063652781515);
 
-    LoadPoisson loadPoisson = new LoadPoisson();
-
-    poissonMap = loadPoisson.allPoisson(testNodesLoadCount);
+    poissonMap = allPoisson(testNodesLoadCount);
 
     assertEquals(poissonMap, testPoissonMap);
   }
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerTest.java b/app/src/test/java/org/astraea/partitioner/smoothPartitioner/PartitionerTest.java
similarity index 86%
rename from app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerTest.java
rename to app/src/test/java/org/astraea/partitioner/smoothPartitioner/PartitionerTest.java
index 6bbbca23..c8d65f6b 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerTest.java
+++ b/app/src/test/java/org/astraea/partitioner/smoothPartitioner/PartitionerTest.java
@@ -1,9 +1,9 @@
-package org.astraea.partitioner.nodeLoadMetric;
+package org.astraea.partitioner.smoothPartitioner;
 
 import static org.astraea.Utils.requireField;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.*;
 
 import java.lang.reflect.Field;
 import java.time.Duration;
@@ -22,7 +22,7 @@ import org.astraea.concurrent.ThreadPool;
 import org.astraea.consumer.Consumer;
 import org.astraea.consumer.Deserializer;
 import org.astraea.consumer.Header;
-import org.astraea.partitioner.smoothPartitioner.SmoothWeightPartitioner;
+import org.astraea.partitioner.nodeLoadMetric.NodeLoadClient;
 import org.astraea.producer.Producer;
 import org.astraea.producer.Serializer;
 import org.astraea.service.RequireBrokerCluster;
@@ -182,6 +182,30 @@ public class PartitionerTest extends RequireBrokerCluster {
     }
   }
 
+  @Test
+  void testUpdateWeightIfNeed() throws NoSuchFieldException, IllegalAccessException {
+    SmoothWeightPartitioner smoothWeightPartitioner = new SmoothWeightPartitioner();
+    var loadCount = new HashMap<Integer, Integer>();
+    loadCount.put(0, 10);
+    loadCount.put(1, 12);
+    var nodeLoadClient = mock(NodeLoadClient.class);
+    var field = smoothWeightPartitioner.getClass().getDeclaredField("nodeLoadClient");
+    field.setAccessible(true);
+    field.set(smoothWeightPartitioner, nodeLoadClient);
+    when(nodeLoadClient.thoughPutComparison(anyInt())).thenReturn(1.0);
+    smoothWeightPartitioner.updateWeightIfNeed(loadCount);
+    var firstBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0)[0];
+    smoothWeightPartitioner.updateWeightIfNeed(loadCount);
+    var secondBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0)[0];
+    Assertions.assertEquals(firstBrokersWeight, secondBrokersWeight);
+    loadCount.put(0, 5);
+    loadCount.put(2, 5);
+    sleep(2);
+    smoothWeightPartitioner.updateWeightIfNeed(loadCount);
+    var thirdBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0)[0];
+    Assertions.assertNotEquals(secondBrokersWeight, thirdBrokersWeight);
+  }
+
   private ThreadPool.Executor producerExecutor(
       Producer<String, byte[]> producer, String topic, String key, Header header, long timeStamp) {
     return new ThreadPool.Executor() {