From 5b1a036bc11a893980ece40872feead3472f98d9 Mon Sep 17 00:00:00 2001
From: Xiaxuan Gao <xiaxgao@linkedin.com>
Date: Wed, 9 Oct 2024 10:11:17 -0700
Subject: [PATCH] Fix BestPossibleExternalViewVerifier for WAGED resource
 (#2939)

Fix the bug where BestPossibleExternalViewVerifier fails to accurately verify WAGED resources. If the user has requested the WAGED resources, the verifier should run against the resources of whole cluster.
---
 .../BestPossibleExternalViewVerifier.java     | 34 ++++++++++++++-----
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/tools/ClusterVerifiers/BestPossibleExternalViewVerifier.java b/helix-core/src/main/java/org/apache/helix/tools/ClusterVerifiers/BestPossibleExternalViewVerifier.java
index 0b0926dda3..cae5cb010d 100644
--- a/helix-core/src/main/java/org/apache/helix/tools/ClusterVerifiers/BestPossibleExternalViewVerifier.java
+++ b/helix-core/src/main/java/org/apache/helix/tools/ClusterVerifiers/BestPossibleExternalViewVerifier.java
@@ -27,12 +27,15 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import org.apache.helix.HelixDefinedState;
 import org.apache.helix.PropertyKey;
 import org.apache.helix.controller.common.PartitionStateMap;
 import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
+import org.apache.helix.controller.rebalancer.util.WagedValidationUtil;
 import org.apache.helix.controller.rebalancer.waged.ReadOnlyWagedRebalancer;
 import org.apache.helix.controller.stages.AttributeName;
 import org.apache.helix.controller.stages.BestPossibleStateCalcStage;
@@ -274,8 +277,21 @@ protected synchronized boolean verifyState() {
 
       // Filter resources if requested
       if (_resources != null && !_resources.isEmpty()) {
-        idealStates.keySet().retainAll(_resources);
-        extViews.keySet().retainAll(_resources);
+        // Find if there are waged-enabled resources among the requested resources
+        boolean hasRequestedWagedResources = _resources.stream().anyMatch(
+            resourceEntry -> WagedValidationUtil.isWagedEnabled(idealStates.get(resourceEntry)));
+        Set<String> resourcesToRetain = new HashSet<>(_resources);
+
+        if (hasRequestedWagedResources) {
+          // If waged-enabled resources are found, retain all the waged-enabled resources and the
+          // user requested resources.
+          resourcesToRetain.addAll(idealStates.keySet().stream().filter(
+                  resourceEntry -> WagedValidationUtil.isWagedEnabled(idealStates.get(resourceEntry)))
+              .collect(Collectors.toSet()));
+        }
+
+        idealStates.keySet().retainAll(resourcesToRetain);
+        extViews.keySet().retainAll(resourcesToRetain);
       }
 
       // if externalView is not empty and idealState doesn't exist
@@ -290,7 +306,7 @@ protected synchronized boolean verifyState() {
       }
 
       // calculate best possible state
-      BestPossibleStateOutput bestPossOutput = calcBestPossState(_dataProvider, _resources);
+      BestPossibleStateOutput bestPossOutput = calcBestPossState(_dataProvider, idealStates);
       Map<String, Map<Partition, Map<String, String>>> bestPossStateMap =
           bestPossOutput.getStateMap();
 
@@ -408,26 +424,26 @@ private Map<String, Map<String, String>> convertBestPossibleState(
    * kick off the BestPossibleStateCalcStage we are providing an empty current state map
    *
    * @param cache
-   * @param resources
+   * @param resourceToIdealStateMap
    * @return
    * @throws Exception
    */
-  private BestPossibleStateOutput calcBestPossState(ResourceControllerDataProvider cache, Set<String> resources)
-      throws Exception {
+  private BestPossibleStateOutput calcBestPossState(ResourceControllerDataProvider cache,
+      Map<String, IdealState> resourceToIdealStateMap) throws Exception {
     ClusterEvent event = new ClusterEvent(_clusterName, ClusterEventType.StateVerifier);
     event.addAttribute(AttributeName.ControllerDataProvider.name(), cache);
 
     RebalanceUtil.runStage(event, new ResourceComputationStage());
 
-    if (resources != null && !resources.isEmpty()) {
+    if (resourceToIdealStateMap != null && !resourceToIdealStateMap.isEmpty()) {
       // Filtering out all non-required resources
       final Map<String, Resource> resourceMap = event.getAttribute(AttributeName.RESOURCES.name());
-      resourceMap.keySet().retainAll(resources);
+      resourceMap.keySet().retainAll(resourceToIdealStateMap.keySet());
       event.addAttribute(AttributeName.RESOURCES.name(), resourceMap);
 
       final Map<String, Resource> resourceMapToRebalance =
           event.getAttribute(AttributeName.RESOURCES_TO_REBALANCE.name());
-      resourceMapToRebalance.keySet().retainAll(resources);
+      resourceMapToRebalance.keySet().retainAll(resourceToIdealStateMap.keySet());
       event.addAttribute(AttributeName.RESOURCES_TO_REBALANCE.name(), resourceMapToRebalance);
     }