From 535967373bfce88df802ca1e2fbbd85c3186685b Mon Sep 17 00:00:00 2001
From: Charanya Sudharsanan <csudhars@csudhars-mn2.linkedin.biz>
Date: Thu, 30 May 2024 11:07:32 -0700
Subject: [PATCH] [apache/helix] --> Update logic for metric calculation when
 replica is set to ANY_LIVEINSTANCE

---
 .../monitoring/mbeans/ResourceMonitor.java    | 22 ++++++++++++++-----
 .../mbeans/TestResourceMonitor.java           | 21 ++++++++++++++++++
 2 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/monitoring/mbeans/ResourceMonitor.java b/helix-core/src/main/java/org/apache/helix/monitoring/mbeans/ResourceMonitor.java
index 5064c64812..70a4ffe2d7 100644
--- a/helix-core/src/main/java/org/apache/helix/monitoring/mbeans/ResourceMonitor.java
+++ b/helix-core/src/main/java/org/apache/helix/monitoring/mbeans/ResourceMonitor.java
@@ -35,6 +35,7 @@
 import org.apache.helix.HelixDefinedState;
 import org.apache.helix.model.ExternalView;
 import org.apache.helix.model.IdealState;
+import org.apache.helix.model.ResourceConfig;
 import org.apache.helix.model.StateModelDefinition;
 import org.apache.helix.monitoring.mbeans.dynamicMBeans.DynamicMBeanProvider;
 import org.apache.helix.monitoring.mbeans.dynamicMBeans.DynamicMetric;
@@ -274,13 +275,22 @@ public void updateResourceState(ExternalView externalView, IdealState idealState
     Set<String> partitions = idealState.getPartitionSet();
     int replica;
     try {
-      replica = Integer.valueOf(idealState.getReplicas());
-    } catch (NumberFormatException e) {
-      _logger.info("Unspecified replica count for {}, skip updating the ResourceMonitor Mbean: {}", _resourceName,
-          idealState.getReplicas());
-      return;
+      String replicaCount = idealState.getReplicas();
+      if (!replicaCount.equals(ResourceConfig.ResourceConfigConstants.ANY_LIVEINSTANCE.name())) {
+        try {
+          replica = Integer.parseInt(replicaCount);
+        } catch (NumberFormatException ex) {
+          _logger.info(
+              "Unspecified replica count for {}, skip updating the ResourceMonitor Mbean: {}",
+              _resourceName, idealState.getReplicas());
+          return;
+        }
+      } else {
+        replica = -1;
+      }
     } catch (Exception ex) {
-      _logger.warn("Failed to get replica count for {}, cannot update the ResourceMonitor Mbean.", _resourceName);
+      _logger.warn("Failed to get replica count for {}, cannot update the ResourceMonitor Mbean.",
+          _resourceName);
       return;
     }
 
diff --git a/helix-core/src/test/java/org/apache/helix/monitoring/mbeans/TestResourceMonitor.java b/helix-core/src/test/java/org/apache/helix/monitoring/mbeans/TestResourceMonitor.java
index 0fb0e09371..355cad4501 100644
--- a/helix-core/src/test/java/org/apache/helix/monitoring/mbeans/TestResourceMonitor.java
+++ b/helix-core/src/test/java/org/apache/helix/monitoring/mbeans/TestResourceMonitor.java
@@ -37,6 +37,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import org.apache.helix.TestHelper;
+import org.apache.helix.model.ResourceConfig;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.helix.model.BuiltInStateModelDefinitions;
 import org.apache.helix.model.ExternalView;
@@ -220,6 +221,26 @@ public void testReportData() throws JMException {
       monitor.setRebalanceState(ResourceMonitor.RebalanceStatus.INTERMEDIATE_STATE_CAL_FAILED);
       Assert.assertEquals(monitor.getRebalanceState(),
           ResourceMonitor.RebalanceStatus.INTERMEDIATE_STATE_CAL_FAILED.name());
+
+      // test when replica is set to ANY_LIVEINSTANCE and all instances are taken offline.
+      idealState.setReplicas(ResourceConfig.ResourceConfigConstants.ANY_LIVEINSTANCE.name());
+
+      for (int i = 0; i < _partitions; i++) {
+        String partition = _dbName + "_" + i;
+        Map<String, String> externalViewStateMap = externalView.getStateMap(partition);
+        for (String key : externalViewStateMap.keySet()) {
+          if (externalViewStateMap.get(key).equalsIgnoreCase("MASTER")) {
+            externalViewStateMap.put(key, "OFFLINE");
+          }
+        }
+        externalView.setStateMap(partition, externalViewStateMap);
+      }
+
+      monitor.updateResourceState(externalView, idealState, stateModelDef);
+
+      Assert.assertEquals(monitor.getMissingTopStatePartitionGauge(), _partitions);
+      Assert.assertEquals(monitor.getMissingReplicaPartitionGauge(), 0);
+      Assert.assertEquals(monitor.getMissingMinActiveReplicaPartitionGauge(), 0);
     } finally {
       // Has to unregister this monitor to clean up. Otherwise, later tests may be affected and fail.
       monitor.unregister();