Smooth weight Round Robin algorithm

app/src/main/java/org/astraea/Utils.java

@@ -141,6 +141,10 @@ public static boolean overSecond(long lastTime, int second) {
     return (lastTime + Duration.ofSeconds(second).toMillis()) < System.currentTimeMillis();
   }
 
+  public static boolean overSecond(long lastTime, Duration second) {
+    return (lastTime + second.toMillis()) < System.currentTimeMillis();
+  }
+
   public static void sleep(Duration duration) {
     try {
       TimeUnit.MILLISECONDS.sleep(duration.toMillis());

app/src/main/java/org/astraea/cost/Periodic.java

@@ -1,5 +1,6 @@
 package org.astraea.cost;
 
+import java.time.Duration;
 import java.util.function.Supplier;
 import org.astraea.Utils;
 
@@ -25,4 +26,27 @@ protected Value tryUpdate(Supplier<Value> updater) {
   protected long currentTime() {
     return System.currentTimeMillis();
   }
+
+  /**
+   * Updates the value interval second.
+   *
+   * @param updater Methods that need to be updated regularly.
+   * @param interval Required interval.
+   * @return an object of type Value created from the parameter value.
+   */
+  protected Value tryUpdate(Supplier<Value> updater, int interval) {
+    if (Utils.overSecond(lastUpdate, interval)) {
+      value = updater.get();
+      lastUpdate = currentTime();
+    }
+    return value;
+  }
+
+  protected Value tryUpdate(Supplier<Value> updater, Duration interval) {
+    if (Utils.overSecond(lastUpdate, interval)) {
+      value = updater.get();
+      lastUpdate = currentTime();
+    }
+    return value;
+  }
 }

app/src/main/java/org/astraea/cost/broker/CostUtils.java

@@ -43,4 +43,10 @@ public static Map<Integer, Double> TScore(Map<Integer, Double> metrics) {
                   return score;
                 }));
   }
+
+  public static Map<Integer, Double> normalize(Map<Integer, Double> score) {
+    var sum = score.values().stream().mapToDouble(d -> d).sum();
+    return score.entrySet().stream()
+        .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getKey() / sum));
+  }
 }

app/src/main/java/org/astraea/partitioner/smooth/SmoothWeightRoundRobin.java

@@ -0,0 +1,111 @@
+package org.astraea.partitioner.smooth;
+
+import java.time.Duration;
+import java.util.Comparator;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.astraea.cost.Periodic;
+import org.astraea.cost.broker.CostUtils;
+
+/**
+ * Given initial key-score pair, it will output a preferred key with the highest current weight. The
+ * current weight of the chosen key will decrease the sum of effective weight. And all current
+ * weight will increment by its effective weight. It may result in "higher score with higher chosen
+ * rate". For example:
+ *
+ * <p>||========================================||============================================||
+ * ||------------ Broker in cluster ------------||------------- Effective weight -------------||
+ * ||------------------Broker1------------------||-------------------- 5 ---------------------||
+ * ||------------------Broker2------------------||-------------------- 1 ---------------------||
+ * ||------------------Broker3------------------||-------------------- 1 ---------------------||
+ * ||===========================================||============================================||
+ *
+ * <p>||===================||=======================||===============||======================||
+ * ||--- Request Number ---|| Before current weight || Target Broker || After current weight ||
+ * ||----------1-----------||------ {5, 1, 1} ------||----Broker1----||----- {-2, 1, 1} -----||
+ * ||----------2-----------||------ {3, 2, 2} ------||----Broker1----||----- {-4, 2, 2} -----||
+ * ||----------3-----------||------ {1, 3, 3} ------||----Broker2----||----- { 1,-4, 3} -----||
+ * ||----------4-----------||------ {6,-3, 4} ------||----Broker1----||----- {-1,-3, 4} -----||
+ * ||----------5-----------||------ {4,-2, 5} ------||----Broker3----||----- { 4,-2,-2} -----||
+ * ||----------6-----------||------ {9,-1,-1} ------||----Broker1----||----- { 2,-1,-1} -----||
+ * ||----------7-----------||------ {7, 0, 0} ------||----Broker1----||----- { 0, 0, 0} -----||
+ * ||======================||=======================||===============||======================||
+ */
+public final class SmoothWeightRoundRobin
+    extends Periodic<SmoothWeightRoundRobin.EffectiveWeightResult> {
+  private EffectiveWeightResult effectiveWeightResult;
+  public Map<Integer, Double> currentWeight;
+
+  public SmoothWeightRoundRobin(Map<Integer, Double> effectiveWeight) {
+    effectiveWeightResult =
+        new EffectiveWeightResult(
+            effectiveWeight.entrySet().stream()
+                .collect(Collectors.toMap(Map.Entry::getKey, ignored -> 1.0)));
+    currentWeight =
+        effectiveWeight.entrySet().stream()
+            .collect(Collectors.toMap(Map.Entry::getKey, ignored -> 0.0));
+  }
+
+  public synchronized void init(Map<Integer, Double> brokerScore) {
+    effectiveWeightResult =
+        tryUpdate(
+            () -> {
+              var normalizationLoad = CostUtils.normalize(brokerScore);
+              return new EffectiveWeightResult(
+                  this.effectiveWeightResult.effectiveWeight.entrySet().stream()
+                      .collect(
+                          Collectors.toMap(
+                              entry -> entry.getKey(),
+                              entry -> {
+                                var nLoad = normalizationLoad.get(entry.getKey());
+                                var weight =
+                                    entry.getValue()
+                                        * (nLoad.isNaN()
+                                            ? 1.0
+                                            : ((nLoad + 1) > 0 ? nLoad + 1 : 0.1));
+                                if (weight > 2.0) return 2.0;
+                                return Math.max(weight, 0.0);
+                              })));
+            },
+            Duration.ofSeconds(10));
+  }
+
+  /**
+   * Get the preferred ID, and update the state.
+   *
+   * @return the preferred ID
+   */
+  public synchronized int getAndChoose() {
+    this.currentWeight =
+        this.currentWeight.entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    e -> e.getValue() + effectiveWeightResult.effectiveWeight.get(e.getKey())));
+    var maxID =
+        this.currentWeight.entrySet().stream()
+            .max(Comparator.comparingDouble(Map.Entry::getValue))
+            .map(Map.Entry::getKey)
+            .orElse(0);
+    this.currentWeight =
+        this.currentWeight.entrySet().stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    e ->
+                        e.getKey().equals(maxID)
+                            ? e.getValue() - effectiveWeightResult.effectiveWeightSum
+                            : e.getValue()));
+    return maxID;
+  }
+
+  public static class EffectiveWeightResult {
+    private final Map<Integer, Double> effectiveWeight;
+    private final double effectiveWeightSum;
+
+    EffectiveWeightResult(Map<Integer, Double> effectiveWeight) {
+      this.effectiveWeight = effectiveWeight;
+      this.effectiveWeightSum = effectiveWeight.values().stream().mapToDouble(i -> i).sum();
+    }
+  }
+}