From 96711482867cf045726631efd14b249faa5e088e Mon Sep 17 00:00:00 2001
From: ngangwar <ngangwar@linkedin.com>
Date: Fri, 28 Nov 2025 12:14:07 +0530
Subject: [PATCH] Fix NPE due to resource deletion

---
 .../rebalancer/waged/WagedRebalancer.java     | 21 ++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java
index 39a197bff5..5c2b3b1475 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/WagedRebalancer.java
@@ -276,9 +276,17 @@ public Map<String, IdealState> computeNewIdealStates(ResourceControllerDataProvi
     // the target assignment.
     newIdealStates.values().parallelStream().forEach(idealState -> {
       String resourceName = idealState.getResourceName();
+      Resource resource = resourceMap.get(resourceName);
+
+      // This can happen when a resource is deleted but still exists in cached assignment
+      if (resource == null) {
+        LOG.warn("Resource {} not found in current resourceMap. Skipping ideal state calculation.", resourceName);
+        return;
+      }
+
       // Adjust the states according to the current state.
       ResourceAssignment finalAssignment = _mappingCalculator
-          .computeBestPossiblePartitionState(clusterData, idealState, resourceMap.get(resourceName),
+          .computeBestPossiblePartitionState(clusterData, idealState, resource,
               currentStateOutput);
 
       // Clean up the state mapping fields. Use the final assignment that is calculated by the
@@ -291,6 +299,17 @@ public Map<String, IdealState> computeNewIdealStates(ResourceControllerDataProvi
             newStateMap == null ? Collections.emptyMap() : newStateMap);
       }
     });
+
+    // Remove any resources from newIdealStates that don't exist in the current resourceMap
+    // This ensures we don't return stale IdealStates for deleted resources
+    newIdealStates.keySet().removeIf(resourceName -> {
+      if (!resourceMap.containsKey(resourceName)) {
+        LOG.warn("Removing stale resource {} from ideal states as it no longer exists in cluster.", resourceName);
+        return true;
+      }
+      return false;
+    });
+
     LOG.info("Finish computing new ideal states for resources: {}",
         resourceMap.keySet().toString());
     return newIdealStates;