From d3887bd42f215402cdf1a41a6d55420c072ded50 Mon Sep 17 00:00:00 2001
From: Zachary Pinto <zapinto@linkedin.com>
Date: Thu, 26 Oct 2023 23:53:30 -0400
Subject: [PATCH] Add handling for clusterConfig == null in
 updateSwappingInstances and fix AssignableNode to check for
 clusterTopologyConfig when attempting to get logicalId.

---
 .../dataproviders/BaseControllerDataProvider.java     | 11 +++++++++--
 .../rebalancer/waged/model/AssignableNode.java        |  2 +-
 .../rebalancer/waged/TestWagedRebalancer.java         |  2 +-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/controller/dataproviders/BaseControllerDataProvider.java b/helix-core/src/main/java/org/apache/helix/controller/dataproviders/BaseControllerDataProvider.java
index 0a3e163cce..9dd5173841 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/dataproviders/BaseControllerDataProvider.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/dataproviders/BaseControllerDataProvider.java
@@ -888,10 +888,17 @@ private void updateDisabledInstances(Collection<InstanceConfig> instanceConfigs,
 
   private void updateSwappingInstances(Collection<InstanceConfig> instanceConfigs,
       Set<String> liveEnabledInstances, ClusterConfig clusterConfig) {
-    ClusterTopologyConfig clusterTopologyConfig =
-        ClusterTopologyConfig.createFromClusterConfig(clusterConfig);
     _swapOutInstanceNameToSwapInInstanceName.clear();
     _enabledLiveSwapInInstanceNames.clear();
+
+    if (clusterConfig == null) {
+      logger.warn("Skip refreshing swapping instances because clusterConfig is null.");
+      return;
+    }
+
+    ClusterTopologyConfig clusterTopologyConfig =
+        ClusterTopologyConfig.createFromClusterConfig(clusterConfig);
+
     Map<String, String> swapOutLogicalIdsByInstanceName = new HashMap<>();
     Map<String, String> swapInInstancesByLogicalId = new HashMap<>();
     instanceConfigs.forEach(instanceConfig -> {
diff --git a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/AssignableNode.java b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/AssignableNode.java
index 7cea4b45eb..e4679bc12e 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/AssignableNode.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/rebalancer/waged/model/AssignableNode.java
@@ -78,7 +78,7 @@ public class AssignableNode implements Comparable<AssignableNode> {
       InstanceConfig instanceConfig, String instanceName) {
     _instanceName = instanceName;
     _logicaId =
-        clusterConfig != null ? instanceConfig.getLogicalId(clusterTopologyConfig.getEndNodeType())
+        clusterTopologyConfig != null ? instanceConfig.getLogicalId(clusterTopologyConfig.getEndNodeType())
             : instanceName;
     Map<String, Integer> instanceCapacity = fetchInstanceCapacity(clusterConfig, instanceConfig);
     _faultZone = computeFaultZone(clusterConfig, instanceConfig);
diff --git a/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/TestWagedRebalancer.java b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/TestWagedRebalancer.java
index 608a4d3afe..881b0ca0e2 100644
--- a/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/TestWagedRebalancer.java
+++ b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/TestWagedRebalancer.java
@@ -738,7 +738,7 @@ public void testRebalanceOverwrite() throws HelixRebalanceException, IOException
     instances.add(offlineInstance);
     when(clusterData.getAllInstances()).thenReturn(instances);
     when(clusterData.getEnabledInstances()).thenReturn(instances);
-    when(clusterData.getEnabledLiveInstances()).thenReturn(ImmutableSet.of(instance0, instance1, instance2));
+    when(clusterData.getEnabledLiveInstances()).thenReturn(new HashSet<>(Set.of(instance0, instance1, instance2)));
     Map<String, Long> instanceOfflineTimeMap = new HashMap<>();
     instanceOfflineTimeMap.put(offlineInstance, System.currentTimeMillis() + Integer.MAX_VALUE);
     when(clusterData.getInstanceOfflineTimeMap()).thenReturn(instanceOfflineTimeMap);