From b136847e1e7b911fc8153c9efa1ddea96f2f518c Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Thu, 24 Oct 2024 16:46:12 -0700
Subject: [PATCH 01/12] check and warn for invalid partition weight

---
 .../balancer/servers/ZooKeeperAnnouncer.java  | 57 ++++++++++++-------
 1 file changed, 36 insertions(+), 21 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index 04152ca74..4a85f00a1 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -20,6 +20,7 @@
 
 package com.linkedin.d2.balancer.servers;
 
+import java.math.BigDecimal;
 import java.net.URI;
 import java.util.ArrayDeque;
 import java.util.Collections;
@@ -93,6 +94,7 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   private final AtomicLong _markDownStartAtRef = new AtomicLong(Long.MAX_VALUE);
 
   private volatile Map<Integer, PartitionData> _partitionDataMap;
+  private final Integer _maxWeight; // max weight for the weight in PartitionData, if null, no max weight is set.
   private volatile Map<String, Object> _uriSpecificProperties;
 
   private ServiceDiscoveryEventEmitter _eventEmitter;
@@ -123,13 +125,13 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   private volatile boolean _markUpFailed;
 
   // ScheduledExecutorService to schedule the end of dark warm-up, defaults to null
-  private ScheduledExecutorService _executorService;
+  private final ScheduledExecutorService _executorService;
 
   // Boolean flag to indicate if dark warm-up is enabled, defaults to false
-  private boolean _isDarkWarmupEnabled;
+  private final boolean _isDarkWarmupEnabled;
 
   // String to store the name of the dark warm-up cluster, defaults to null
-  private String _warmupClusterName;
+  private final String _warmupClusterName;
   // Similar as _znodePath and _znodeData above but for the warm up cluster.
   private final AtomicReference<String> _warmupClusterZnodePathRef = new AtomicReference<>();
   private final AtomicReference<String> _warmupClusterZnodeDataRef = new AtomicReference<>();
@@ -139,7 +141,7 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   private final AtomicLong _warmupClusterMarkDownStartAtRef = new AtomicLong(Long.MAX_VALUE);
 
   // Field to store the dark warm-up time duration in seconds, defaults to zero
-  private int _warmupDuration;
+  private final int _warmupDuration;
 
   /**
    * @deprecated Use the constructor {@link #ZooKeeperAnnouncer(LoadBalancerServer)} instead.
@@ -195,26 +197,18 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
   public ZooKeeperAnnouncer(ZooKeeperServer server, boolean initialIsUp,
                             boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService, ServiceDiscoveryEventEmitter eventEmitter)
   {
-    _server = server;
-    // initialIsUp is used for delay mark up. If it's false, there won't be markup when the announcer is started.
-    _isUp = initialIsUp;
-    _isWarmingUp = false;
-    _isRetryWarmup = false;
-    _pendingMarkDown = new ArrayDeque<>();
-    _pendingMarkUp = new ArrayDeque<>();
-    _pendingWarmupMarkDown = new ArrayDeque<>();
-
-    _isDarkWarmupEnabled = isDarkWarmupEnabled;
-    _warmupClusterName = warmupClusterName;
-    _warmupDuration = warmupDuration;
-    _executorService = executorService;
-    _eventEmitter = eventEmitter;
-
-    server.setServiceDiscoveryEventHelper(this);
+    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null);
   }
 
   public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
       boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService, ServiceDiscoveryEventEmitter eventEmitter)
+  {
+    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null);
+  }
+
+  public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
+      boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService,
+      ServiceDiscoveryEventEmitter eventEmitter, Integer maxWeight)
   {
     _server = server;
     // initialIsUp is used for delay mark up. If it's false, there won't be markup when the announcer is started.
@@ -230,6 +224,7 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
     _warmupDuration = warmupDuration;
     _executorService = executorService;
     _eventEmitter = eventEmitter;
+    _maxWeight = maxWeight;
 
     if (server instanceof ZooKeeperServer)
     {
@@ -750,11 +745,12 @@ public void setWeight(double weight)
 
     Map<Integer, PartitionData> partitionDataMap = new HashMap<>(1);
     partitionDataMap.put(partitionId, new PartitionData(weight));
-    _partitionDataMap = Collections.unmodifiableMap(partitionDataMap);
+    setPartitionData(partitionDataMap);
   }
 
   public void setPartitionData(Map<Integer, PartitionData> partitionData)
   {
+    validatePartitionData(partitionData);
     _partitionDataMap = Collections.unmodifiableMap(new HashMap<>(partitionData));
   }
 
@@ -832,5 +828,24 @@ private ImmutablePair<String, String> getZnodePathAndData(String cluster) {
    */
   public boolean isWarmingUp() {
     return _isWarmingUp;
+
+  private void validatePartitionData(Map<Integer, PartitionData> partitionData) {
+    for (Map.Entry<Integer, PartitionData> entry : partitionData.entrySet()) {
+      double weight = entry.getValue().getWeight();
+      // check max weight
+      if (_maxWeight != null && weight > _maxWeight) {
+        _log.warn("", new IllegalArgumentException(String.format("[ACTION NEEDED] Weight %f in Partition %d is greater"
+            + " than the max weight allowed: %d. Please correct the weight. It will be force-capped to the max weight "
+            + "in the future.", weight, entry.getKey(), _maxWeight)));
+        // TODO: throw exception or cap weight to max weight
+      }
+
+      // check decimal places
+      if (BigDecimal.valueOf(weight).scale() > 2) {
+        _log.warn("", new IllegalArgumentException(String.format("Weight %f in Partition %d has more than 2 decimal "
+            + "places. It will be rounded to 2 decimal places in the future.", weight, entry.getKey())));
+        // TODO: round weight to 2 decimal places.
+      }
+    }
   }
 }

From d2dd399ad37392125acf73641dd9d1e3bc01846e Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Mon, 28 Oct 2024 14:18:50 -0700
Subject: [PATCH 02/12] set configurable actions on weight validation breach

---
 .../balancer/servers/ZooKeeperAnnouncer.java  | 91 +++++++++++++++----
 1 file changed, 73 insertions(+), 18 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index 4a85f00a1..cbc346266 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -21,6 +21,7 @@
 package com.linkedin.d2.balancer.servers;
 
 import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.net.URI;
 import java.util.ArrayDeque;
 import java.util.Collections;
@@ -94,7 +95,14 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   private final AtomicLong _markDownStartAtRef = new AtomicLong(Long.MAX_VALUE);
 
   private volatile Map<Integer, PartitionData> _partitionDataMap;
-  private final Integer _maxWeight; // max weight for the weight in PartitionData, if null, no max weight is set.
+  /**
+   * If not null, it defines two rules for d2 weight validation:
+   * 1. The maximum d2 weight allowed.
+   * 2. The maximum number of decimal places allowed. Use 0s on decimal places to indicate it.
+   * For example, 100.00 means the max weight allowed is 100 and the max number of decimal places is 2.
+   */
+  private final BigDecimal _maxWeight;
+  private ActOnWeightBreach _actOnWeightBreach = ActOnWeightBreach.IGNORE;
   private volatile Map<String, Object> _uriSpecificProperties;
 
   private ServiceDiscoveryEventEmitter _eventEmitter;
@@ -143,6 +151,18 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   // Field to store the dark warm-up time duration in seconds, defaults to zero
   private final int _warmupDuration;
 
+  /**
+   * The action to take when d2 weight breaches validation rules.
+   */
+  public enum ActOnWeightBreach {
+    // Ignore and no op.
+    IGNORE,
+    //
+    WARN,
+    THROW,
+    RECTIFY
+  }
+
   /**
    * @deprecated Use the constructor {@link #ZooKeeperAnnouncer(LoadBalancerServer)} instead.
    */
@@ -197,18 +217,18 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
   public ZooKeeperAnnouncer(ZooKeeperServer server, boolean initialIsUp,
                             boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService, ServiceDiscoveryEventEmitter eventEmitter)
   {
-    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null);
+    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null, ActOnWeightBreach.IGNORE);
   }
 
   public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
       boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService, ServiceDiscoveryEventEmitter eventEmitter)
   {
-    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null);
+    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null, ActOnWeightBreach.IGNORE);
   }
 
   public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
       boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService,
-      ServiceDiscoveryEventEmitter eventEmitter, Integer maxWeight)
+      ServiceDiscoveryEventEmitter eventEmitter, BigDecimal maxWeight, ActOnWeightBreach actOnWeightBreach)
   {
     _server = server;
     // initialIsUp is used for delay mark up. If it's false, there won't be markup when the announcer is started.
@@ -225,6 +245,7 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
     _executorService = executorService;
     _eventEmitter = eventEmitter;
     _maxWeight = maxWeight;
+    _actOnWeightBreach = actOnWeightBreach;
 
     if (server instanceof ZooKeeperServer)
     {
@@ -750,8 +771,7 @@ public void setWeight(double weight)
 
   public void setPartitionData(Map<Integer, PartitionData> partitionData)
   {
-    validatePartitionData(partitionData);
-    _partitionDataMap = Collections.unmodifiableMap(new HashMap<>(partitionData));
+    _partitionDataMap = Collections.unmodifiableMap(new HashMap<>(validatePartitionData(partitionData)));
   }
 
   public Map<Integer, PartitionData> getPartitionData()
@@ -829,23 +849,58 @@ private ImmutablePair<String, String> getZnodePathAndData(String cluster) {
   public boolean isWarmingUp() {
     return _isWarmingUp;
 
-  private void validatePartitionData(Map<Integer, PartitionData> partitionData) {
-    for (Map.Entry<Integer, PartitionData> entry : partitionData.entrySet()) {
-      double weight = entry.getValue().getWeight();
+  private Map<Integer, PartitionData> validatePartitionData(Map<Integer, PartitionData> partitionData) {
+    if (_maxWeight == null || _actOnWeightBreach == ActOnWeightBreach.IGNORE) {
+      return partitionData;
+    }
+
+    Map<Integer, PartitionData> res = new HashMap<>(partitionData); // modifiable copy in case the input is unmodifiable
+    for (Map.Entry<Integer, PartitionData> entry : res.entrySet()) {
+      BigDecimal weight = BigDecimal.valueOf(entry.getValue().getWeight());
       // check max weight
-      if (_maxWeight != null && weight > _maxWeight) {
-        _log.warn("", new IllegalArgumentException(String.format("[ACTION NEEDED] Weight %f in Partition %d is greater"
-            + " than the max weight allowed: %d. Please correct the weight. It will be force-capped to the max weight "
-            + "in the future.", weight, entry.getKey(), _maxWeight)));
-        // TODO: throw exception or cap weight to max weight
+      if (weight.compareTo(_maxWeight) > 0) {
+        switch (_actOnWeightBreach) {
+          case WARN:
+            _log.warn("", getMaxWeightBreachException(weight, entry.getKey()));
+            break;
+          case THROW:
+            throw getMaxWeightBreachException(weight, entry.getKey());
+          case RECTIFY:
+            entry.setValue(new PartitionData(_maxWeight.intValue()));
+            _log.warn("Capped weight {} in Partition {} to the max weight allowed: {}.", weight.doubleValue(),
+                entry.getKey(), _maxWeight.intValue());
+            break;
+          default:
+            break;
+        }
       }
 
       // check decimal places
-      if (BigDecimal.valueOf(weight).scale() > 2) {
-        _log.warn("", new IllegalArgumentException(String.format("Weight %f in Partition %d has more than 2 decimal "
-            + "places. It will be rounded to 2 decimal places in the future.", weight, entry.getKey())));
-        // TODO: round weight to 2 decimal places.
+      if (weight.scale() > _maxWeight.scale()) {
+        switch (_actOnWeightBreach) {
+          case WARN: // both WARN and THROW only log the warning. Don't throw exception for decimal places.
+          case THROW:
+            _log.warn("", new IllegalArgumentException(String.format("Weight %f in Partition %d has more than %d"
+                + " decimal places. It will be rounded in the future.", weight.doubleValue(), entry.getKey(),
+                _maxWeight.scale())));
+            break;
+          case RECTIFY:
+            double newWeight = weight.setScale(_maxWeight.scale(), RoundingMode.HALF_UP).doubleValue();
+            entry.setValue(new PartitionData(newWeight));
+            _log.warn("Rounded weight {} in Partition {} to {} decimal places: {}.", weight.doubleValue(),
+                entry.getKey(), _maxWeight.scale(), newWeight);
+            break;
+          default:
+            break;
+        }
       }
     }
+    return res;
+  }
+
+  private IllegalArgumentException getMaxWeightBreachException(BigDecimal weight, int partition) {
+    return new IllegalArgumentException(String.format("[ACTION NEEDED] Weight %f in Partition %d is greater"
+        + " than the max weight allowed: %d. Please correct the weight. It will be force-capped to the max weight "
+        + "in the future.", weight.doubleValue(), partition, _maxWeight.intValue()));
   }
 }

From 6481440926f2d097cb118daa9b35d37e2b6509c0 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Mon, 28 Oct 2024 14:29:56 -0700
Subject: [PATCH 03/12] rebase master

---
 CHANGELOG.md                                                | 6 +++++-
 .../linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java    | 1 +
 gradle.properties                                           | 2 +-
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9de306a72..5bfa211c9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,9 @@ and what APIs have changed, if applicable.
 
 ## [Unreleased]
 
+## [29.62.0] - 2024-10-28
+- Check and take configurable action for invalid partition weight
+
 ## [29.61.0] - 2024-10-24
 - Disable dark traffic dispatching during dark warmup
 
@@ -5752,7 +5755,8 @@ patch operations can re-use these classes for generating patch messages.
 
 ## [0.14.1]
 
-[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.61.0...master
+[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.62.0...master
+[29.62.0]: https://github.com/linkedin/rest.li/compare/v29.61.0...v29.62.0
 [29.61.0]: https://github.com/linkedin/rest.li/compare/v29.60.0...v29.61.0
 [29.60.0]: https://github.com/linkedin/rest.li/compare/v29.59.0...v29.60.0
 [29.59.0]: https://github.com/linkedin/rest.li/compare/v29.58.11...v29.59.0
diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index cbc346266..8f78343d3 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -848,6 +848,7 @@ private ImmutablePair<String, String> getZnodePathAndData(String cluster) {
    */
   public boolean isWarmingUp() {
     return _isWarmingUp;
+  }
 
   private Map<Integer, PartitionData> validatePartitionData(Map<Integer, PartitionData> partitionData) {
     if (_maxWeight == null || _actOnWeightBreach == ActOnWeightBreach.IGNORE) {
diff --git a/gradle.properties b/gradle.properties
index b3a11e91d..be86bdb5d 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -1,4 +1,4 @@
-version=29.61.0
+version=29.62.0
 group=com.linkedin.pegasus
 org.gradle.configureondemand=true
 org.gradle.parallel=true

From c9b8f594cb6f4464df7219882421b99ac92fcfd3 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Mon, 28 Oct 2024 14:36:50 -0700
Subject: [PATCH 04/12] complete doc

---
 .../d2/balancer/servers/ZooKeeperAnnouncer.java  | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index 8f78343d3..949f2c2d1 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -102,6 +102,9 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
    * For example, 100.00 means the max weight allowed is 100 and the max number of decimal places is 2.
    */
   private final BigDecimal _maxWeight;
+  /**
+   * The action to take when d2 weight breaches validation rules.
+   */
   private ActOnWeightBreach _actOnWeightBreach = ActOnWeightBreach.IGNORE;
   private volatile Map<String, Object> _uriSpecificProperties;
 
@@ -151,15 +154,14 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   // Field to store the dark warm-up time duration in seconds, defaults to zero
   private final int _warmupDuration;
 
-  /**
-   * The action to take when d2 weight breaches validation rules.
-   */
   public enum ActOnWeightBreach {
     // Ignore and no op.
     IGNORE,
-    //
+    // only log warnings
     WARN,
+    // throw exceptions
     THROW,
+    // rectify the invalid weight (e.g: cap to the max, round to the nearest valid decimal places)
     RECTIFY
   }
 
@@ -245,7 +247,11 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
     _executorService = executorService;
     _eventEmitter = eventEmitter;
     _maxWeight = maxWeight;
-    _actOnWeightBreach = actOnWeightBreach;
+    
+    if (actOnWeightBreach != null)
+    {
+      _actOnWeightBreach = actOnWeightBreach;
+    }
 
     if (server instanceof ZooKeeperServer)
     {

From e719789fb0e8206e2d04e38431b7ade2225468f1 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Mon, 28 Oct 2024 15:18:04 -0700
Subject: [PATCH 05/12] adjust param type

---
 .../com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index 949f2c2d1..d6b470f7a 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -230,7 +230,7 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
 
   public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
       boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService,
-      ServiceDiscoveryEventEmitter eventEmitter, BigDecimal maxWeight, ActOnWeightBreach actOnWeightBreach)
+      ServiceDiscoveryEventEmitter eventEmitter, String maxWeight, ActOnWeightBreach actOnWeightBreach)
   {
     _server = server;
     // initialIsUp is used for delay mark up. If it's false, there won't be markup when the announcer is started.
@@ -246,8 +246,7 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
     _warmupDuration = warmupDuration;
     _executorService = executorService;
     _eventEmitter = eventEmitter;
-    _maxWeight = maxWeight;
-    
+    _maxWeight = maxWeight == null ? null : new BigDecimal(maxWeight);
     if (actOnWeightBreach != null)
     {
       _actOnWeightBreach = actOnWeightBreach;

From 6b0f02de0d12afafc0f52d5d9a231f8f372ade47 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Tue, 29 Oct 2024 12:14:58 -0700
Subject: [PATCH 06/12] fix edge case that breaches both max and decimal places

---
 .../com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index d6b470f7a..8ff6e8564 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -246,6 +246,8 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
     _warmupDuration = warmupDuration;
     _executorService = executorService;
     _eventEmitter = eventEmitter;
+
+    // take in max weight as a string to be precise at decimal places when creating a BigDecimal.
     _maxWeight = maxWeight == null ? null : new BigDecimal(maxWeight);
     if (actOnWeightBreach != null)
     {
@@ -873,6 +875,7 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
             throw getMaxWeightBreachException(weight, entry.getKey());
           case RECTIFY:
             entry.setValue(new PartitionData(_maxWeight.intValue()));
+            weight = _maxWeight;
             _log.warn("Capped weight {} in Partition {} to the max weight allowed: {}.", weight.doubleValue(),
                 entry.getKey(), _maxWeight.intValue());
             break;

From e333b9cda9de45c3a28284ac699b12ba534e4684 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Wed, 30 Oct 2024 15:01:01 -0700
Subject: [PATCH 07/12] add jmx methods for weight breach count and markup
 status

---
 .../balancer/servers/ZooKeeperAnnouncer.java  | 44 ++++++++++++++++++-
 .../d2/jmx/ZooKeeperAnnouncerJmx.java         | 23 ++++++++++
 .../d2/jmx/ZooKeeperAnnouncerJmxMXBean.java   | 21 +++++++++
 3 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index 8ff6e8564..056f96be5 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -46,6 +46,8 @@
 import com.linkedin.util.ArgumentUtil;
 
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.Nullable;
@@ -106,6 +108,10 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
    * The action to take when d2 weight breaches validation rules.
    */
   private ActOnWeightBreach _actOnWeightBreach = ActOnWeightBreach.IGNORE;
+
+  private final AtomicInteger _maxWeightBreachedCount = new AtomicInteger(0);
+  private final AtomicInteger _weightDecimalPlacesBreachedCount = new AtomicInteger(0);
+
   private volatile Map<String, Object> _uriSpecificProperties;
 
   private ServiceDiscoveryEventEmitter _eventEmitter;
@@ -115,6 +121,10 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
    * it will try to bring up the server again on ZK if the connection goes down, or a new store is set
    */
   private boolean _isUp;
+  /**
+   * Whether the current announcement status is up. It's up only if the announcement completed successfully.
+   */
+  private final AtomicBoolean _isMarkedUp = new AtomicBoolean(false);
 
   // Field to indicate if warm up was started. If it is true, it will try to end the warm up
   // by marking down on ZK if the connection goes down
@@ -140,6 +150,10 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
 
   // Boolean flag to indicate if dark warm-up is enabled, defaults to false
   private final boolean _isDarkWarmupEnabled;
+  /**
+   * Whether the dark warmup cluster announcement status is up. It's up only if the announcement completed successfully.
+   */
+  private final AtomicBoolean _isDarkWarmupMarkedUp = new AtomicBoolean(false);
 
   // String to store the name of the dark warm-up cluster, defaults to null
   private final String _warmupClusterName;
@@ -377,6 +391,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
+        _isMarkedUp.set(true);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_cluster, true, true, _markUpStartAtRef.get());
         _markUpFailed = false;
         _log.info("markUp for uri = {} on cluster {} succeeded.", _uri, _cluster);
@@ -436,6 +451,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
+        _isDarkWarmupMarkedUp.set(false);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_warmupClusterName, false, true, _warmupClusterMarkDownStartAtRef.get());
         // Mark _isWarmingUp to false to indicate warm up has completed
         _isWarmingUp = false;
@@ -486,6 +502,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
+        _isDarkWarmupMarkedUp.set(true);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_warmupClusterName, true, true, _warmupClusterMarkUpStartAtRef.get());
         _log.info("markUp for uri {} on warm-up cluster {} succeeded", _uri, _warmupClusterName);
         // Mark _isWarmingUp to true to indicate warm up is in progress
@@ -560,6 +577,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
+        _isMarkedUp.set(false);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_cluster, false, true, _markDownStartAtRef.get());
         _log.info("markDown for uri = {} succeeded.", _uri);
         // Note that the pending callbacks we see at this point are
@@ -797,6 +815,26 @@ public boolean isMarkUpFailed()
     return _markUpFailed;
   }
 
+  public boolean isMarkedUp()
+  {
+    return _isMarkedUp.get();
+  }
+
+  public boolean isDarkWarmupMarkedUp()
+  {
+    return _isDarkWarmupMarkedUp.get();
+  }
+
+  public int getMaxWeightBreachedCount()
+  {
+    return _maxWeightBreachedCount.get();
+  }
+
+  public int getWeightDecimalPlacesBreachedCount()
+  {
+    return _weightDecimalPlacesBreachedCount.get();
+  }
+
   public void setEventEmitter(ServiceDiscoveryEventEmitter emitter) {
     _eventEmitter = emitter;
   }
@@ -858,7 +896,7 @@ public boolean isWarmingUp() {
   }
 
   private Map<Integer, PartitionData> validatePartitionData(Map<Integer, PartitionData> partitionData) {
-    if (_maxWeight == null || _actOnWeightBreach == ActOnWeightBreach.IGNORE) {
+    if (_maxWeight == null) {
       return partitionData;
     }
 
@@ -867,6 +905,7 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
       BigDecimal weight = BigDecimal.valueOf(entry.getValue().getWeight());
       // check max weight
       if (weight.compareTo(_maxWeight) > 0) {
+        _maxWeightBreachedCount.incrementAndGet();
         switch (_actOnWeightBreach) {
           case WARN:
             _log.warn("", getMaxWeightBreachException(weight, entry.getKey()));
@@ -879,6 +918,7 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
             _log.warn("Capped weight {} in Partition {} to the max weight allowed: {}.", weight.doubleValue(),
                 entry.getKey(), _maxWeight.intValue());
             break;
+          case IGNORE:
           default:
             break;
         }
@@ -886,6 +926,7 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
 
       // check decimal places
       if (weight.scale() > _maxWeight.scale()) {
+        _weightDecimalPlacesBreachedCount.incrementAndGet();
         switch (_actOnWeightBreach) {
           case WARN: // both WARN and THROW only log the warning. Don't throw exception for decimal places.
           case THROW:
@@ -899,6 +940,7 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
             _log.warn("Rounded weight {} in Partition {} to {} decimal places: {}.", weight.doubleValue(),
                 entry.getKey(), _maxWeight.scale(), newWeight);
             break;
+          case IGNORE:
           default:
             break;
         }
diff --git a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java
index f10211e35..351c96159 100644
--- a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java
+++ b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java
@@ -183,4 +183,27 @@ public Map<Integer, PartitionData> getPartitionData()
   public boolean isMarkUpFailed() {
     return _announcer.isMarkUpFailed();
   }
+
+  @Override
+  public boolean isMarkedUp()
+  {
+    return _announcer.isMarkedUp();
+  }
+
+  @Override
+  public boolean isDarkWarmupMarkedUp() {
+    return _announcer.isDarkWarmupMarkedUp();
+  }
+
+  @Override
+  public int getMaxWeightBreachedCount()
+  {
+    return _announcer.getMaxWeightBreachedCount();
+  }
+
+  @Override
+  public int getWeightDecimalPlacesBreachedCount()
+  {
+    return _announcer.getWeightDecimalPlacesBreachedCount();
+  }
 }
diff --git a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java
index 53c85c079..0289a3bc7 100644
--- a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java
+++ b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java
@@ -72,4 +72,25 @@ void setPartitionDataUsingJson(String partitionDataJson)
   void setPartitionData(Map<Integer, PartitionData> partitionData);
 
   boolean isMarkUpFailed();
+
+  /**
+   * @return true if the announcer is marked up.
+   */
+  boolean isMarkedUp();
+
+  /**
+   * @return true if the announcer marked up the dark warmup cluster.
+   */
+  boolean isDarkWarmupMarkedUp();
+
+  /**
+   * @return the times that the max weight has been breached.
+   */
+  int getMaxWeightBreachedCount();
+
+  /**
+   *
+   * @return the times that the max number of decimal places on weight have been breached.
+   */
+  int getWeightDecimalPlacesBreachedCount();
 }

From 6dcce29f0efdc776ebf6a968ff753725de114662 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Thu, 31 Oct 2024 15:44:54 -0700
Subject: [PATCH 08/12] adjust var names and types

---
 .../balancer/servers/ZooKeeperAnnouncer.java  | 74 ++++++++-------
 .../servers/TestZooKeeperAnnouncer.java       | 90 +++++++++++++++++++
 2 files changed, 132 insertions(+), 32 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index 056f96be5..a03cf5e28 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -20,6 +20,7 @@
 
 package com.linkedin.d2.balancer.servers;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.math.BigDecimal;
 import java.math.RoundingMode;
 import java.net.URI;
@@ -102,12 +103,14 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
    * 1. The maximum d2 weight allowed.
    * 2. The maximum number of decimal places allowed. Use 0s on decimal places to indicate it.
    * For example, 100.00 means the max weight allowed is 100 and the max number of decimal places is 2.
+   * CAUTION: BigDecimal yields accurate scale when constructed with a string of the number, instead of a double/float.
+   * E.g: new BigDecimal("100.00") instead of new BigDecimal(100.00).
    */
   private final BigDecimal _maxWeight;
   /**
    * The action to take when d2 weight breaches validation rules.
    */
-  private ActOnWeightBreach _actOnWeightBreach = ActOnWeightBreach.IGNORE;
+  private ActionOnWeightBreach _actionOnWeightBreach = ActionOnWeightBreach.IGNORE;
 
   private final AtomicInteger _maxWeightBreachedCount = new AtomicInteger(0);
   private final AtomicInteger _weightDecimalPlacesBreachedCount = new AtomicInteger(0);
@@ -122,9 +125,11 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
    */
   private boolean _isUp;
   /**
-   * Whether the current announcement status is up. It's up only if the announcement completed successfully.
+   * Whether the announcer has completed sending a markup intent. NOTE THAT a mark-up intent sent does NOT mean the
+   * announcement status on service discovery registry is up. Service discovery registry may further process the host
+   * and determine its status. Check on service discovery registry for the final status.
    */
-  private final AtomicBoolean _isMarkedUp = new AtomicBoolean(false);
+  private final AtomicBoolean _isMarkUpIntentSent = new AtomicBoolean(false);
 
   // Field to indicate if warm up was started. If it is true, it will try to end the warm up
   // by marking down on ZK if the connection goes down
@@ -151,7 +156,7 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   // Boolean flag to indicate if dark warm-up is enabled, defaults to false
   private final boolean _isDarkWarmupEnabled;
   /**
-   * Whether the dark warmup cluster announcement status is up. It's up only if the announcement completed successfully.
+   * Whether the announcer has completed sending a dark warmup cluster markup intent.
    */
   private final AtomicBoolean _isDarkWarmupMarkedUp = new AtomicBoolean(false);
 
@@ -168,7 +173,7 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   // Field to store the dark warm-up time duration in seconds, defaults to zero
   private final int _warmupDuration;
 
-  public enum ActOnWeightBreach {
+  public enum ActionOnWeightBreach {
     // Ignore and no op.
     IGNORE,
     // only log warnings
@@ -233,18 +238,18 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
   public ZooKeeperAnnouncer(ZooKeeperServer server, boolean initialIsUp,
                             boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService, ServiceDiscoveryEventEmitter eventEmitter)
   {
-    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null, ActOnWeightBreach.IGNORE);
+    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null, ActionOnWeightBreach.IGNORE);
   }
 
   public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
       boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService, ServiceDiscoveryEventEmitter eventEmitter)
   {
-    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null, ActOnWeightBreach.IGNORE);
+    this(server, initialIsUp, isDarkWarmupEnabled, warmupClusterName, warmupDuration, executorService, eventEmitter, null, ActionOnWeightBreach.IGNORE);
   }
 
   public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
       boolean isDarkWarmupEnabled, String warmupClusterName, int warmupDuration, ScheduledExecutorService executorService,
-      ServiceDiscoveryEventEmitter eventEmitter, String maxWeight, ActOnWeightBreach actOnWeightBreach)
+      ServiceDiscoveryEventEmitter eventEmitter, BigDecimal maxWeight, ActionOnWeightBreach actionOnWeightBreach)
   {
     _server = server;
     // initialIsUp is used for delay mark up. If it's false, there won't be markup when the announcer is started.
@@ -261,11 +266,10 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
     _executorService = executorService;
     _eventEmitter = eventEmitter;
 
-    // take in max weight as a string to be precise at decimal places when creating a BigDecimal.
-    _maxWeight = maxWeight == null ? null : new BigDecimal(maxWeight);
-    if (actOnWeightBreach != null)
+    _maxWeight = maxWeight;
+    if (actionOnWeightBreach != null)
     {
-      _actOnWeightBreach = actOnWeightBreach;
+      _actionOnWeightBreach = actionOnWeightBreach;
     }
 
     if (server instanceof ZooKeeperServer)
@@ -391,7 +395,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
-        _isMarkedUp.set(true);
+        _isMarkUpIntentSent.set(true);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_cluster, true, true, _markUpStartAtRef.get());
         _markUpFailed = false;
         _log.info("markUp for uri = {} on cluster {} succeeded.", _uri, _cluster);
@@ -577,7 +581,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
-        _isMarkedUp.set(false);
+        _isMarkUpIntentSent.set(false);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_cluster, false, true, _markDownStartAtRef.get());
         _log.info("markDown for uri = {} succeeded.", _uri);
         // Note that the pending callbacks we see at this point are
@@ -817,7 +821,7 @@ public boolean isMarkUpFailed()
 
   public boolean isMarkedUp()
   {
-    return _isMarkedUp.get();
+    return _isMarkUpIntentSent.get();
   }
 
   public boolean isDarkWarmupMarkedUp()
@@ -895,18 +899,25 @@ public boolean isWarmingUp() {
     return _isWarmingUp;
   }
 
-  private Map<Integer, PartitionData> validatePartitionData(Map<Integer, PartitionData> partitionData) {
-    if (_maxWeight == null) {
-      return partitionData;
-    }
-
+  @VisibleForTesting
+  Map<Integer, PartitionData> validatePartitionData(Map<Integer, PartitionData> partitionData) {
     Map<Integer, PartitionData> res = new HashMap<>(partitionData); // modifiable copy in case the input is unmodifiable
     for (Map.Entry<Integer, PartitionData> entry : res.entrySet()) {
       BigDecimal weight = BigDecimal.valueOf(entry.getValue().getWeight());
+      // check negative weight
+      if (weight.compareTo(BigDecimal.ZERO) < 0) {
+        throw new IllegalArgumentException(String.format("Weight %s in Partition %d is negative. Please correct it.",
+            weight, entry.getKey()));
+      }
+
+      if (_maxWeight == null) {
+        break;
+      }
+
       // check max weight
       if (weight.compareTo(_maxWeight) > 0) {
         _maxWeightBreachedCount.incrementAndGet();
-        switch (_actOnWeightBreach) {
+        switch (_actionOnWeightBreach) {
           case WARN:
             _log.warn("", getMaxWeightBreachException(weight, entry.getKey()));
             break;
@@ -915,8 +926,8 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
           case RECTIFY:
             entry.setValue(new PartitionData(_maxWeight.intValue()));
             weight = _maxWeight;
-            _log.warn("Capped weight {} in Partition {} to the max weight allowed: {}.", weight.doubleValue(),
-                entry.getKey(), _maxWeight.intValue());
+            _log.warn("Capped weight {} in Partition {} to the max weight allowed: {}.", weight, entry.getKey(),
+                _maxWeight);
             break;
           case IGNORE:
           default:
@@ -927,18 +938,17 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
       // check decimal places
       if (weight.scale() > _maxWeight.scale()) {
         _weightDecimalPlacesBreachedCount.incrementAndGet();
-        switch (_actOnWeightBreach) {
+        switch (_actionOnWeightBreach) {
           case WARN: // both WARN and THROW only log the warning. Don't throw exception for decimal places.
           case THROW:
-            _log.warn("", new IllegalArgumentException(String.format("Weight %f in Partition %d has more than %d"
-                + " decimal places. It will be rounded in the future.", weight.doubleValue(), entry.getKey(),
-                _maxWeight.scale())));
+            _log.warn("", new IllegalArgumentException(String.format("Weight %s in Partition %d has more than %d"
+                + " decimal places. It will be rounded in the future.", weight, entry.getKey(), _maxWeight.scale())));
             break;
           case RECTIFY:
             double newWeight = weight.setScale(_maxWeight.scale(), RoundingMode.HALF_UP).doubleValue();
             entry.setValue(new PartitionData(newWeight));
-            _log.warn("Rounded weight {} in Partition {} to {} decimal places: {}.", weight.doubleValue(),
-                entry.getKey(), _maxWeight.scale(), newWeight);
+            _log.warn("Rounded weight {} in Partition {} to {} decimal places: {}.", weight, entry.getKey(),
+                _maxWeight.scale(), newWeight);
             break;
           case IGNORE:
           default:
@@ -950,8 +960,8 @@ private Map<Integer, PartitionData> validatePartitionData(Map<Integer, Partition
   }
 
   private IllegalArgumentException getMaxWeightBreachException(BigDecimal weight, int partition) {
-    return new IllegalArgumentException(String.format("[ACTION NEEDED] Weight %f in Partition %d is greater"
-        + " than the max weight allowed: %d. Please correct the weight. It will be force-capped to the max weight "
-        + "in the future.", weight.doubleValue(), partition, _maxWeight.intValue()));
+    return new IllegalArgumentException(String.format("[ACTION NEEDED] Weight %s in Partition %d is greater"
+        + " than the max weight allowed: %s. Please correct the weight. It will be force-capped to the max weight "
+        + "in the future.", weight, partition, _maxWeight));
   }
 }
diff --git a/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java b/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
index 00ade0b0b..954ec18b4 100644
--- a/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
+++ b/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
@@ -2,13 +2,20 @@
 
 import com.linkedin.common.callback.Callback;
 import com.linkedin.common.util.None;
+import com.linkedin.d2.balancer.properties.PartitionData;
 import com.linkedin.d2.balancer.properties.PropertyKeys;
 
+import com.linkedin.d2.discovery.event.LogOnlyServiceDiscoveryEventEmitter;
+import java.math.BigDecimal;
+import java.util.Collections;
+import java.util.Map;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import static org.junit.Assert.*;
 import static org.mockito.Mockito.*;
 
 
@@ -24,6 +31,13 @@ public class TestZooKeeperAnnouncer
   @Mock
   private Callback<None> _callback;
 
+  private static final Map<Integer, PartitionData> MAX_WEIGHT_BREACH_PARTITION_DATA =
+      Collections.singletonMap(0, new PartitionData(1000));
+  private static final Map<Integer, PartitionData> DECIMAL_PLACES_BREACH_PARTITION_DATA =
+      Collections.singletonMap(0, new PartitionData(5.345));
+  private static final Map<Integer, PartitionData> MAX_WEIGHT_AND_DECIMAL_PLACES_BREACH_PARTITION_DATA =
+      Collections.singletonMap(0, new PartitionData(10.89));
+
   @BeforeMethod
   public void setUp()
   {
@@ -43,4 +57,80 @@ public void testSetDoNotLoadBalance()
 
     verify(_server).addUriSpecificProperty(any(), any(), any(), any(), eq(PropertyKeys.DO_NOT_LOAD_BALANCE), eq(false), any());
   }
+
+  @DataProvider(name = "validatePartitionDataDataProvider")
+  public Object[][] getValidatePartitionDataDataProvider()
+  {
+    return new Object[][] {
+        {
+          // no weight rules
+          null, null, MAX_WEIGHT_BREACH_PARTITION_DATA, MAX_WEIGHT_BREACH_PARTITION_DATA, null, 0, 0
+        },
+        {
+          // no action default to IGNORE, which won't correct the value BUT will increment the counts
+          "10.0", null, MAX_WEIGHT_BREACH_PARTITION_DATA, MAX_WEIGHT_BREACH_PARTITION_DATA, null, 1, 0
+        },
+        {
+          // warn action won't correct the value
+          "10.0", ZooKeeperAnnouncer.ActionOnWeightBreach.WARN, MAX_WEIGHT_BREACH_PARTITION_DATA,
+            MAX_WEIGHT_BREACH_PARTITION_DATA, null, 1, 0
+        },
+        {
+          // max weight breach, correct the value
+          "10.0", ZooKeeperAnnouncer.ActionOnWeightBreach.RECTIFY, MAX_WEIGHT_BREACH_PARTITION_DATA,
+            Collections.singletonMap(0, new PartitionData(10)), null, 1, 0
+        },
+        {
+          // decimal places breach, correct the value
+          "10.0", ZooKeeperAnnouncer.ActionOnWeightBreach.RECTIFY, DECIMAL_PLACES_BREACH_PARTITION_DATA,
+            Collections.singletonMap(0, new PartitionData(5.3)), null, 0, 1
+        },
+        {
+          // max weight and decimal places breach, correct the value
+          "10.0", ZooKeeperAnnouncer.ActionOnWeightBreach.RECTIFY, MAX_WEIGHT_AND_DECIMAL_PLACES_BREACH_PARTITION_DATA,
+            Collections.singletonMap(0, new PartitionData(10)), null, 1, 0
+        },
+        {
+          // throw action throws for max weight breach
+          "10.0", ZooKeeperAnnouncer.ActionOnWeightBreach.THROW, MAX_WEIGHT_BREACH_PARTITION_DATA, null,
+            new IllegalArgumentException("[ACTION NEEDED] Weight 1000.0 in Partition 0 is greater than the max weight "
+                + "allowed: 10.0. Please correct the weight. It will be force-capped to the max weight in the future."),
+            1, 0
+        },
+        {
+          // throw action doesn not throw for decimal places breach
+          "10.0", ZooKeeperAnnouncer.ActionOnWeightBreach.THROW, DECIMAL_PLACES_BREACH_PARTITION_DATA,
+            DECIMAL_PLACES_BREACH_PARTITION_DATA, null, 0, 1
+        }
+    };
+  }
+  @Test(dataProvider = "validatePartitionDataDataProvider")
+  public void testValidatePartitionData(String maxWeight, ZooKeeperAnnouncer.ActionOnWeightBreach action,
+      Map<Integer, PartitionData> input, Map<Integer, PartitionData> expected, Exception expectedException,
+      int expectedMaxWeightBreachedCount, int expectedWeightDecimalPlacesBreachedCount)
+  {
+    ZooKeeperAnnouncer announcer = new ZooKeeperAnnouncer(_server, true, false, null, 0,
+        null, new LogOnlyServiceDiscoveryEventEmitter(),
+        maxWeight == null ? null : new BigDecimal(maxWeight), action);
+
+    if (expectedException != null)
+    {
+      try
+      {
+        announcer.validatePartitionData(input);
+        fail("Expected exception not thrown");
+      }
+      catch (Exception ex)
+      {
+        assertTrue(ex instanceof IllegalArgumentException);
+        assertEquals(expectedException.getMessage(), ex.getMessage());
+      }
+    }
+    else
+    {
+      assertEquals(expected, announcer.validatePartitionData(input));
+    }
+    assertEquals(expectedMaxWeightBreachedCount, announcer.getMaxWeightBreachedCount());
+    assertEquals(expectedWeightDecimalPlacesBreachedCount, announcer.getWeightDecimalPlacesBreachedCount());
+  }
 }

From 8fb85a8d8679f6b4d2a48b7544626c717c5523a7 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Thu, 31 Oct 2024 15:57:39 -0700
Subject: [PATCH 09/12] add test case of negative weight

---
 .../d2/balancer/servers/TestZooKeeperAnnouncer.java        | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java b/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
index 954ec18b4..fdca9f452 100644
--- a/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
+++ b/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
@@ -66,6 +66,11 @@ public Object[][] getValidatePartitionDataDataProvider()
           // no weight rules
           null, null, MAX_WEIGHT_BREACH_PARTITION_DATA, MAX_WEIGHT_BREACH_PARTITION_DATA, null, 0, 0
         },
+        {
+          // negative weight throws
+          null, null, Collections.singletonMap(0, new PartitionData(-1.0)), null,
+            new IllegalArgumentException("Weight -1.0 in Partition 0 is negative. Please correct it."), 0, 0
+        },
         {
           // no action default to IGNORE, which won't correct the value BUT will increment the counts
           "10.0", null, MAX_WEIGHT_BREACH_PARTITION_DATA, MAX_WEIGHT_BREACH_PARTITION_DATA, null, 1, 0
@@ -98,7 +103,7 @@ public Object[][] getValidatePartitionDataDataProvider()
             1, 0
         },
         {
-          // throw action doesn not throw for decimal places breach
+          // throw action does not throw for decimal places breach
           "10.0", ZooKeeperAnnouncer.ActionOnWeightBreach.THROW, DECIMAL_PLACES_BREACH_PARTITION_DATA,
             DECIMAL_PLACES_BREACH_PARTITION_DATA, null, 0, 1
         }

From bdba636dde2e7f49b741f9c05eed000dfca7d585 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Thu, 31 Oct 2024 16:14:19 -0700
Subject: [PATCH 10/12] update jmx and mbean method names

---
 .../d2/balancer/servers/ZooKeeperAnnouncer.java      | 12 ++++++------
 .../com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java   |  8 ++++----
 .../linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java | 12 +++++++-----
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index a03cf5e28..315727a6d 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -158,7 +158,7 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   /**
    * Whether the announcer has completed sending a dark warmup cluster markup intent.
    */
-  private final AtomicBoolean _isDarkWarmupMarkedUp = new AtomicBoolean(false);
+  private final AtomicBoolean _isDarkWarmupMarkUpIntentSent = new AtomicBoolean(false);
 
   // String to store the name of the dark warm-up cluster, defaults to null
   private final String _warmupClusterName;
@@ -455,7 +455,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
-        _isDarkWarmupMarkedUp.set(false);
+        _isDarkWarmupMarkUpIntentSent.set(false);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_warmupClusterName, false, true, _warmupClusterMarkDownStartAtRef.get());
         // Mark _isWarmingUp to false to indicate warm up has completed
         _isWarmingUp = false;
@@ -506,7 +506,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
-        _isDarkWarmupMarkedUp.set(true);
+        _isDarkWarmupMarkUpIntentSent.set(true);
         emitSDStatusActiveUpdateIntentAndWriteEvents(_warmupClusterName, true, true, _warmupClusterMarkUpStartAtRef.get());
         _log.info("markUp for uri {} on warm-up cluster {} succeeded", _uri, _warmupClusterName);
         // Mark _isWarmingUp to true to indicate warm up is in progress
@@ -819,14 +819,14 @@ public boolean isMarkUpFailed()
     return _markUpFailed;
   }
 
-  public boolean isMarkedUp()
+  public boolean isMarkUpIntentSent()
   {
     return _isMarkUpIntentSent.get();
   }
 
-  public boolean isDarkWarmupMarkedUp()
+  public boolean isDarkWarmupMarkUpIntentSent()
   {
-    return _isDarkWarmupMarkedUp.get();
+    return _isDarkWarmupMarkUpIntentSent.get();
   }
 
   public int getMaxWeightBreachedCount()
diff --git a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java
index 351c96159..0fcbe33bf 100644
--- a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java
+++ b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmx.java
@@ -185,14 +185,14 @@ public boolean isMarkUpFailed() {
   }
 
   @Override
-  public boolean isMarkedUp()
+  public boolean isMarkUpIntentSent()
   {
-    return _announcer.isMarkedUp();
+    return _announcer.isMarkUpIntentSent();
   }
 
   @Override
-  public boolean isDarkWarmupMarkedUp() {
-    return _announcer.isDarkWarmupMarkedUp();
+  public boolean isDarkWarmupMarkUpIntentSent() {
+    return _announcer.isDarkWarmupMarkUpIntentSent();
   }
 
   @Override
diff --git a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java
index 0289a3bc7..70cd1c716 100644
--- a/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java
+++ b/d2/src/main/java/com/linkedin/d2/jmx/ZooKeeperAnnouncerJmxMXBean.java
@@ -74,14 +74,16 @@ void setPartitionDataUsingJson(String partitionDataJson)
   boolean isMarkUpFailed();
 
   /**
-   * @return true if the announcer is marked up.
+   * @return true if the announcer has completed sending a markup intent. NOTE THAT a mark-up intent sent does NOT mean the
+   * announcement status on service discovery registry is up. Service discovery registry may further process the host
+   * and determine its status. Check on service discovery registry for the final status.
    */
-  boolean isMarkedUp();
+  boolean isMarkUpIntentSent();
 
   /**
-   * @return true if the announcer marked up the dark warmup cluster.
+   * @return true if the announcer has completed sending a dark warmup cluster markup intent.
    */
-  boolean isDarkWarmupMarkedUp();
+  boolean isDarkWarmupMarkUpIntentSent();
 
   /**
    * @return the times that the max weight has been breached.
@@ -90,7 +92,7 @@ void setPartitionDataUsingJson(String partitionDataJson)
 
   /**
    *
-   * @return the times that the max number of decimal places on weight have been breached.
+   * @return the times that the max number of decimal places on weight has been breached.
    */
   int getWeightDecimalPlacesBreachedCount();
 }

From c5801fb798927668aa700fb104d165cc17f9c951 Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Mon, 4 Nov 2024 10:15:53 -0800
Subject: [PATCH 11/12] add a test case of valid weight

---
 .../d2/balancer/servers/TestZooKeeperAnnouncer.java         | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java b/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
index fdca9f452..2af6532e1 100644
--- a/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
+++ b/d2/src/test/java/com/linkedin/d2/balancer/servers/TestZooKeeperAnnouncer.java
@@ -37,6 +37,8 @@ public class TestZooKeeperAnnouncer
       Collections.singletonMap(0, new PartitionData(5.345));
   private static final Map<Integer, PartitionData> MAX_WEIGHT_AND_DECIMAL_PLACES_BREACH_PARTITION_DATA =
       Collections.singletonMap(0, new PartitionData(10.89));
+  private static final Map<Integer, PartitionData> VALID_PARTITION_DATA =
+      Collections.singletonMap(0, new PartitionData(2.3));
 
   @BeforeMethod
   public void setUp()
@@ -71,6 +73,10 @@ public Object[][] getValidatePartitionDataDataProvider()
           null, null, Collections.singletonMap(0, new PartitionData(-1.0)), null,
             new IllegalArgumentException("Weight -1.0 in Partition 0 is negative. Please correct it."), 0, 0
         },
+        {
+          // valid weight
+          "3.0", null, VALID_PARTITION_DATA, VALID_PARTITION_DATA, null, 0, 0
+        },
         {
           // no action default to IGNORE, which won't correct the value BUT will increment the counts
           "10.0", null, MAX_WEIGHT_BREACH_PARTITION_DATA, MAX_WEIGHT_BREACH_PARTITION_DATA, null, 1, 0

From bbb1c20b8458769e0299db28cd3715896937035d Mon Sep 17 00:00:00 2001
From: Bohan Yang <boyang@linkedin.com>
Date: Mon, 4 Nov 2024 15:27:13 -0800
Subject: [PATCH 12/12] make var final

---
 .../linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java   | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
index 315727a6d..78fbfb0b4 100644
--- a/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
+++ b/d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java
@@ -110,7 +110,7 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
   /**
    * The action to take when d2 weight breaches validation rules.
    */
-  private ActionOnWeightBreach _actionOnWeightBreach = ActionOnWeightBreach.IGNORE;
+  private final ActionOnWeightBreach _actionOnWeightBreach;
 
   private final AtomicInteger _maxWeightBreachedCount = new AtomicInteger(0);
   private final AtomicInteger _weightDecimalPlacesBreachedCount = new AtomicInteger(0);
@@ -267,10 +267,7 @@ public ZooKeeperAnnouncer(LoadBalancerServer server, boolean initialIsUp,
     _eventEmitter = eventEmitter;
 
     _maxWeight = maxWeight;
-    if (actionOnWeightBreach != null)
-    {
-      _actionOnWeightBreach = actionOnWeightBreach;
-    }
+    _actionOnWeightBreach = actionOnWeightBreach != null ? actionOnWeightBreach : ActionOnWeightBreach.IGNORE;
 
     if (server instanceof ZooKeeperServer)
     {