From 5691d3a5ec7133b47f058c22fe5250562835cd0e Mon Sep 17 00:00:00 2001
From: Huizhi L <ihuizhi.lu@gmail.com>
Date: Tue, 5 Nov 2019 17:18:25 -0800
Subject: [PATCH] Filter resource map with ideal states for instance capacity
 metrics. (#574)

ResourceToReblance map also has resources from current states. And this causes null pointer exceptions at parsing all replicas stage when the resource is not in ideal states. This diff fixes the issue by only using the resources in ideal states to parse all replicas.
---
 .../stages/CurrentStateComputationStage.java  | 20 ++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/controller/stages/CurrentStateComputationStage.java b/helix-core/src/main/java/org/apache/helix/controller/stages/CurrentStateComputationStage.java
index 5c70e133bd..890fa570ae 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/stages/CurrentStateComputationStage.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/stages/CurrentStateComputationStage.java
@@ -24,6 +24,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import org.apache.helix.controller.LogUtil;
 import org.apache.helix.controller.dataproviders.BaseControllerDataProvider;
@@ -33,8 +34,14 @@
 import org.apache.helix.controller.rebalancer.waged.model.AssignableNode;
 import org.apache.helix.controller.rebalancer.waged.model.ClusterModel;
 import org.apache.helix.controller.rebalancer.waged.model.ClusterModelProvider;
-import org.apache.helix.model.*;
+import org.apache.helix.model.CurrentState;
+import org.apache.helix.model.IdealState;
+import org.apache.helix.model.LiveInstance;
+import org.apache.helix.model.Message;
 import org.apache.helix.model.Message.MessageType;
+import org.apache.helix.model.Partition;
+import org.apache.helix.model.Resource;
+import org.apache.helix.model.ResourceAssignment;
 import org.apache.helix.monitoring.mbeans.ClusterStatusMonitor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -238,10 +245,17 @@ private void reportInstanceCapacityMetrics(ClusterStatusMonitor clusterStatusMon
       CurrentStateOutput currentStateOutput) {
     asyncExecute(dataProvider.getAsyncTasksThreadPool(), () -> {
       try {
+        // ResourceToRebalance map also has resources from current states.
+        // Only use the resources in ideal states to parse all replicas.
+        Map<String, IdealState> idealStateMap = dataProvider.getIdealStates();
+        Map<String, Resource> resourceToMonitorMap = resourceMap.entrySet().stream()
+            .filter(resourceName -> idealStateMap.containsKey(resourceName))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
         Map<String, ResourceAssignment> currentStateAssignment =
-            currentStateOutput.getAssignment(resourceMap.keySet());
+            currentStateOutput.getAssignment(resourceToMonitorMap.keySet());
         ClusterModel clusterModel = ClusterModelProvider.generateClusterModelFromCurrentState(
-            dataProvider, resourceMap, currentStateAssignment);
+            dataProvider, resourceToMonitorMap, currentStateAssignment);
 
         Map<String, Double> maxUsageMap = new HashMap<>();
         for (AssignableNode node : clusterModel.getAssignableNodes().values()) {