Repository: skiptests/astraea
Pull request #198 (2022-01-24 15:50:22)

Run an init inside for zookeeper, broker, kafka-tool, and spark
spark在任務提交到k8s的腳本仰賴`tini`，而為了保持衛生乾淨，我將其他服務也都置於`tini`底下


Patch:
diff --git a/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java b/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java
index dac0c30f..cf5ed419 100644
--- a/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java
+++ b/app/src/main/java/org/astraea/metrics/collector/BeanCollector.java
@@ -1,7 +1,12 @@
 package org.astraea.metrics.collector;
 
 import java.time.Duration;
-import java.util.*;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentSkipListMap;
 import java.util.concurrent.locks.Lock;
@@ -121,7 +126,7 @@ public class BeanCollector {
                 }
               }
 
-              private void tryUpdate() {
+              private synchronized void tryUpdate() {
                 var needUpdate =
                     objects.keySet().stream()
                         .max((Long::compare))
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 03542570..95909263 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
@@ -1,24 +1,21 @@
 package org.astraea.partitioner.nodeLoadMetric;
 
-import static java.lang.Double.sum;
-
 import java.io.IOException;
 import java.net.InetAddress;
 import java.net.UnknownHostException;
+import java.time.Duration;
 import java.util.ArrayList;
-import java.util.Comparator;
 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.function.Function;
 import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 import org.apache.kafka.common.Cluster;
-import org.apache.kafka.common.Node;
 import org.astraea.Utils;
 import org.astraea.metrics.collector.Receiver;
+import org.astraea.metrics.java.HasJvmMemory;
 import org.astraea.metrics.kafka.KafkaMetrics;
 
 /**
@@ -28,12 +25,15 @@ import org.astraea.metrics.kafka.KafkaMetrics;
 public class NodeLoadClient {
 
   private final List<Receiver> receiverList;
-
-  private Map<Integer, Map<String, Receiver>> eachNodeIDMetrics;
-
   private final Map<String, Integer> currentJmxAddresses;
-
   private static final ReceiverFactory RECEIVER_FACTORY = new ReceiverFactory();
+  private long lastTime = -1;
+  // visible for testing
+  List<Broker> brokers;
+  private Map<Integer, Integer> currentLoadNode;
+  private int referenceBrokerID;
+  private CountDownLatch countDownLatch;
+  private boolean notInMethod = true;
 
   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)$";
@@ -45,193 +45,294 @@ public class NodeLoadClient {
   }
 
   /**
-   * @param cluster the cluster from the partitioner
-   * @return each node overLoad count in preset time
+   * When the interval between getting the load twice exceeds one second, update the load and
+   * return, otherwise directly return to the existing load.
+   *
+   * @param cluster from partitioner
+   * @return each node load count in preset time
    */
-  public Map<Integer, Integer> nodesOverLoad(Cluster cluster) throws UnknownHostException {
-    var nodeIDReceiver = new HashMap<Integer, List<Receiver>>();
-    var addresses =
-        cluster.nodes().stream()
-            .collect(Collectors.toMap(Node::id, node -> Map.entry(node.host(), node.port())));
-
-    for (Map.Entry<Integer, Map.Entry<String, Integer>> hostPort : addresses.entrySet()) {
-      List<Receiver> matchingNode;
-      var host = ipAddress(hostPort);
-      matchingNode =
-          receiverList.stream()
-              .filter(
-                  receiver ->
-                      (Objects.equals(nodeIDReceiver.size(), 0)
-                              || (nodeIDReceiver.values().stream()
-                                      .map(
-                                          n ->
-                                              Objects.equals(
-                                                  n.stream().findAny().get().host(), host))
-                                      .anyMatch(n -> n.equals(true))
-                                  && nodeIDReceiver.values().stream()
-                                      .map(
-                                          n ->
-                                              !Objects.equals(
-                                                  n.stream().findAny().get().port(),
-                                                  hostPort.getValue().getValue()))
-                                      .anyMatch(n -> n.equals(true))))
-                          && Objects.equals(receiver.host(), host))
-              .collect(Collectors.toList());
-
-      if (matchingNode.size() > 0) {
-        if (!nodeIDReceiver.containsKey(hostPort.getKey()))
-          nodeIDReceiver.put(hostPort.getKey(), matchingNode);
-        else matchingNode.forEach(node -> nodeIDReceiver.get(hostPort.getKey()).add(node));
-      }
+  public synchronized Map<Integer, Integer> loadSituation(Cluster cluster)
+      throws UnknownHostException {
+    if (overOneSecond() && notInMethod) {
+      notInMethod = false;
+      nodesOverLoad(cluster);
+    }
+    try {
+      countDownLatch.await();
+    } catch (InterruptedException e) {
+      e.printStackTrace();
     }
+    return currentLoadNode;
+  }
 
-    eachNodeIDMetrics = metricsNameForReceiver(nodeIDReceiver);
-    var eachBrokerMsgPerSec = brokersMsgPerSec();
-    var overLoadCount = new HashMap<Integer, Integer>();
-    eachBrokerMsgPerSec.keySet().forEach(e -> overLoadCount.put(e, 0));
-    var i = 0;
-    var minRecordSec =
-        eachBrokerMsgPerSec.values().stream()
-            .map(n -> n.size())
-            .min(Comparator.comparing(Integer::intValue))
-            .get();
-    while (i < minRecordSec) {
-      var avg = avgMsgSec(i, eachBrokerMsgPerSec);
-      var eachMsg = brokersMsgSec(i, eachBrokerMsgPerSec);
-      var standardDeviation = standardDeviationImperative(eachMsg, avg);
-      eachMsg.forEach(
-          (nodeID, msg) -> {
-            if (msg > (avg + standardDeviation)) {
-              overLoadCount.put(nodeID, overLoadCount.get(nodeID) + 1);
-            }
+  /** @param cluster the cluster from the partitioner */
+  private synchronized void nodesOverLoad(Cluster cluster) {
+    var nodes = cluster.nodes();
+    if (lastTime == -1) {
+      countDownLatch = new CountDownLatch(1);
+      brokers =
+          nodes.stream()
+              .map(node -> new Broker(node.id(), node.host(), node.port()))
+              .collect(Collectors.toList());
+      referenceBrokerID = brokers.stream().findFirst().get().brokerID;
+      var nodeIDReceiver = new HashMap<Integer, List<Receiver>>();
+      receiverList.forEach(
+          receiver -> {
+            var brokersID = brokerIDOfReceiver(receiver.host());
+            brokersID.forEach(
+                brokerID -> {
+                  if (nodeIDReceiver.containsKey(brokerID)) {
+                    nodeIDReceiver.get(brokerID).add(receiver);
+                  } else {
+                    nodeIDReceiver.put(brokerID, new ArrayList<>(List.of(receiver)));
+                  }
+                });
           });
-      i++;
+      metricsNameForReceiver(nodeIDReceiver);
     }
-    return overLoadCount;
-  }
-
-  /** @return data transfer per second per node */
-  private Map<Integer, List<Double>> brokersMsgPerSec() {
-    var eachMsg = new HashMap<Integer, List<Double>>();
-
-    for (Map.Entry<Integer, Map<String, Receiver>> entry : eachNodeIDMetrics.entrySet()) {
-      List<Double> sumList = new ArrayList<>();
-      var outMsg = entry.getValue().get(KafkaMetrics.BrokerTopic.BytesOutPerSec.toString());
-      var inMsg = entry.getValue().get(KafkaMetrics.BrokerTopic.BytesInPerSec.toString());
-      var outMsgBean = outMsg.current();
-      var inMsgBean = inMsg.current();
-      IntStream.range(0, Math.min(outMsgBean.size(), inMsgBean.size()))
-          .forEach(
-              i ->
-                  sumList.add(
-                      sum(
-                          Double.parseDouble(
-                              outMsgBean
-                                  .get(i)
-                                  .beanObject()
-                                  .getAttributes()
-                                  .get("MeanRate")
-                                  .toString()),
-                          Double.parseDouble(
-                              inMsgBean
-                                  .get(i)
-                                  .beanObject()
-                                  .getAttributes()
-                                  .get("MeanRate")
-                                  .toString()))));
-      eachMsg.put(entry.getKey(), sumList);
-    }
-    return eachMsg;
+
+    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);
+
+    AtomicReference<Double> totalBrokerSituation = new AtomicReference<>(0.0);
+    brokers.forEach(
+        broker -> {
+          broker.maxThoughPut();
+          broker.inputNormalized(totalInput);
+          broker.outputNormalized(totalOutput);
+          broker.thoughPutComparison(findBroker(referenceBrokerID).maxThoughPut);
+          broker.brokerSituation();
+          totalBrokerSituation.updateAndGet(v -> v + broker.brokerSituation);
+        });
+
+    var totalSituation = totalBrokerSituation.get();
+
+    brokers.forEach(broker -> broker.brokerSituationNormalized(totalSituation));
+
+    var avg = totalSituation / brokers.size();
+
+    // TODO Regarding countermeasures against cluster conditions
+    var standardDeviation = standardDeviationImperative(avg);
+
+    brokers.forEach(broker -> broker.load(brokerLoad(broker.brokerSituation, avg)));
+    currentLoadNode =
+        brokers.stream().collect(Collectors.toMap(broker -> broker.brokerID, Broker::totalLoad));
+    lastTime = System.currentTimeMillis();
+    countDownLatch.countDown();
+    notInMethod = true;
   }
 
-  public List<Double> avgBrokersMsgPerSec(Map<Integer, List<Double>> eachMsg) {
-    return IntStream.range(0, eachMsg.values().stream().findFirst().get().size())
-        .mapToDouble(i -> eachMsg.values().stream().map(s -> s.get(i)).reduce(0.0, Double::sum))
-        .map(sum -> sum / eachMsg.size())
-        .boxed()
-        .collect(Collectors.toList());
+  private void brokersMsgPerSec() {
+    brokers.forEach(
+        broker -> {
+          broker.output =
+              Double.parseDouble(
+                  broker
+                      .metrics
+                      .get(KafkaMetrics.BrokerTopic.BytesOutPerSec.name())
+                      .current()
+                      .get(0)
+                      .beanObject()
+                      .getAttributes()
+                      .get("MeanRate")
+                      .toString());
+          broker.input =
+              Double.parseDouble(
+                  broker
+                      .metrics
+                      .get(KafkaMetrics.BrokerTopic.BytesInPerSec.name())
+                      .current()
+                      .get(0)
+                      .beanObject()
+                      .getAttributes()
+                      .get("MeanRate")
+                      .toString());
+        });
   }
 
-  public double standardDeviationImperative(
-      Map<Integer, Double> eachMsg, double avgBrokersMsgPerSec) {
+  public double standardDeviationImperative(double avgBrokersSituation) {
     var variance = new AtomicReference<>(0.0);
-    eachMsg.forEach(
-        (nodeID, msg) ->
+    brokers.forEach(
+        broker ->
             variance.updateAndGet(
-                v -> v + (msg - avgBrokersMsgPerSec) * (msg - avgBrokersMsgPerSec)));
-    return Math.sqrt(variance.get() / eachMsg.size());
+                v ->
+                    v
+                        + (broker.brokerSituationNormalized - avgBrokersSituation)
+                            * (broker.brokerSituationNormalized - avgBrokersSituation)));
+    return Math.sqrt(variance.get() / brokers.size());
   }
 
-  /**
-   * @param nodeIDReceiver nodes metrics that exists in topic
-   * @return the name and corresponding data of the metrics obtained by each node
-   */
-  public Map<Integer, Map<String, Receiver>> metricsNameForReceiver(
-      HashMap<Integer, List<Receiver>> nodeIDReceiver) {
-
-    var result = new HashMap<Integer, Map<String, Receiver>>();
-
-    nodeIDReceiver.forEach(
-        (nodeID, receivers) -> {
-          result.put(
-              nodeID,
-              receivers.stream()
-                  .map(
-                      receiver ->
-                          receiver.current().stream()
-                              .map(
-                                  bean ->
-                                      bean.beanObject().getProperties().values().stream()
-                                          .findAny()
-                                          .get())
-                              .findAny()
-                              .get())
-                  .collect(Collectors.toList())
-                  .stream()
-                  .collect(
-                      Collectors.toMap(
-                          Function.identity(),
-                          name ->
-                              receivers.stream()
-                                  .filter(
-                                      mbean ->
-                                          mbean.current().stream()
-                                              .findAny()
-                                              .get()
-                                              .beanObject()
-                                              .getProperties()
-                                              .values()
-                                              .stream()
-                                              .findAny()
-                                              .get()
-                                              .equals(name))
-                                  .findAny()
-                                  .get())));
-        });
-    return result;
+  /** @param brokerIDReceiver nodes metrics that exists in topic */
+  private void metricsNameForReceiver(HashMap<Integer, List<Receiver>> brokerIDReceiver) {
+    brokerIDReceiver.forEach(
+        (brokerID, receivers) ->
+            receivers.stream()
+                .map(
+                    receiver ->
+                        receiver.current().stream()
+                            .map(
+                                bean ->
+                                    bean.beanObject().getProperties().values().stream()
+                                        .findAny()
+                                        .get())
+                            .findAny()
+                            .get())
+                .collect(Collectors.toList())
+                .forEach(
+                    name ->
+                        findBroker(brokerID)
+                            .metrics
+                            .put(
+                                name,
+                                receivers.stream()
+                                    .filter(
+                                        mbean ->
+                                            mbean.current().stream()
+                                                .findAny()
+                                                .get()
+                                                .beanObject()
+                                                .getProperties()
+                                                .values()
+                                                .stream()
+                                                .findAny()
+                                                .get()
+                                                .equals(name))
+                                    .findAny()
+                                    .get())));
   }
 
-  private Map<Integer, Double> brokersMsgSec(int index, Map<Integer, List<Double>> eachMsg) {
-    return eachMsg.entrySet().stream()
-        .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().get(index)));
+  private Broker findBroker(int brokerID) {
+    return brokers.stream().filter(broker -> broker.brokerID == brokerID).findAny().get();
   }
 
-  private double avgMsgSec(int index, Map<Integer, List<Double>> eachMsg) {
-    return avgBrokersMsgPerSec(eachMsg).get(index);
-  }
-
-  private String ipAddress(Map.Entry<Integer, Map.Entry<String, Integer>> hostPort)
-      throws UnknownHostException {
-    var host = "-1.-1.-1.-1";
-    if (!hostPort.getValue().getKey().matches(regex)) {
-      host = InetAddress.getByName(hostPort.getValue().getKey()).toString().split("/")[1];
+  private String ipAddress(String host) {
+    var correctHost = "-1.-1.-1.-1";
+    if (notIPAddress(host)) {
+      try {
+        correctHost = String.valueOf(InetAddress.getByName(host)).split("/")[1];
+      } catch (UnknownHostException e) {
+        e.printStackTrace();
+      }
     } else {
-      host = hostPort.getValue().getKey();
+      correctHost = host;
     }
-    return host;
+    return correctHost;
+  }
+
+  private boolean notIPAddress(String host) {
+    return !host.matches(regex);
+  }
+
+  private boolean overOneSecond() {
+    return lastTime + Duration.ofSeconds(1).toMillis() <= System.currentTimeMillis();
   }
 
   public void close() {
     RECEIVER_FACTORY.close(currentJmxAddresses);
   }
+
+  private List<Integer> brokerIDOfReceiver(String host) {
+    return brokers.stream()
+        .filter(broker -> Objects.equals(ipAddress(broker.host), host))
+        .map(broker -> broker.brokerID)
+        .collect(Collectors.toList());
+  }
+
+  // visible of test
+  int brokerLoad(double brokerSituation, double avg) {
+    var initialization = 5;
+    var loadDeviationRate = (brokerSituation - avg) / avg;
+    var load = (int) Math.round(initialization + loadDeviationRate * 10);
+    if (load > 10) load = 10;
+    else if (load < 0) load = 0;
+    return load;
+  }
+
+  public double thoughPutComparison(int brokerID) {
+    return findBroker(brokerID).thoughPutComparison;
+  }
+
+  public double memoryUsage(int brokerID) {
+    return findBroker(brokerID).memoryUsage();
+  }
+
+  // visible of test
+  static class Broker {
+    private final int brokerID;
+    private final String host;
+    private final int port;
+    private int count = 0;
+    private final Map<Integer, Integer> load = new HashMap<>();
+    private final Map<String, Receiver> metrics = new HashMap<>();
+    private double input;
+    private double output;
+    private double maxThoughPut;
+    private double brokerSituation;
+    private double thoughPutComparison;
+    private double inputNormalized;
+    private double outputNormalized;
+    private double brokerSituationNormalized;
+    private static final double inputWeights = 0.5;
+    private static final double outputWeights = 0.5;
+
+    // visible of test
+    Broker(int brokerID, String host, int port) {
+      this.brokerID = brokerID;
+      this.host = host;
+      this.port = port;
+    }
+
+    private int totalLoad() {
+      return load.values().stream().reduce(0, Integer::sum);
+    }
+
+    private void load(Integer currentLoad) {
+      this.load.put(count % 10, currentLoad);
+      count++;
+    }
+
+    private void brokerSituationNormalized(double totalSituation) {
+      this.brokerSituationNormalized = brokerSituation / totalSituation;
+    }
+
+    private void inputNormalized(double brokersInput) {
+      this.inputNormalized = (input + 1) / (brokersInput + 1);
+    }
+
+    private void outputNormalized(double brokersOutput) {
+      this.outputNormalized = (output + 1) / (brokersOutput + 1);
+    }
+
+    private void brokerSituation() {
+      this.brokerSituation =
+          (inputNormalized * inputWeights + outputNormalized * outputWeights) / thoughPutComparison;
+    }
+
+    private void maxThoughPut() {
+      this.maxThoughPut = Math.max(maxThoughPut, input + output);
+    }
+
+    private void thoughPutComparison(double referenceThoughPut) {
+      if (count >= 20) {
+        this.thoughPutComparison =
+            (maxThoughPut - referenceThoughPut) / (referenceThoughPut + 1) + 1;
+      } else if (count >= 10) {
+        this.thoughPutComparison =
+            (maxThoughPut - referenceThoughPut) / (referenceThoughPut + 1) / 2 + 1;
+      } else if (count >= 0) {
+        this.thoughPutComparison = 1;
+      }
+    }
+
+    private double memoryUsage() {
+      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/ReceiverFactory.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/ReceiverFactory.java
index ac309631..eb230376 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/ReceiverFactory.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/ReceiverFactory.java
@@ -23,7 +23,7 @@ public class ReceiverFactory {
     beanCollector =
         BeanCollector.builder()
             .interval(Duration.ofSeconds(1))
-            .numberOfObjectsPerNode(10)
+            .numberOfObjectsPerNode(1)
             .clientCreator(MBeanClient::jndi)
             .build();
   }
@@ -51,6 +51,13 @@ public class ReceiverFactory {
                       .port(port)
                       .metricsGetter(KafkaMetrics.BrokerTopic.BytesOutPerSec::fetch)
                       .build());
+              receiversList.add(
+                  beanCollector
+                      .register()
+                      .host(host)
+                      .port(port)
+                      .metricsGetter(KafkaMetrics.Host::jvmMemory)
+                      .build());
               count.put(nodeKey(host, port), 1);
             } else {
               count.put(nodeKey(host, port), count.get(nodeKey(host, port)) + 1);
diff --git a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/DependencyClient.java b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/DependencyClient.java
index 777dfc7c..5667d1a7 100644
--- a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/DependencyClient.java
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/DependencyClient.java
@@ -1,6 +1,14 @@
 package org.astraea.partitioner.smoothPartitioner;
 
-// TODO
+/**
+ * The interface for the {@link SmoothWeightPartitioner}
+ *
+ * @see SmoothWeightPartitioner Enable users to send dependent data.
+ */
 public interface DependencyClient {
+  void startDependency();
+
   void finishDependency();
+
+  String state();
 }
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 4942327a..287af45e 100644
--- a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
@@ -27,7 +27,8 @@ import org.astraea.partitioner.nodeLoadMetric.NodeLoadClient;
  * <pre>{@code
  * KafkaProducer producer = new KafkaProducer(props);
  *
- * var dependencyControl = SmoothWeightPartitioner.beginDependency(producer);
+ * var dependencyControl = SmoothWeightPartitioner.dependencyClient(producer);
+ * dependencyControl.startDependency();
  * try{
  *     producer.send();
  *     producer.send();
@@ -43,10 +44,10 @@ public class SmoothWeightPartitioner implements Partitioner, DependencyClient {
    * Record the current weight of each node according to Poisson calculation and the weight after
    * partitioner calculation.
    */
-  private Map<Integer, int[]> brokerHashMap = new HashMap<>();
+  private final Map<Integer, int[]> brokersWeight = new HashMap<>();
 
   private NodeLoadClient nodeLoadClient;
-  private static final DependencyManager dependencyManager = new DependencyManager();
+  private DependencyManager dependencyManager;
   /**
    * This parameter only works in dependency mode.And it will be specified when the dependency is
    * executed for the first time, and no changes will be made afterwards.
@@ -56,51 +57,55 @@ public class SmoothWeightPartitioner implements Partitioner, DependencyClient {
   @Override
   public int partition(
       String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
-    Map<Integer, Integer> overLoadCount;
+    Map<Integer, Integer> loadCount;
     var rand = new Random();
     switch (dependencyManager.currentState) {
-      case Start_Dependency:
-        overLoadCount = overLoadCount(cluster);
-        Objects.requireNonNull(overLoadCount, "OverLoadCount should not be null.");
-        var minOverLoadNode =
-            overLoadCount.entrySet().stream().min(Map.Entry.comparingByValue()).get();
+      case START_DEPENDENCY:
+        loadCount = loadCount(cluster);
+        Objects.requireNonNull(loadCount, "OverLoadCount should not be null.");
+        var minOverLoadNode = loadCount.entrySet().stream().min(Map.Entry.comparingByValue()).get();
         var minOverLoadPartition = cluster.partitionsForNode(minOverLoadNode.getKey());
         targetPartition =
             minOverLoadPartition.get(rand.nextInt(minOverLoadPartition.size())).partition();
-        dependencyManager.transitionTo(DependencyManager.State.IN_Dependency);
+        dependencyManager.transitionTo(DependencyManager.State.IN_DEPENDENCY);
         return targetPartition;
-      case IN_Dependency:
+      case IN_DEPENDENCY:
         return targetPartition;
       default:
-        overLoadCount = overLoadCount(cluster);
+        loadCount = loadCount(cluster);
         var loadPoisson = new LoadPoisson();
 
-        Objects.requireNonNull(overLoadCount, "OverLoadCount should not be null.");
-        brokerHashMap(loadPoisson.allPoisson(overLoadCount));
+        Objects.requireNonNull(loadCount, "OverLoadCount should not be null.");
+        brokersWeight(loadPoisson.allPoisson(loadCount));
         AtomicReference<Map.Entry<Integer, int[]>> maxWeightServer = new AtomicReference<>();
         var allWeight = allNodesWeight();
-        var currentBrokerHashMap = brokerHashMap();
-        currentBrokerHashMap.forEach(
+        brokersWeight.forEach(
             (nodeID, weight) -> {
               if (maxWeightServer.get() == null
                   || weight[1] > maxWeightServer.get().getValue()[1]) {
                 maxWeightServer.set(Map.entry(nodeID, weight));
               }
-              currentBrokerHashMap.put(nodeID, new int[] {weight[0], weight[1] + weight[0]});
+              brokersWeight.put(nodeID, new int[] {weight[0], weight[1] + weight[0]});
             });
         Objects.requireNonNull(maxWeightServer.get(), "MaxWeightServer should not be null.");
-        currentBrokerHashMap.put(
+        brokersWeight.put(
             maxWeightServer.get().getKey(),
             new int[] {
               maxWeightServer.get().getValue()[0], maxWeightServer.get().getValue()[1] - allWeight
             });
-        currentBrokerHashMap(currentBrokerHashMap);
+        brokersWeight
+            .keySet()
+            .forEach(
+                brokerID -> {
+                  if (memoryWarning(brokerID)) {
+                    subBrokerWeight(brokerID, 100);
+                  }
+                });
 
         ArrayList<Integer> partitionList = new ArrayList<>();
         cluster
             .partitionsForNode(maxWeightServer.get().getKey())
             .forEach(partitionInfo -> partitionList.add(partitionInfo.partition()));
-
         return partitionList.get(rand.nextInt(partitionList.size()));
     }
   }
@@ -126,6 +131,7 @@ public class SmoothWeightPartitioner implements Partitioner, DependencyClient {
       Objects.requireNonNull(mapAddress, "You must configure jmx_servers correctly.");
 
       nodeLoadClient = new NodeLoadClient(mapAddress);
+      dependencyManager = new DependencyManager();
     } catch (IOException e) {
       throw new RuntimeException();
     }
@@ -133,38 +139,37 @@ public class SmoothWeightPartitioner implements Partitioner, DependencyClient {
 
   /** Change the weight of the node according to the current Poisson. */
   // visible for test
-  public synchronized void brokerHashMap(Map<Integer, Double> poissonMap) {
+  public synchronized void brokersWeight(Map<Integer, Double> poissonMap) {
     poissonMap.forEach(
         (key, value) -> {
-          if (!brokerHashMap.containsKey(key)) {
-            brokerHashMap.put(key, new int[] {(int) ((1 - value) * 20), 0});
+          var thoughPutAbility = nodeLoadClient.thoughPutComparison(key);
+          if (!brokersWeight.containsKey(key)) {
+            brokersWeight.put(key, new int[] {(int) ((1 - value) * 20 * thoughPutAbility), 0});
           } else {
-            brokerHashMap.put(key, new int[] {(int) ((1 - value) * 20), brokerHashMap.get(key)[1]});
+            brokersWeight.put(
+                key,
+                new int[] {(int) ((1 - value) * 20 * thoughPutAbility), brokersWeight.get(key)[1]});
           }
         });
   }
 
-  private synchronized int allNodesWeight() {
-    return brokerHashMap.values().stream().mapToInt(vs -> vs[0]).sum();
-  }
-
-  // visible for test
-  public synchronized Map<Integer, int[]> brokerHashMap() {
-    return brokerHashMap;
+  private void subBrokerWeight(int brokerID, int subNumber) {
+    brokersWeight.put(
+        brokerID,
+        new int[] {brokersWeight.get(brokerID)[0], brokersWeight.get(brokerID)[1] - subNumber});
   }
 
-  private synchronized void currentBrokerHashMap(Map<Integer, int[]> currentBrokerHashMap) {
-    brokerHashMap = currentBrokerHashMap;
+  private synchronized int allNodesWeight() {
+    return brokersWeight.values().stream().mapToInt(vs -> vs[0]).sum();
   }
 
-  // visible for test
-  public void brokerHashMapValue(Integer x, int y) {
-    brokerHashMap.put(x, new int[] {0, y});
+  public static DependencyClient dependencyClient(KafkaProducer<?, ?> producer) {
+    return (DependencyClient) Utils.requireField(producer, "partitioner");
   }
 
-  public static synchronized DependencyClient beginDependency(KafkaProducer<?, ?> producer) {
-    dependencyManager.beginDependency();
-    return (DependencyClient) Utils.requireField(producer, "partitioner");
+  @Override
+  public synchronized void startDependency() {
+    dependencyManager.startDependency();
   }
 
   @Override
@@ -172,22 +177,31 @@ public class SmoothWeightPartitioner implements Partitioner, DependencyClient {
     dependencyManager.finishDependency();
   }
 
-  private Map<Integer, Integer> overLoadCount(Cluster cluster) {
+  private boolean memoryWarning(int brokerID) {
+    return nodeLoadClient.memoryUsage(brokerID) >= 0.8;
+  }
+
+  private Map<Integer, Integer> loadCount(Cluster cluster) {
     try {
-      return nodeLoadClient.nodesOverLoad(cluster);
+      return nodeLoadClient.loadSituation(cluster);
     } catch (UnknownHostException e) {
       throw new IllegalArgumentException(e);
     }
   }
 
+  // visible of testing
+  public String state() {
+    return dependencyManager.currentState.toString();
+  }
+
   /**
    * Store the state related to the partitioner dependency and avoid wrong transitions between them.
    */
   private static class DependencyManager {
     private volatile State currentState = State.UNINITIALIZED;
 
-    private synchronized void beginDependency() {
-      transitionTo(State.Start_Dependency);
+    private synchronized void startDependency() {
+      transitionTo(State.START_DEPENDENCY);
     }
 
     private synchronized void finishDependency() {
@@ -196,19 +210,19 @@ public class SmoothWeightPartitioner implements Partitioner, DependencyClient {
 
     private enum State {
       UNINITIALIZED,
-      Start_Dependency,
-      IN_Dependency,
+      START_DEPENDENCY,
+      IN_DEPENDENCY,
       FATAL_ERROR;
 
       private boolean isTransitionValid(
           DependencyManager.State source, DependencyManager.State target) {
         switch (target) {
           case UNINITIALIZED:
-            return source == Start_Dependency || source == IN_Dependency;
-          case Start_Dependency:
+            return source == START_DEPENDENCY || source == IN_DEPENDENCY;
+          case START_DEPENDENCY:
             return source == UNINITIALIZED;
-          case IN_Dependency:
-            return source == Start_Dependency;
+          case IN_DEPENDENCY:
+            return source == START_DEPENDENCY;
           default:
             return true;
         }
@@ -216,12 +230,16 @@ public class SmoothWeightPartitioner implements Partitioner, DependencyClient {
     }
 
     private void transitionTo(State target) {
-      if (!currentState.isTransitionValid(currentState, target)) {
-        throw new KafkaException(
-            "Invalid transition attempted from state "
-                + currentState.name()
-                + " to state "
-                + target.name());
+      try {
+        if (!currentState.isTransitionValid(currentState, target)) {
+          throw new KafkaException(
+              "Invalid transition attempted from state "
+                  + currentState.name()
+                  + " to state "
+                  + target.name());
+        }
+      } catch (KafkaException e) {
+        // TODO
       }
       currentState = target;
     }
diff --git a/docker/start_broker.sh b/docker/start_broker.sh
index 05aa1e59..80af7463 100755
--- a/docker/start_broker.sh
+++ b/docker/start_broker.sh
@@ -91,9 +91,6 @@ RUN apt-get update && apt-get upgrade -y && DEBIAN_FRONTEND=noninteractive apt-g
 # add user
 RUN groupadd $USER && useradd -ms /bin/bash -g $USER $USER
 
-# change user
-USER $USER
-
 # download jmx exporter
 RUN mkdir /tmp/jmx_exporter
 WORKDIR /tmp/jmx_exporter
@@ -105,9 +102,16 @@ RUN git clone https://github.com/apache/kafka /tmp/kafka
 WORKDIR /tmp/kafka
 RUN git checkout $VERSION
 RUN ./gradlew clean releaseTarGz
-RUN mkdir /home/$USER/kafka
-RUN tar -zxvf \$(find ./core/build/distributions/ -maxdepth 1 -type f -name kafka_*SNAPSHOT.tgz) -C /home/$USER/kafka --strip-components=1
-WORKDIR "/home/$USER/kafka"
+RUN mkdir /opt/kafka
+RUN tar -zxvf \$(find ./core/build/distributions/ -maxdepth 1 -type f -name kafka_*SNAPSHOT.tgz) -C /opt/kafka --strip-components=1
+WORKDIR /opt/kafka
+
+# export ENV
+ENV KAFKA_HOME /opt/kafka
+
+# change user
+RUN chown -R $USER:$USER /opt/kafka
+USER $USER
 " >"$DOCKERFILE"
 }
 
@@ -121,9 +125,6 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y openjdk-11-jdk wg
 # add user
 RUN groupadd $USER && useradd -ms /bin/bash -g $USER $USER
 
-# change user
-USER $USER
-
 # download jmx exporter
 RUN mkdir /tmp/jmx_exporter
 WORKDIR /tmp/jmx_exporter
@@ -133,9 +134,16 @@ RUN wget https://REPO1.maven.org/maven2/io/prometheus/jmx/jmx_prometheus_javaage
 # download kafka
 WORKDIR /tmp
 RUN wget https://archive.apache.org/dist/kafka/${VERSION}/kafka_2.13-${VERSION}.tgz
-RUN mkdir /home/$USER/kafka
-RUN tar -zxvf kafka_2.13-${VERSION}.tgz -C /home/$USER/kafka --strip-components=1
-WORKDIR "/home/$USER/kafka"
+RUN mkdir /opt/kafka
+RUN tar -zxvf kafka_2.13-${VERSION}.tgz -C /opt/kafka --strip-components=1
+WORKDIR /opt/kafka
+
+# export ENV
+ENV KAFKA_HOME /opt/kafka
+
+# change user
+RUN chown -R $USER:$USER /opt/kafka
+USER $USER
 " >"$DOCKERFILE"
 }
 
@@ -287,7 +295,7 @@ setPropertyIfEmpty "offsets.topic.replication.factor" "1"
 setPropertyIfEmpty "transaction.state.log.min.isr" "1"
 setLogDirs
 
-docker run -d \
+docker run -d --init \
   --name $CONTAINER_NAME \
   -e KAFKA_HEAP_OPTS="$HEAP_OPTS" \
   -e KAFKA_JMX_OPTS="$JMX_OPTS" \
diff --git a/docker/start_kafka_tool.sh b/docker/start_kafka_tool.sh
index d2afe920..b51a116d 100755
--- a/docker/start_kafka_tool.sh
+++ b/docker/start_kafka_tool.sh
@@ -72,7 +72,7 @@ function buildImageIfNeed() {
 
 function runContainer() {
   local args=$1
-  docker run --rm \
+  docker run --rm --init \
     $IMAGE_NAME \
     /bin/bash -c "java -jar /tmp/app.jar $args"
 }
diff --git a/docker/start_spark.sh b/docker/start_spark.sh
index 7ae25739..83a205d0 100755
--- a/docker/start_spark.sh
+++ b/docker/start_spark.sh
@@ -10,11 +10,6 @@ declare -r BUILD=${BUILD:-false}
 declare -r RUN=${RUN:-true}
 declare -r SPARK_PORT=${SPARK_PORT:-$(($(($RANDOM % 10000)) + 10000))}
 declare -r SPARK_UI_PORT=${SPARK_UI_PORT:-$(($(($RANDOM % 10000)) + 10000))}
-declare -r IVY_VERSION=2.5.0
-declare -r DELTA_VERSION=${DELTA_VERSION:-1.0.0}
-declare -r PYTHON_KAFKA_VERSION=${PYTHON_KAFKA_VERSION:-1.7.0}
-# hardcode hadoop version to avoid NPE (see https://issues.apache.org/jira/browse/HADOOP-16410)
-declare -r HADOOP_VERSION=3.2.2
 declare -r DOCKER_FOLDER=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
 declare -r DOCKERFILE=$DOCKER_FOLDER/spark.dockerfile
 declare -r ADDRESS=$([[ "$(which ipconfig)" != "" ]] && ipconfig getifaddr en0 || hostname -i)
@@ -76,33 +71,25 @@ FROM ubuntu:20.04
 ENV DEBIAN_FRONTEND noninteractive
 
 # install tools
-RUN apt-get update && apt-get upgrade -y && apt-get install -y openjdk-11-jdk wget python3 python3-pip unzip
+RUN apt-get update && apt-get upgrade -y && apt-get install -y openjdk-11-jdk wget python3 python3-pip unzip tini
 
 # add user
 RUN groupadd $USER && useradd -ms /bin/bash -g $USER $USER
 
-# change user
-USER $USER
-
-# install python dependencies
-RUN pip3 install confluent-kafka==$PYTHON_KAFKA_VERSION delta-spark==$DELTA_VERSION pyspark==$VERSION
-
-# install java dependencies
-WORKDIR /tmp
-RUN wget https://dlcdn.apache.org//ant/ivy/${IVY_VERSION}/apache-ivy-${IVY_VERSION}-bin.zip
-RUN unzip apache-ivy-${IVY_VERSION}-bin.zip
-WORKDIR apache-ivy-${IVY_VERSION}
-RUN java -jar ./ivy-${IVY_VERSION}.jar -dependency io.delta delta-core_2.12 $DELTA_VERSION
-RUN java -jar ./ivy-${IVY_VERSION}.jar -dependency org.apache.spark spark-sql-kafka-0-10_2.12 $VERSION
-RUN java -jar ./ivy-${IVY_VERSION}.jar -dependency org.apache.spark spark-token-provider-kafka-0-10_2.12 $VERSION
-RUN java -jar ./ivy-${IVY_VERSION}.jar -dependency org.apache.hadoop hadoop-azure $HADOOP_VERSION
-
 # download spark
 WORKDIR /tmp
 RUN wget https://archive.apache.org/dist/spark/spark-${VERSION}/spark-${VERSION}-bin-hadoop3.2.tgz
-RUN mkdir /home/$USER/spark
-RUN tar -zxvf spark-${VERSION}-bin-hadoop3.2.tgz -C /home/$USER/spark --strip-components=1
-WORKDIR /home/$USER/spark
+RUN mkdir /opt/spark
+RUN tar -zxvf spark-${VERSION}-bin-hadoop3.2.tgz -C /opt/spark --strip-components=1
+
+# export ENV
+ENV SPARK_HOME /opt/spark
+
+# change user
+RUN chown -R $USER:$USER /opt/spark
+USER $USER
+
+WORKDIR /opt/spark
 " >"$DOCKERFILE"
 }
 
@@ -122,18 +109,23 @@ RUN pip3 install confluent-kafka==$PYTHON_KAFKA_VERSION
 # add user
 RUN groupadd $USER && useradd -ms /bin/bash -g $USER $USER
 
-# change user
-USER $USER
-
 # build spark from source code
 RUN git clone https://github.com/apache/spark /tmp/spark
 WORKDIR /tmp/spark
 RUN git checkout $VERSION
 RUN ./dev/make-distribution.sh --pip --tgz
-RUN mkdir /home/$USER/spark
-RUN tar -zxvf \$(find ./ -maxdepth 1 -type f -name spark-*SNAPSHOT*.tgz) -C /home/$USER/spark --strip-components=1
+RUN mkdir /opt/spark
+RUN tar -zxvf \$(find ./ -maxdepth 1 -type f -name spark-*SNAPSHOT*.tgz) -C /opt/spark --strip-components=1
 RUN ./build/mvn install -DskipTests
-WORKDIR "/home/$USER/spark"
+
+# export ENV
+ENV SPARK_HOME /opt/spark
+
+# change user
+RUN chown -R $USER:$USER /opt/spark
+USER $USER
+
+WORKDIR /opt/spark
 " >"$DOCKERFILE"
 }
 
@@ -189,7 +181,7 @@ checkOs
 
 if [[ -n "$master_url" ]]; then
   checkConflictContainer $WORKER_NAME "worker"
-  docker run -d \
+  docker run -d --init \
     -e SPARK_WORKER_WEBUI_PORT=$SPARK_UI_PORT \
     -e SPARK_WORKER_PORT=$SPARK_PORT \
     -e SPARK_NO_DAEMONIZE=true \
@@ -203,7 +195,7 @@ if [[ -n "$master_url" ]]; then
   echo "================================================="
 else
   checkConflictContainer $MASTER_NAME "master"
-  docker run -d \
+  docker run -d --init \
     -e SPARK_MASTER_WEBUI_PORT=$SPARK_UI_PORT \
     -e SPARK_MASTER_PORT=$SPARK_PORT \
     -e SPARK_NO_DAEMONIZE=true \
diff --git a/docker/start_zookeeper.sh b/docker/start_zookeeper.sh
index 37271c57..5e5e0eaf 100755
--- a/docker/start_zookeeper.sh
+++ b/docker/start_zookeeper.sh
@@ -51,20 +51,24 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y openjdk-11-jdk wg
 # add user
 RUN groupadd $USER && useradd -ms /bin/bash -g $USER $USER
 
-# change user
-USER $USER
-
 # download zookeeper
 WORKDIR /tmp
 RUN wget https://archive.apache.org/dist/zookeeper/zookeeper-${VERSION}/apache-zookeeper-${VERSION}-bin.tar.gz
-RUN mkdir /home/$USER/zookeeper
-RUN tar -zxvf apache-zookeeper-${VERSION}-bin.tar.gz -C /home/$USER/zookeeper --strip-components=1
-WORKDIR /home/$USER/zookeeper
+RUN mkdir /opt/zookeeper
+RUN tar -zxvf apache-zookeeper-${VERSION}-bin.tar.gz -C /opt/zookeeper --strip-components=1
+WORKDIR /opt/zookeeper
 
 # create config file
 RUN echo "tickTime=2000" >> ./conf/zoo.cfg
 RUN echo "dataDir=$DATA_FOLDER_IN_CONTAINER" >> ./conf/zoo.cfg
 RUN echo "clientPort=2181" >> ./conf/zoo.cfg
+
+# export ENV
+ENV ZOOKEEPER_HOME /opt/zookeeper
+
+# change user
+RUN chown -R $USER:$USER /opt/zookeeper
+USER $USER
 " >"$DOCKERFILE"
 }
 
@@ -109,13 +113,13 @@ checkNetwork
 if [[ -n "$DATA_FOLDER" ]]; then
   mkdir -p "$DATA_FOLDER"
   echo "mount $DATA_FOLDER to container"
-  docker run -d \
+  docker run -d --init \
     --name $CONTAINER_NAME \
     -p $ZOOKEEPER_PORT:2181 \
     -v "$DATA_FOLDER":$DATA_FOLDER_IN_CONTAINER \
     $IMAGE_NAME ./bin/zkServer.sh start-foreground
 else
-  docker run -d \
+  docker run -d --init \
     -p $ZOOKEEPER_PORT:2181 \
     $IMAGE_NAME ./bin/zkServer.sh start-foreground
 fi


Test patch:
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/DependencyPartitionTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/DependencyPartitionTest.java
index ea2c5fb2..1e8ac86d 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/DependencyPartitionTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/DependencyPartitionTest.java
@@ -46,7 +46,8 @@ public class DependencyPartitionTest extends RequireBrokerCluster {
         ps.entrySet().stream()
             .collect(Collectors.toMap(e -> e.getKey().toString(), Map.Entry::getValue));
     try (var producer = Producer.builder().configs(props).build()) {
-      var dependencyClient = SmoothWeightPartitioner.beginDependency(producer.kafkaProducer());
+      var dependencyClient = SmoothWeightPartitioner.dependencyClient(producer.kafkaProducer());
+      dependencyClient.startDependency();
       var targetPartition = 0;
       var i = 0;
       while (i < 20) {
@@ -86,4 +87,18 @@ public class DependencyPartitionTest extends RequireBrokerCluster {
       Assertions.assertArrayEquals(header.value(), actualHeader.value());
     }
   }
+
+  @Test
+  void testDependencyStateChange() {
+    var props =
+        initProConfig().entrySet().stream()
+            .collect(Collectors.toMap(e -> e.getKey().toString(), Map.Entry::getValue));
+    var producer1 = Producer.builder().configs(props).build();
+    var producer2 = Producer.builder().configs(props).build();
+    var client1 = SmoothWeightPartitioner.dependencyClient(producer1.kafkaProducer());
+    var client2 = SmoothWeightPartitioner.dependencyClient(producer2.kafkaProducer());
+    client1.startDependency();
+    Assertions.assertEquals(client1.state(), "START_DEPENDENCY");
+    Assertions.assertEquals(client2.state(), "UNINITIALIZED");
+  }
 }
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 85c8692c..8fc7e8b2 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
@@ -1,31 +1,31 @@
 package org.astraea.partitioner.nodeLoadMetric;
 
+import static org.mockito.Mockito.when;
+
 import java.io.IOException;
-import java.util.Arrays;
+import java.lang.reflect.Field;
+import java.net.UnknownHostException;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.Properties;
-import org.apache.kafka.clients.producer.ProducerConfig;
-import org.apache.kafka.common.serialization.ByteArraySerializer;
-import org.astraea.partitioner.smoothPartitioner.SmoothWeightPartitioner;
+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.service.RequireBrokerCluster;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestInstance;
+import org.mockito.Mockito;
 
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 public class NodeLoadClientTest extends RequireBrokerCluster {
-  private final String brokerList = bootstrapServers();
   private NodeLoadClient nodeLoadClient;
-  private String jmxAddress;
-  private Map props;
 
   @BeforeAll
   void setUp() throws IOException {
-    props = initProConfig();
     var map = new HashMap<String, Integer>();
     map.put(jmxServiceURL().getHost(), jmxServiceURL().getPort());
     nodeLoadClient = new NodeLoadClient(map);
@@ -36,25 +36,27 @@ public class NodeLoadClientTest extends RequireBrokerCluster {
     nodeLoadClient.close();
   }
 
-  Properties initProConfig() {
-    Properties props = new Properties();
-    jmxAddress = jmxServiceURL().getHost() + ":" + jmxServiceURL().getPort();
-    props.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList);
-    props.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
-    props.put(ProducerConfig.CLIENT_ID_CONFIG, "id1");
-    props.put(ProducerConfig.PARTITIONER_CLASS_CONFIG, SmoothWeightPartitioner.class.getName());
-    props.put("jmx_servers", jmxAddress);
-    return props;
+  private Field field(Object object, String fieldName) {
+    Field field = null;
+    try {
+      field = object.getClass().getDeclaredField(fieldName);
+      field.setAccessible(true);
+    } catch (NoSuchFieldException e) {
+      e.printStackTrace();
+    }
+    return field;
   }
 
   @Test
-  void testAvgBrokersMsgPerSec() {
-    var testMap = new HashMap<Integer, List<Double>>();
-    testMap.put(0, Arrays.asList(20.0, 30.0, 40.0));
-    testMap.put(1, Arrays.asList(30.0, 20.0, 30.0));
-    testMap.put(2, Arrays.asList(40.0, 40.0, 20.0));
-    Assertions.assertEquals(
-        nodeLoadClient.avgBrokersMsgPerSec(testMap), Arrays.asList(30.0, 30.0, 30.0));
+  void testBrokerLoad() {
+    List<NodeLoadClient.Broker> brokers = new ArrayList<>();
+    brokers.add(new NodeLoadClient.Broker(0, "0.0.0.0", 111));
+    brokers.add(new NodeLoadClient.Broker(1, "0.0.0.0", 222));
+    brokers.add(new NodeLoadClient.Broker(2, "0.0.0.0", 333));
+    setBrokers(brokers);
+    Assertions.assertEquals(nodeLoadClient.brokerLoad(0.37, 1.0 / 3), 6);
+    Assertions.assertEquals(nodeLoadClient.brokerLoad(0.01, 1.0 / 3), 0);
+    Assertions.assertEquals(nodeLoadClient.brokerLoad(0.8, 1.0 / 3), 10);
   }
 
   @Test
@@ -65,7 +67,73 @@ public class NodeLoadClientTest extends RequireBrokerCluster {
     testMap.put(2, 25.0);
     testMap.put(3, 20.0);
     testMap.put(4, 20.0);
-    Assertions.assertEquals(
-        nodeLoadClient.standardDeviationImperative(testMap, 20), 3.1622776601683795);
+
+    List<NodeLoadClient.Broker> brokers = new ArrayList<>();
+    brokers.add(0, setSituationNormalized(0, "0.0.0.0", 111, 15.0));
+    brokers.add(1, setSituationNormalized(0, "0.0.0.0", 222, 20.0));
+    brokers.add(2, setSituationNormalized(0, "0.0.0.0", 333, 25.0));
+    brokers.add(3, setSituationNormalized(0, "0.0.0.0", 444, 20.0));
+    brokers.add(4, setSituationNormalized(0, "0.0.0.0", 555, 20.0));
+    setBrokers(brokers);
+
+    Assertions.assertEquals(nodeLoadClient.standardDeviationImperative(20), 3.1622776601683795);
+  }
+
+  @Test
+  void testLoadSituation() throws UnknownHostException {
+    var bootstrapServers = List.of(bootstrapServers().split(","));
+    List<Node> 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])));
+          count.getAndIncrement();
+        });
+
+    Cluster cluster = Mockito.mock(Cluster.class);
+    when(cluster.nodes()).thenReturn(nodes);
+    var load = nodeLoadClient.loadSituation(cluster);
+    Assertions.assertEquals(load.get(0), 5);
+    Assertions.assertEquals(load.get(1), 5);
+    Assertions.assertEquals(load.get(2), 5);
+    load = nodeLoadClient.loadSituation(cluster);
+    Assertions.assertEquals(load.get(0), 5);
+    Assertions.assertEquals(load.get(1), 5);
+    Assertions.assertEquals(load.get(2), 5);
+    sleep(1);
+    load = nodeLoadClient.loadSituation(cluster);
+    Assertions.assertEquals(load.get(0), 10);
+    Assertions.assertEquals(load.get(1), 10);
+    Assertions.assertEquals(load.get(2), 10);
+  }
+
+  private NodeLoadClient.Broker setSituationNormalized(
+      int brokerID, String host, int port, double situation) {
+    var broker = new NodeLoadClient.Broker(brokerID, host, port);
+    var brokerSituation = field(broker, "brokerSituationNormalized");
+    try {
+      brokerSituation.set(broker, situation);
+    } catch (IllegalAccessException e) {
+      e.printStackTrace();
+    }
+    return broker;
+  }
+
+  private void setBrokers(List<NodeLoadClient.Broker> brokers) {
+    var brokersField = field(nodeLoadClient, "brokers");
+    try {
+      brokersField.set(nodeLoadClient, brokers);
+    } catch (IllegalAccessException e) {
+      e.printStackTrace();
+    }
+  }
+
+  private static void sleep(int seconds) {
+    try {
+      TimeUnit.SECONDS.sleep(seconds);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
   }
 }
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerTest.java
index 3869d56c..6bbbca23 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerTest.java
@@ -1,7 +1,11 @@
 package org.astraea.partitioner.nodeLoadMetric;
 
+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 java.lang.reflect.Field;
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.List;
@@ -9,9 +13,12 @@ import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.astraea.concurrent.ThreadPool;
 import org.astraea.consumer.Consumer;
 import org.astraea.consumer.Deserializer;
 import org.astraea.consumer.Header;
@@ -22,11 +29,13 @@ import org.astraea.service.RequireBrokerCluster;
 import org.astraea.topic.TopicAdmin;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInstance;
+import org.mockito.Mockito;
 
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
 public class PartitionerTest extends RequireBrokerCluster {
   private final String brokerList = bootstrapServers();
   TopicAdmin admin = TopicAdmin.of(bootstrapServers());
-  private final String topicName = "address";
 
   private Properties initProConfig() {
     Properties props = new Properties();
@@ -40,8 +49,31 @@ public class PartitionerTest extends RequireBrokerCluster {
     return props;
   }
 
+  @Test
+  void testSetBrokerHashMap() {
+    var nodeLoadClient = Mockito.mock(NodeLoadClient.class);
+    when(nodeLoadClient.thoughPutComparison(anyInt())).thenReturn(1.0);
+    var poissonMap = new HashMap<Integer, Double>();
+    poissonMap.put(0, 0.5);
+    poissonMap.put(1, 0.8);
+    poissonMap.put(2, 0.3);
+
+    var smoothWeightPartitioner = new SmoothWeightPartitioner();
+    setNodeLoadClient(nodeLoadClient, smoothWeightPartitioner);
+    smoothWeightPartitioner.brokersWeight(poissonMap);
+
+    var brokerWeight = (Map<Integer, int[]>) requireField(smoothWeightPartitioner, "brokersWeight");
+    assertEquals(brokerWeight.get(0)[0], 10);
+    assertEquals(brokerWeight.get(1)[0], 3);
+
+    brokerWeight.put(0, new int[] {0, 8});
+    smoothWeightPartitioner.brokersWeight(poissonMap);
+    assertEquals(brokerWeight.get(0)[1], 8);
+  }
+
   @Test
   void testPartitioner() {
+    var topicName = "address";
     admin.creator().topic(topicName).numberOfPartitions(10).create();
     var key = "tainan";
     var timestamp = System.currentTimeMillis() + 10;
@@ -54,7 +86,7 @@ public class PartitionerTest extends RequireBrokerCluster {
                     .collect(Collectors.toMap(e -> e.getKey().toString(), Map.Entry::getValue)))
             .build()) {
       var i = 0;
-      while (i < 20) {
+      while (i < 300) {
         var metadata =
             producer
                 .sender()
@@ -72,7 +104,7 @@ public class PartitionerTest extends RequireBrokerCluster {
     } catch (ExecutionException | InterruptedException e) {
       e.printStackTrace();
     }
-
+    sleep(1);
     try (var consumer =
         Consumer.builder()
             .brokers(bootstrapServers())
@@ -80,8 +112,12 @@ public class PartitionerTest extends RequireBrokerCluster {
             .topics(Set.of(topicName))
             .keyDeserializer(Deserializer.STRING)
             .build()) {
-      var records = consumer.poll(Duration.ofSeconds(10));
-      assertEquals(20, records.size());
+      var records = consumer.poll(Duration.ofSeconds(20));
+      var recordsCount = records.size();
+      while (recordsCount < 300) {
+        recordsCount += consumer.poll(Duration.ofSeconds(20)).size();
+      }
+      assertEquals(300, recordsCount);
       var record = records.iterator().next();
       assertEquals(topicName, record.topic());
       assertEquals("tainan", record.key());
@@ -93,21 +129,117 @@ public class PartitionerTest extends RequireBrokerCluster {
   }
 
   @Test
-  void testSetBrokerHashMap() {
-    var poissonMap = new HashMap<Integer, Double>();
-    poissonMap.put(0, 0.5);
-    poissonMap.put(1, 0.8);
-    poissonMap.put(2, 0.3);
+  void testMultipleProducer() {
+    var topicName = "addr";
+    admin.creator().topic(topicName).numberOfPartitions(10).create();
+    var key = "tainan";
+    var timestamp = System.currentTimeMillis() + 10;
+    var header = Header.of("a", "b".getBytes());
+    try (ThreadPool threadPool =
+        ThreadPool.builder()
+            .executors(
+                IntStream.range(0, 10)
+                    .mapToObj(
+                        i ->
+                            producerExecutor(
+                                Producer.builder()
+                                    .keySerializer(Serializer.STRING)
+                                    .configs(
+                                        initProConfig().entrySet().stream()
+                                            .collect(
+                                                Collectors.toMap(
+                                                    e -> e.getKey().toString(),
+                                                    Map.Entry::getValue)))
+                                    .build(),
+                                topicName,
+                                key,
+                                header,
+                                timestamp))
+                    .collect(Collectors.toUnmodifiableList()))
+            .build()) {
+      threadPool.waitAll();
+    }
+    try (var consumer =
+        Consumer.builder()
+            .brokers(bootstrapServers())
+            .fromBeginning()
+            .topics(Set.of(topicName))
+            .keyDeserializer(Deserializer.STRING)
+            .build()) {
+      var records = consumer.poll(Duration.ofSeconds(20));
+      var recordsCount = records.size();
+      while (recordsCount < 1000) {
+        recordsCount += consumer.poll(Duration.ofSeconds(20)).size();
+      }
+      assertEquals(1000, recordsCount);
+      var record = records.iterator().next();
+      assertEquals(topicName, record.topic());
+      assertEquals("tainan", record.key());
+      assertEquals(1, record.headers().size());
+      var actualHeader = record.headers().iterator().next();
+      assertEquals(header.key(), actualHeader.key());
+      Assertions.assertArrayEquals(header.value(), actualHeader.value());
+    }
+  }
 
-    var smoothWeightPartitioner = new SmoothWeightPartitioner();
+  private ThreadPool.Executor producerExecutor(
+      Producer<String, byte[]> producer, String topic, String key, Header header, long timeStamp) {
+    return new ThreadPool.Executor() {
+      int i = 0;
+
+      @Override
+      public State execute() throws InterruptedException {
+        if (i > 99) return State.DONE;
+        try {
+          producer
+              .sender()
+              .topic(topic)
+              .key(key)
+              .timestamp(timeStamp)
+              .headers(List.of(header))
+              .run()
+              .toCompletableFuture()
+              .get();
+        } catch (ExecutionException e) {
+          e.printStackTrace();
+        }
+        i++;
+        return State.RUNNING;
+      }
+
+      @Override
+      public void close() {
+        producer.close();
+      }
+    };
+  }
 
-    smoothWeightPartitioner.brokerHashMap(poissonMap);
+  private void setNodeLoadClient(
+      NodeLoadClient nodeLoadClient, SmoothWeightPartitioner partitioner) {
+    var target = field(partitioner, "nodeLoadClient");
+    try {
+      target.set(partitioner, nodeLoadClient);
+    } catch (IllegalAccessException e) {
+      e.printStackTrace();
+    }
+  }
 
-    assertEquals(smoothWeightPartitioner.brokerHashMap().get(0)[0], 10);
-    assertEquals(smoothWeightPartitioner.brokerHashMap().get(1)[0], 3);
+  private Field field(Object object, String fieldName) {
+    Field field = null;
+    try {
+      field = object.getClass().getDeclaredField(fieldName);
+      field.setAccessible(true);
+    } catch (NoSuchFieldException e) {
+      e.printStackTrace();
+    }
+    return field;
+  }
 
-    smoothWeightPartitioner.brokerHashMapValue(0, 8);
-    smoothWeightPartitioner.brokerHashMap(poissonMap);
-    assertEquals(smoothWeightPartitioner.brokerHashMap().get(0)[1], 8);
+  private static void sleep(int seconds) {
+    try {
+      TimeUnit.SECONDS.sleep(seconds);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
   }
 }