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

Improve Smooth weight algorithm
我梳理了一遍整個演算發，確保它每一步都有按預期的來。修復和改善了許多問題。有一些問題諸如0.20*20=3.999999……結果我用（int）轉型使他結果變爲了3，這種喚起了我大學程式課記憶的問題。
主要修改的方面：
修改了一下抓取到的metrics，使其更符合演算法需求
一些數學的運算公式及應該何時運算
調整了一些變量到AtomicInteger讓其在高頻調用中保持正常

現在資料的寫入能夠均勻的分布在各partition上，且經過測試速率相比之前也稍有進步。
![Screenshot from 2022-02-28 13-07-51](https://user-images.githubusercontent.com/43372856/155936620-7371a367-e121-4a0d-a8ba-d641d6328b1e.png)



Patch:
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 f1269bd5..32c30388 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClient.java
@@ -9,7 +9,11 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 import org.astraea.Utils;
 import org.astraea.metrics.collector.Receiver;
@@ -25,15 +29,15 @@ public class NodeLoadClient {
 
   private final List<Receiver> receiverList;
   private final Map<String, Integer> currentJmxAddresses;
+  private final Lock lock = new ReentrantLock();
   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 boolean notInMethod = true;
-  private double totalInput;
-  private double totalOutput;
+  private int totalInput;
+  private int 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)$";
@@ -52,17 +56,21 @@ public class NodeLoadClient {
    * @return each node load count in preset time
    */
   public Map<Integer, Integer> loadSituation(ClusterInfo cluster) {
-    if (overSecond(lastTime, 1) && notInMethod) {
-      notInMethod = false;
-      nodesOverLoad(cluster);
-      notInMethod = true;
+    if (overSecond(lastTime, 1) && lock.tryLock()) {
+      try {
+        lock.lock();
+        nodesOverLoad(cluster);
+      } finally {
+        lock.unlock();
+        lastTime = System.currentTimeMillis();
+      }
     }
     Utils.waitFor(() -> currentLoadNode != null);
     return currentLoadNode;
   }
 
   /** @param cluster the cluster from the partitioner */
-  private synchronized void nodesOverLoad(ClusterInfo cluster) {
+  private void nodesOverLoad(ClusterInfo cluster) {
     if (lastTime == -1) {
       var nodes = cluster.nodes();
       brokers =
@@ -107,27 +115,26 @@ public class NodeLoadClient {
         broker ->
             broker.load(
                 brokerLoad(broker.brokerSituation, totalBrokerSituation.get() / brokers.size())));
-    currentLoadNode =
+    this.currentLoadNode =
         brokers.stream().collect(Collectors.toMap(broker -> broker.brokerID, Broker::totalLoad));
-    lastTime = System.currentTimeMillis();
   }
 
   private void brokersMsgPerSec() {
     brokers.forEach(
         broker -> {
-          broker.output =
-              Double.parseDouble(
+          var currentOutput =
+              Long.parseLong(
                   broker
                       .metrics
-                      .get(KafkaMetrics.BrokerTopic.BytesOutPerSec.name())
+                      .get(KafkaMetrics.BrokerTopic.BytesInPerSec.name())
                       .current()
                       .get(0)
                       .beanObject()
                       .getAttributes()
-                      .get("MeanRate")
+                      .get("Count")
                       .toString());
-          broker.input =
-              Double.parseDouble(
+          var currentInput =
+              Long.parseLong(
                   broker
                       .metrics
                       .get(KafkaMetrics.BrokerTopic.BytesInPerSec.name())
@@ -135,8 +142,19 @@ public class NodeLoadClient {
                       .get(0)
                       .beanObject()
                       .getAttributes()
-                      .get("MeanRate")
+                      .get("Count")
                       .toString());
+          if (broker.lastInput == -1) {
+            broker.lastInput = currentInput;
+            broker.lastOutPut = currentOutput;
+          }
+
+          broker.output =
+              Integer.parseInt(String.valueOf((currentOutput - broker.lastOutPut) / 1024));
+          broker.lastOutPut = currentOutput;
+          broker.input = Integer.parseInt(String.valueOf((currentInput - broker.lastInput) / 1024));
+          broker.lastInput = currentInput;
+
           broker.jvm =
               (HasJvmMemory)
                   broker.metrics.get("Memory").current().stream().findAny().orElse(broker.jvm);
@@ -228,22 +246,22 @@ public class NodeLoadClient {
         .collect(Collectors.toList());
   }
 
-  private double totalInput() {
-    return brokers.stream().map(broker -> broker.input).reduce(0.0, Double::sum);
+  private int totalInput() {
+    return brokers.stream().map(broker -> broker.input).reduce(0, Integer::sum);
   }
 
-  private double totalOutput() {
-    return brokers.stream().map(broker -> broker.output).reduce(0.0, Double::sum);
+  private int totalOutput() {
+    return brokers.stream().map(broker -> broker.output).reduce(0, Integer::sum);
   }
 
-  // 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;
+    if (brokerSituation <= (1 - 0.5) * avg) {
+      return 0;
+    } else if (brokerSituation >= (1 + 0.5) * avg) {
+      return 2;
+    } else {
+      return 1;
+    }
   }
 
   public double thoughPutComparison(int brokerID) {
@@ -254,16 +272,17 @@ public class NodeLoadClient {
     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 final ConcurrentMap<Integer, Integer> load = new ConcurrentHashMap<>();
+    private final ConcurrentMap<String, Receiver> metrics = new ConcurrentHashMap<>();
+    private int input;
+    private long lastInput = -1;
+    private int output;
+    private long lastOutPut = -1;
     private double maxThoughPut;
     private double brokerSituation;
     private double thoughPutComparison;
@@ -286,7 +305,7 @@ public class NodeLoadClient {
     }
 
     private void load(Integer currentLoad) {
-      this.load.put(count % 10, currentLoad);
+      load.put(count % 10, currentLoad);
       count++;
     }
 
@@ -294,12 +313,12 @@ public class NodeLoadClient {
       this.brokerSituationNormalized = brokerSituation / totalSituation;
     }
 
-    private void inputNormalized(double brokersInput) {
-      this.inputNormalized = (input + 1) / (brokersInput + 1);
+    private void inputNormalized(int brokersInput) {
+      this.inputNormalized = (input + 1.0) / (brokersInput + 1);
     }
 
-    private void outputNormalized(double brokersOutput) {
-      this.outputNormalized = (output + 1) / (brokersOutput + 1);
+    private void outputNormalized(int brokersOutput) {
+      this.outputNormalized = (output + 1.0) / (brokersOutput + 1);
     }
 
     private void brokerSituation() {
diff --git a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
index 81411eab..9036314c 100644
--- a/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
+++ b/app/src/main/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtils.java
@@ -14,6 +14,11 @@ public class PartitionerUtils {
     return poissonMap;
   }
 
+  public static int weightPoisson(Double value, Double thoughPutAbility) {
+    if (value < 0.8) return (int) (Math.round((1 - value) * 20 * thoughPutAbility));
+    else return (int) Math.pow(Math.round((1 - value) * 20 * thoughPutAbility), 2) / 10;
+  }
+
   static double doPoisson(int lambda, int x) {
     var Probability = 0.0;
     var ans = 0.0;
@@ -37,7 +42,7 @@ public class PartitionerUtils {
   private static int avgLoadCount(Map<Integer, Integer> overLoadCount) {
     var avgLoadCount =
         overLoadCount.values().stream().mapToDouble(Integer::doubleValue).average().orElse(0);
-    return (int) avgLoadCount;
+    return (int) Math.round(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 5d2a53d6..b1632239 100644
--- a/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
+++ b/app/src/main/java/org/astraea/partitioner/smoothPartitioner/SmoothWeightPartitioner.java
@@ -2,14 +2,19 @@ package org.astraea.partitioner.smoothPartitioner;
 
 import static org.astraea.Utils.overSecond;
 import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.allPoisson;
+import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.weightPoisson;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Random;
-import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import org.apache.kafka.clients.producer.Partitioner;
 import org.apache.kafka.common.Cluster;
 import org.astraea.partitioner.ClusterInfo;
@@ -26,49 +31,46 @@ public class SmoothWeightPartitioner implements Partitioner {
    * Record the current weight of each node according to Poisson calculation and the weight after
    * partitioner calculation.
    */
-  private final Map<Integer, int[]> brokersWeight = new HashMap<>();
+  private final ConcurrentMap<Integer, SmoothWeightServer> brokersWeight =
+      new ConcurrentHashMap<>();
 
   private NodeLoadClient nodeLoadClient;
-
   private long lastTime = -1;
   private final Random rand = new Random();
+  private final Lock lock = new ReentrantLock();
+  private int weightSum = -1;
 
   @Override
   public int partition(
       String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
-    Map<Integer, Integer> loadCount;
-    loadCount = nodeLoadClient.loadSituation(ClusterInfo.of(cluster));
+    var loadCount = nodeLoadClient.loadSituation(ClusterInfo.of(cluster));
     Objects.requireNonNull(loadCount, "OverLoadCount should not be null.");
     updateWeightIfNeed(loadCount);
-    AtomicReference<Map.Entry<Integer, int[]>> maxWeightServer = new AtomicReference<>();
-    brokersWeight.forEach(
-        (nodeID, weight) -> {
-          if (maxWeightServer.get() == null || weight[1] > maxWeightServer.get().getValue()[1]) {
-            maxWeightServer.set(Map.entry(nodeID, weight));
-          }
-          brokersWeight.put(nodeID, new int[] {weight[0], weight[1] + weight[0]});
-        });
-    Objects.requireNonNull(maxWeightServer.get(), "MaxWeightServer should not be null.");
-    brokersWeight.put(
-        maxWeightServer.get().getKey(),
-        new int[] {
-          maxWeightServer.get().getValue()[0],
-          maxWeightServer.get().getValue()[1] - allNodesWeight()
-        });
-    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()));
+
+    SmoothWeightServer maxWeightServer = null;
+
+    Iterator<Integer> iterator = brokersWeight.keySet().iterator();
+    var currentWrightSum = weightSum;
+    while (iterator.hasNext()) {
+      var smoothWeightServer = brokersWeight.get(iterator.next());
+      if (smoothWeightServer != null) {
+        smoothWeightServer.originalWeight.updateAndGet((v) -> v + smoothWeightServer.currentWeight);
+
+        if (maxWeightServer == null) {
+          maxWeightServer = smoothWeightServer;
+        }
+        if (smoothWeightServer.originalWeight() > maxWeightServer.originalWeight()) {
+          maxWeightServer = smoothWeightServer;
+        }
+      }
+    }
+
+    Objects.requireNonNull(maxWeightServer, "MaxWeightServer should not be null.");
+    maxWeightServer.updateOriginalWeight(currentWrightSum);
+
+    var partitions = cluster.partitionsForNode(maxWeightServer.brokerID);
+
+    return partitions.get(rand.nextInt(partitions.size())).partition();
   }
 
   @Override
@@ -94,43 +96,76 @@ public class SmoothWeightPartitioner implements Partitioner {
   }
 
   /** Change the weight of the node according to the current Poisson. */
-  // visible for test
-  public synchronized void brokersWeight(Map<Integer, Double> poissonMap) {
+  synchronized void brokersWeight(Map<Integer, Double> poissonMap) {
+    AtomicInteger sum = new AtomicInteger(0);
     poissonMap.forEach(
         (key, value) -> {
           var thoughPutAbility = nodeLoadClient.thoughPutComparison(key);
-          if (!brokersWeight.containsKey(key)) {
-            brokersWeight.put(key, new int[] {(int) ((1 - value) * 20 * thoughPutAbility), 0});
-          } else {
-            brokersWeight.put(
-                key,
-                new int[] {(int) ((1 - value) * 20 * thoughPutAbility), brokersWeight.get(key)[1]});
+          if (!brokersWeight.containsKey(key))
+            brokersWeight.putIfAbsent(
+                key, new SmoothWeightServer(key, weightPoisson(value, thoughPutAbility)));
+          else {
+            var broker = brokersWeight.get(key);
+            broker.currentWeight(weightPoisson(value, thoughPutAbility));
+          }
+          if (memoryWarning(key)) {
+            subBrokerWeight(key, 100);
           }
+          sum.addAndGet(brokersWeight.get(key).currentWeight);
         });
+
+    weightSum = sum.get();
   }
 
   void updateWeightIfNeed(Map<Integer, Integer> loadCount) {
-    if (overSecond(lastTime, 1)) {
-      brokersWeight(allPoisson(loadCount));
-      lastTime = System.currentTimeMillis();
+    if (overSecond(lastTime, 1) && lock.tryLock()) {
+      try {
+        lock.lock();
+        brokersWeight(allPoisson(loadCount));
+      } finally {
+        lock.unlock();
+        lastTime = System.currentTimeMillis();
+      }
     }
   }
 
-  Map<Integer, int[]> brokersWeight() {
+  Map<Integer, SmoothWeightServer> brokersWeight() {
     return brokersWeight;
   }
 
   private void subBrokerWeight(int brokerID, int subNumber) {
-    brokersWeight.put(
-        brokerID,
-        new int[] {brokersWeight.get(brokerID)[0], brokersWeight.get(brokerID)[1] - subNumber});
-  }
-
-  private synchronized int allNodesWeight() {
-    return brokersWeight.values().stream().mapToInt(vs -> vs[0]).sum();
+    brokersWeight.get(brokerID).updateOriginalWeight(subNumber);
   }
 
   private boolean memoryWarning(int brokerID) {
     return nodeLoadClient.memoryUsage(brokerID) >= 0.8;
   }
+
+  static class SmoothWeightServer {
+    private final int brokerID;
+    private AtomicInteger originalWeight;
+    private int currentWeight;
+
+    SmoothWeightServer(int brokerID, int currentWeight) {
+      this.brokerID = brokerID;
+      this.currentWeight = currentWeight;
+      this.originalWeight = new AtomicInteger(currentWeight);
+    }
+
+    int updateOriginalWeight(int subWeight) {
+      return originalWeight.getAndUpdate((v) -> v - subWeight);
+    }
+
+    int originalWeight() {
+      return originalWeight.get();
+    }
+
+    void currentWeight(int currentWeight) {
+      this.currentWeight = currentWeight;
+    }
+
+    int currentWeight() {
+      return currentWeight;
+    }
+  }
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
index 909f7dd1..3c754173 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/NodeLoadClientTest.java
@@ -4,11 +4,9 @@ import static org.mockito.Mockito.when;
 
 import java.io.IOException;
 import java.lang.reflect.Field;
-import java.net.UnknownHostException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.astraea.Utils;
 import org.astraea.partitioner.ClusterInfo;
@@ -55,9 +53,9 @@ public class NodeLoadClientTest extends RequireBrokerCluster {
     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.37, 1.0 / 3), 1);
     Assertions.assertEquals(nodeLoadClient.brokerLoad(0.01, 1.0 / 3), 0);
-    Assertions.assertEquals(nodeLoadClient.brokerLoad(0.8, 1.0 / 3), 10);
+    Assertions.assertEquals(nodeLoadClient.brokerLoad(0.8, 1.0 / 3), 2);
   }
 
   @Test
@@ -81,7 +79,7 @@ public class NodeLoadClientTest extends RequireBrokerCluster {
   }
 
   @Test
-  void testLoadSituation() throws UnknownHostException {
+  void testLoadSituation() {
     var bootstrapServers = List.of(bootstrapServers().split(","));
     List<NodeInfo> nodes = new ArrayList<>();
     AtomicInteger count = new AtomicInteger(0);
@@ -95,19 +93,19 @@ public class NodeLoadClientTest extends RequireBrokerCluster {
     var cluster = Mockito.mock(ClusterInfo.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);
+    Assertions.assertEquals(load.get(0), 1);
+    Assertions.assertEquals(load.get(1), 1);
+    Assertions.assertEquals(load.get(2), 1);
     load = nodeLoadClient.loadSituation(cluster);
-    Assertions.assertEquals(load.get(0), 5);
-    Assertions.assertEquals(load.get(1), 5);
-    Assertions.assertEquals(load.get(2), 5);
+    Assertions.assertEquals(load.get(0), 1);
+    Assertions.assertEquals(load.get(1), 1);
+    Assertions.assertEquals(load.get(2), 1);
     var lastTime = System.currentTimeMillis();
     Utils.waitFor(() -> Utils.overSecond(lastTime, 1));
     load = nodeLoadClient.loadSituation(cluster);
-    Assertions.assertEquals(load.get(0), 10);
-    Assertions.assertEquals(load.get(1), 10);
-    Assertions.assertEquals(load.get(2), 10);
+    Assertions.assertEquals(load.get(0), 2);
+    Assertions.assertEquals(load.get(1), 2);
+    Assertions.assertEquals(load.get(2), 2);
   }
 
   private NodeLoadClient.Broker setSituationNormalized(
@@ -130,12 +128,4 @@ public class NodeLoadClientTest extends RequireBrokerCluster {
       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/PartitionerUtilsTest.java b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
index 8729045e..bc144edd 100644
--- a/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
+++ b/app/src/test/java/org/astraea/partitioner/nodeLoadMetric/PartitionerUtilsTest.java
@@ -4,6 +4,7 @@ import static org.astraea.partitioner.nodeLoadMetric.PartitionerUtils.*;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import java.util.HashMap;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 public class PartitionerUtilsTest {
@@ -39,4 +40,14 @@ public class PartitionerUtilsTest {
 
     assertEquals(poissonMap, testPoissonMap);
   }
+
+  @Test
+  void testWeightPoisson() {
+    Assertions.assertEquals(weightPoisson(0.5, 1.0), 10);
+    Assertions.assertEquals(weightPoisson(1.0, 1.0), 0);
+    Assertions.assertEquals(weightPoisson(0.95, 1.0), 0);
+    Assertions.assertEquals(weightPoisson(0.9, 1.0), 0);
+    Assertions.assertEquals(weightPoisson(0.0, 1.0), 20);
+    Assertions.assertEquals(weightPoisson(0.8, 1.0), 1);
+  }
 }
diff --git a/app/src/test/java/org/astraea/partitioner/smoothPartitioner/PartitionerTest.java b/app/src/test/java/org/astraea/partitioner/smoothPartitioner/PartitionerTest.java
index c8d65f6b..6fc9cf4e 100644
--- a/app/src/test/java/org/astraea/partitioner/smoothPartitioner/PartitionerTest.java
+++ b/app/src/test/java/org/astraea/partitioner/smoothPartitioner/PartitionerTest.java
@@ -3,7 +3,8 @@ 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.*;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.lang.reflect.Field;
 import java.time.Duration;
@@ -18,6 +19,7 @@ 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.Utils;
 import org.astraea.concurrent.ThreadPool;
 import org.astraea.consumer.Consumer;
 import org.astraea.consumer.Deserializer;
@@ -62,13 +64,11 @@ public class PartitionerTest extends RequireBrokerCluster {
     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);
+    var brokerWeight =
+        (Map<Integer, SmoothWeightPartitioner.SmoothWeightServer>)
+            requireField(smoothWeightPartitioner, "brokersWeight");
+    assertEquals(brokerWeight.get(0).originalWeight(), 10);
+    assertEquals((brokerWeight.get(1).originalWeight()), 1);
   }
 
   @Test
@@ -194,18 +194,28 @@ public class PartitionerTest extends RequireBrokerCluster {
     field.set(smoothWeightPartitioner, nodeLoadClient);
     when(nodeLoadClient.thoughPutComparison(anyInt())).thenReturn(1.0);
     smoothWeightPartitioner.updateWeightIfNeed(loadCount);
-    var firstBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0)[0];
+    var firstBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0).currentWeight();
     smoothWeightPartitioner.updateWeightIfNeed(loadCount);
-    var secondBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0)[0];
+    var secondBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0).currentWeight();
     Assertions.assertEquals(firstBrokersWeight, secondBrokersWeight);
     loadCount.put(0, 5);
     loadCount.put(2, 5);
-    sleep(2);
+    var lastTime = System.currentTimeMillis();
+    Utils.waitFor(() -> Utils.overSecond(lastTime, 1));
     smoothWeightPartitioner.updateWeightIfNeed(loadCount);
-    var thirdBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0)[0];
+    var thirdBrokersWeight = smoothWeightPartitioner.brokersWeight().get(0).currentWeight();
     Assertions.assertNotEquals(secondBrokersWeight, thirdBrokersWeight);
   }
 
+  @Test
+  void testUpdateOriginalWeight() {
+    var smoothWeightServer = new SmoothWeightPartitioner.SmoothWeightServer(1, 1);
+
+    Assertions.assertEquals(smoothWeightServer.updateOriginalWeight(5), 1);
+    Assertions.assertEquals(smoothWeightServer.updateOriginalWeight(5), -4);
+    Assertions.assertEquals(smoothWeightServer.originalWeight(), -9);
+  }
+
   private ThreadPool.Executor producerExecutor(
       Producer<String, byte[]> producer, String topic, String key, Header header, long timeStamp) {
     return new ThreadPool.Executor() {