Add endpoint for Helix-computed offline instance count

helix-core/src/main/java/org/apache/helix/controller/dataproviders/BaseControllerDataProvider.java

@@ -64,6 +64,7 @@
 import org.apache.helix.model.StateModelDefinition;
 import org.apache.helix.task.TaskConstants;
 import org.apache.helix.util.HelixUtil;
+import org.apache.helix.util.InstanceUtil;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.helix.zookeeper.zkclient.DataUpdater;
 import org.slf4j.Logger;
@@ -759,16 +760,10 @@ public Set<String> getEnabledInstances() {
    * auto Maintenance Mode (MAX_OFFLINE_INSTANCES_ALLOWED at entry,
    * NUM_OFFLINE_INSTANCES_FOR_AUTO_EXIT at exit).
    *
-   * <p>An instance counts when all three conditions hold:
-   * <ul>
-   *   <li>Its InstanceOperation is routable, i.e. not in
-   *       {@link InstanceConstants#UNROUTABLE_INSTANCE_OPERATIONS}.</li>
-   *   <li>It is not currently enabled-and-live.</li>
-   *   <li>It does not carry a valid (unexpired) instance-operation maintenance marker.</li>
-   * </ul>
-   * EVACUATE and DISABLE instances are included because they cannot accept new ONLINE
-   * replicas; ENABLE+offline instances are included for the same reason. SWAP_IN and
-   * UNKNOWN are excluded because they do not represent assignable cluster capacity.
+   * <p>Membership rules live in
+   * {@link InstanceUtil#getInstancesUnableToAcceptOnlineReplicas(Map, java.util.Collection, long)}
+   * so that the controller and the read-only helix-rest endpoint that reports this number to
+   * clients share one implementation.
    *
    * <p>Used by both BestPossibleStateCalcStage (MM entry) and MaintenanceRecoveryStage
    * (MM exit) so that the same population is measured against each threshold.
@@ -777,18 +772,8 @@ public Set<String> getEnabledInstances() {
    * @return a fresh modifiable set of instance names.
    */
   public Set<String> getInstancesUnableToAcceptOnlineReplicas(long nowMs) {
-    Map<String, InstanceConfig> instanceConfigMap = getInstanceConfigMap();
-    Set<String> result = instanceConfigMap.entrySet().stream()
-        .filter(e -> !InstanceConstants.UNROUTABLE_INSTANCE_OPERATIONS.contains(
-            e.getValue().getInstanceOperation().getOperation()))
-        .map(Map.Entry::getKey)
-        .collect(Collectors.toCollection(HashSet::new));
-    result.removeAll(getEnabledLiveInstances());
-    result.removeIf(name -> {
-      InstanceConfig cfg = instanceConfigMap.get(name);
-      return cfg != null && cfg.isUnderInstanceOperationMaintenance(nowMs);
-    });
-    return result;
+    return InstanceUtil.getInstancesUnableToAcceptOnlineReplicas(getInstanceConfigMap(),
+        getLiveInstances().keySet(), nowMs);
   }
 
   /**

helix-core/src/main/java/org/apache/helix/util/InstanceUtil.java

@@ -19,8 +19,12 @@
  * under the License.
  */
 
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import javax.annotation.Nullable;
@@ -314,6 +318,79 @@ public static void setInstanceOperation(ConfigAccessor configAccessor,
     }
   }
 
+  /**
+   * Returns the live instances that are also marked with
+   * {@link InstanceConstants.InstanceOperation#ENABLE}. These are the instances that can
+   * currently accept ONLINE replicas.
+   *
+   * @param instanceConfigMap all instance configs in the cluster, keyed by instance name.
+   * @param liveInstanceNames names of the currently live instances.
+   * @return a fresh modifiable set of instance names.
+   */
+  private static Set<String> getEnabledLiveInstances(
+      Map<String, InstanceConfig> instanceConfigMap, Collection<String> liveInstanceNames) {
+    if (instanceConfigMap == null || liveInstanceNames == null) {
+      return new HashSet<>();
+    }
+    Set<String> enabledLiveInstances = new HashSet<>();
+    for (String instanceName : liveInstanceNames) {
+      InstanceConfig config = instanceConfigMap.get(instanceName);
+      if (config != null && config.getInstanceOperation().getOperation()
+          == InstanceConstants.InstanceOperation.ENABLE) {
+        enabledLiveInstances.add(instanceName);
+      }
+    }
+    return enabledLiveInstances;
+  }
+
+  /**
+   * Returns the set of instances that count toward the cluster-wide offline budget driving
+   * auto Maintenance Mode ({@code MAX_OFFLINE_INSTANCES_ALLOWED} at entry,
+   * {@code NUM_OFFLINE_INSTANCES_FOR_AUTO_EXIT} at exit).
+   *
+   * <p>An instance counts when all three conditions hold:
+   * <ul>
+   *   <li>Its InstanceOperation is routable, i.e. not in
+   *       {@link InstanceConstants#UNROUTABLE_INSTANCE_OPERATIONS}.</li>
+   *   <li>It is not currently enabled-and-live.</li>
+   *   <li>It does not carry a valid (unexpired) instance-operation maintenance marker.</li>
+   * </ul>
+   * EVACUATE and DISABLE instances are included because they cannot accept new ONLINE
+   * replicas; ENABLE+offline instances are included for the same reason. SWAP_IN and
+   * UNKNOWN are excluded because they do not represent assignable cluster capacity.
+   *
+   * <p>This is the single definition of the offline-budget population. The controller
+   * (MM entry in {@code BestPossibleStateCalcStage}, MM exit in
+   * {@code MaintenanceRecoveryStage}) reaches it through
+   * {@code BaseControllerDataProvider#getInstancesUnableToAcceptOnlineReplicas}, and
+   * helix-rest exposes the same computation read-only so clients never have to reimplement
+   * (and drift from) these rules.
+   *
+   * @param instanceConfigMap all instance configs in the cluster, keyed by instance name.
+   * @param liveInstanceNames names of the currently live instances.
+   * @param nowMs current wall-clock millis used for marker-expiry comparison.
+   * @return a fresh modifiable set of instance names.
+   */
+  public static Set<String> getInstancesUnableToAcceptOnlineReplicas(
+      Map<String, InstanceConfig> instanceConfigMap, Collection<String> liveInstanceNames,
+      long nowMs) {
+    if (instanceConfigMap == null || instanceConfigMap.isEmpty()) {
+      return new HashSet<>();
+    }
+    Set<String> result = instanceConfigMap.entrySet().stream()
+        .filter(e -> e.getValue() != null)
+        .filter(e -> !InstanceConstants.UNROUTABLE_INSTANCE_OPERATIONS.contains(
+            e.getValue().getInstanceOperation().getOperation()))
+        .map(Map.Entry::getKey)
+        .collect(Collectors.toCollection(HashSet::new));
+    result.removeAll(getEnabledLiveInstances(instanceConfigMap, liveInstanceNames));
+    result.removeIf(name -> {
+      InstanceConfig config = instanceConfigMap.get(name);
+      return config != null && config.isUnderInstanceOperationMaintenance(nowMs);
+    });
+    return result;
+  }
+
   private static String formatMatchingInstances(List<InstanceConfig> matchingInstances) {
     return matchingInstances.stream()
         .map(ic -> ic.getInstanceName() + " (operation="

helix-core/src/test/java/org/apache/helix/controller/dataproviders/TestInstancesUnableToAcceptOnlineReplicas.java

@@ -27,6 +27,7 @@
 
 import org.apache.helix.constants.InstanceConstants;
 import org.apache.helix.model.InstanceConfig;
+import org.apache.helix.model.LiveInstance;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -52,7 +53,7 @@ public class TestInstancesUnableToAcceptOnlineReplicas {
   public void testEnableLiveExcluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.ENABLE)),
-        enabledLive("h1"));
+        liveInstances("h1"));
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty(),
         "Healthy ENABLE+live instance must not count toward the offline budget");
   }
@@ -61,7 +62,7 @@ public void testEnableLiveExcluded() {
   public void testEnableOfflineIncluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.ENABLE)),
-        enabledLive());
+        liveInstances());
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("h1"), "ENABLE+offline is the canonical 'real outage' case and must count");
   }
@@ -70,7 +71,7 @@ public void testEnableOfflineIncluded() {
   public void testDisableLiveIncluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.DISABLE)),
-        enabledLive());
+        liveInstances("h1"));
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("h1"),
         "DISABLE instances cannot accept ONLINE replicas and must count regardless of liveness");
@@ -80,7 +81,7 @@ public void testDisableLiveIncluded() {
   public void testDisableOfflineIncluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.DISABLE)),
-        enabledLive());
+        liveInstances());
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("h1"));
   }
@@ -89,7 +90,7 @@ public void testDisableOfflineIncluded() {
   public void testEvacuateLiveIncluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.EVACUATE)),
-        enabledLive());
+        liveInstances("h1"));
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("h1"),
         "EVACUATE is the asymmetric op pre-fix; both entry and exit must now count it");
@@ -99,7 +100,7 @@ public void testEvacuateLiveIncluded() {
   public void testEvacuateOfflineIncluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.EVACUATE)),
-        enabledLive());
+        liveInstances());
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("h1"));
   }
@@ -108,7 +109,7 @@ public void testEvacuateOfflineIncluded() {
   public void testSwapInExcluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.SWAP_IN)),
-        enabledLive());
+        liveInstances());
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty(),
         "SWAP_IN is in UNROUTABLE_INSTANCE_OPERATIONS and must never count");
   }
@@ -117,7 +118,7 @@ public void testSwapInExcluded() {
   public void testUnknownExcluded() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.UNKNOWN)),
-        enabledLive());
+        liveInstances());
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty(),
         "UNKNOWN is in UNROUTABLE_INSTANCE_OPERATIONS and must never count");
   }
@@ -128,7 +129,7 @@ public void testUnknownExcluded() {
   public void testValidMarkerExemptsEnableOffline() {
     BaseControllerDataProvider provider = providerWith(
         configs(configWithMarker("h1", InstanceConstants.InstanceOperation.ENABLE, FUTURE_MS)),
-        enabledLive());
+        liveInstances());
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty(),
         "Valid marker on ENABLE+offline must exempt the instance from the budget");
   }
@@ -137,7 +138,7 @@ public void testValidMarkerExemptsEnableOffline() {
   public void testValidMarkerExemptsEvacuate() {
     BaseControllerDataProvider provider = providerWith(
         configs(configWithMarker("h1", InstanceConstants.InstanceOperation.EVACUATE, FUTURE_MS)),
-        enabledLive());
+        liveInstances());
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty(),
         "Valid marker on EVACUATE must exempt — this is the orchestrator-driven decom case");
   }
@@ -146,15 +147,15 @@ public void testValidMarkerExemptsEvacuate() {
   public void testValidMarkerExemptsDisable() {
     BaseControllerDataProvider provider = providerWith(
         configs(configWithMarker("h1", InstanceConstants.InstanceOperation.DISABLE, FUTURE_MS)),
-        enabledLive());
+        liveInstances());
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty());
   }
 
   @Test
   public void testExpiredMarkerDoesNotExempt() {
     BaseControllerDataProvider provider = providerWith(
         configs(configWithMarker("h1", InstanceConstants.InstanceOperation.ENABLE, PAST_MS)),
-        enabledLive());
+        liveInstances());
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("h1"), "Expired marker must behave as if the marker were absent");
   }
@@ -165,7 +166,7 @@ public void testBoundaryNowEqualsUntil() {
     // already past the window.
     BaseControllerDataProvider provider = providerWith(
         configs(configWithMarker("h1", InstanceConstants.InstanceOperation.ENABLE, NOW_MS)),
-        enabledLive());
+        liveInstances());
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("h1"), "nowMs == untilMs is no longer under maintenance; instance must count");
   }
@@ -175,7 +176,7 @@ public void testMarkerOnSwapInIsIrrelevant() {
     // SWAP_IN is filtered out before the marker check; a marker on it must not change that.
     BaseControllerDataProvider provider = providerWith(
         configs(configWithMarker("h1", InstanceConstants.InstanceOperation.SWAP_IN, FUTURE_MS)),
-        enabledLive());
+        liveInstances());
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty(),
         "Marker on a SWAP_IN instance must not change the (already excluded) outcome");
   }
@@ -184,7 +185,7 @@ public void testMarkerOnSwapInIsIrrelevant() {
 
   @Test
   public void testEmptyCluster() {
-    BaseControllerDataProvider provider = providerWith(configs(), enabledLive());
+    BaseControllerDataProvider provider = providerWith(configs(), liveInstances());
     Assert.assertTrue(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS).isEmpty());
   }
 
@@ -194,11 +195,11 @@ public void testMixedCluster() {
     // 1 DISABLE (counts), 1 EVACUATE (counts), 1 EVACUATE w/marker (exempt),
     // 1 SWAP_IN (excluded), 1 UNKNOWN (excluded).
     Map<String, InstanceConfig> instanceConfigMap = new HashMap<>();
-    Set<String> enabledLive = new HashSet<>();
+    Set<String> liveInstanceNames = new HashSet<>();
     for (int i = 0; i < 8; i++) {
       String name = "enable-live-" + i;
       instanceConfigMap.put(name, config(name, InstanceConstants.InstanceOperation.ENABLE));
-      enabledLive.add(name);
+      liveInstanceNames.add(name);
     }
     instanceConfigMap.put("enable-offline",
         config("enable-offline", InstanceConstants.InstanceOperation.ENABLE));
@@ -217,7 +218,7 @@ public void testMixedCluster() {
     instanceConfigMap.put("unknown",
         config("unknown", InstanceConstants.InstanceOperation.UNKNOWN));
 
-    BaseControllerDataProvider provider = providerWith(instanceConfigMap, enabledLive);
+    BaseControllerDataProvider provider = providerWith(instanceConfigMap, liveInstanceNames);
 
     Assert.assertEquals(provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS),
         setOf("enable-offline", "disable", "evacuate"),
@@ -232,7 +233,7 @@ public void testReturnedSetIsModifiable() {
     // particular) rely on this for downstream filtering. Verify both properties hold.
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.ENABLE)),
-        enabledLive());
+        liveInstances());
     Set<String> result = provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS);
     result.add("h99");
     Assert.assertTrue(result.contains("h99"), "Returned set must be modifiable");
@@ -242,7 +243,7 @@ public void testReturnedSetIsModifiable() {
   public void testReturnedSetIsIndependentOfFutureCalls() {
     BaseControllerDataProvider provider = providerWith(
         configs(config("h1", InstanceConstants.InstanceOperation.ENABLE)),
-        enabledLive());
+        liveInstances());
     Set<String> first = provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS);
     first.clear();
     Set<String> second = provider.getInstancesUnableToAcceptOnlineReplicas(NOW_MS);
@@ -254,22 +255,27 @@ public void testReturnedSetIsIndependentOfFutureCalls() {
 
   /**
    * Builds a {@link BaseControllerDataProvider} whose only contract is to return the
-   * supplied instance-config map and enabled-live set. The accessor under test depends on
+   * supplied instance-config map and live-instance set. The accessor under test depends on
    * exactly those two hooks plus {@link InstanceConfig#isUnderInstanceOperationMaintenance},
    * so overriding the two getters is sufficient and avoids initializing the full cluster
-   * cache machinery.
+   * cache machinery. The enabled-live subset is derived from the configs the same way the
+   * controller derives it, so the ENABLE filter is exercised here rather than stubbed.
    */
   private static BaseControllerDataProvider providerWith(
-      Map<String, InstanceConfig> instanceConfigMap, Set<String> enabledLive) {
+      Map<String, InstanceConfig> instanceConfigMap, Set<String> liveInstanceNames) {
+    Map<String, LiveInstance> liveInstanceMap = new HashMap<>();
+    for (String name : liveInstanceNames) {
+      liveInstanceMap.put(name, new LiveInstance(name));
+    }
     return new BaseControllerDataProvider() {
       @Override
       public Map<String, InstanceConfig> getInstanceConfigMap() {
         return instanceConfigMap;
       }
 
       @Override
-      public Set<String> getEnabledLiveInstances() {
-        return Collections.unmodifiableSet(enabledLive);
+      public Map<String, LiveInstance> getLiveInstances() {
+        return Collections.unmodifiableMap(liveInstanceMap);
       }
     };
   }
@@ -297,7 +303,7 @@ private static Map<String, InstanceConfig> configs(InstanceConfig... cfgs) {
   }
 
   /**
-   * Test-side set builder. Used for both the enabled-live input set and the expected-result
+   * Test-side set builder. Used for both the live-instance input set and the expected-result
    * set so that assertions and setup share the same shape; the two readings sit at the call
    * site through the method name on the input side and explicit assertEquals on the result.
    */
@@ -307,7 +313,7 @@ private static Set<String> setOf(String... names) {
     return set;
   }
 
-  private static Set<String> enabledLive(String... names) {
+  private static Set<String> liveInstances(String... names) {
     return setOf(names);
   }
 }

helix-rest/src/main/java/org/apache/helix/rest/server/resources/AbstractResource.java

@@ -73,6 +73,7 @@ public enum Command {
     delete,
     stoppable,
     instanceOperationMaintenance,
+    getInstancesUnableToAcceptOnlineReplicas,
     rebalance,
     reset,
     resetPartitions,