From 88ba8aeada650b4279eb12917226329095d9f865 Mon Sep 17 00:00:00 2001
From: Jiajun Wang <jjwang@linkedin.com>
Date: Thu, 28 Nov 2019 21:09:28 -0800
Subject: [PATCH] Fix a potential issue in the ResourceChangeSnapshot.

The trim logic in the ResourceChangeSnapshot for cleaning up the IdealState should not clear the whole map. This will cause the WAGED rebalancer ignores changes such as new partitions into the partition list.
Modify the test case accordingly.
---
 .../changedetector/ResourceChangeSnapshot.java     | 14 ++++++++++----
 .../changedetector/TestResourceChangeDetector.java | 11 +++++------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/helix-core/src/main/java/org/apache/helix/controller/changedetector/ResourceChangeSnapshot.java b/helix-core/src/main/java/org/apache/helix/controller/changedetector/ResourceChangeSnapshot.java
index c965124b01..fc8c5c483c 100644
--- a/helix-core/src/main/java/org/apache/helix/controller/changedetector/ResourceChangeSnapshot.java
+++ b/helix-core/src/main/java/org/apache/helix/controller/changedetector/ResourceChangeSnapshot.java
@@ -24,8 +24,10 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import org.apache.helix.HelixConstants;
+import org.apache.helix.ZNRecord;
 import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
 import org.apache.helix.model.ClusterConfig;
 import org.apache.helix.model.IdealState;
@@ -130,18 +132,22 @@ private IdealState trimIdealState(IdealState originalIdealState) {
     // Clone the IdealState to avoid modifying the objects in the Cluster Data Cache, which might
     // be used by the other stages in the pipeline.
     IdealState trimmedIdealState = new IdealState(originalIdealState.getRecord());
+    ZNRecord trimmedIdealStateRecord = trimmedIdealState.getRecord();
     switch (originalIdealState.getRebalanceMode()) {
+      // WARNING: the IdealState copy constructor is not really deep copy. So we should not modify
+      // the values directly or the cached values will be changed.
       case FULL_AUTO:
         // For FULL_AUTO resources, both map fields and list fields are not considered as data input
         // for the controller. The controller will write to these two types of fields for persisting
         // the assignment mapping.
-        trimmedIdealState.getRecord().setListFields(Collections.emptyMap());
-        trimmedIdealState.getRecord().setMapFields(Collections.emptyMap());
-        break;
+        trimmedIdealStateRecord.setListFields(trimmedIdealStateRecord.getListFields().keySet().stream().collect(
+            Collectors.toMap(partition -> partition, partition -> Collections.emptyList())));
+        // Continue to clean up map fields in the SEMI_AUTO case.
       case SEMI_AUTO:
         // For SEMI_AUTO resources, map fields are not considered as data input for the controller.
         // The controller will write to the map fields for persisting the assignment mapping.
-        trimmedIdealState.getRecord().setMapFields(Collections.emptyMap());
+        trimmedIdealStateRecord.setMapFields(trimmedIdealStateRecord.getMapFields().keySet().stream().collect(
+            Collectors.toMap(partition -> partition, partition -> Collections.emptyMap())));
         break;
       default:
         break;
diff --git a/helix-core/src/test/java/org/apache/helix/controller/changedetector/TestResourceChangeDetector.java b/helix-core/src/test/java/org/apache/helix/controller/changedetector/TestResourceChangeDetector.java
index 8567c4c8a8..275e32c8a3 100644
--- a/helix-core/src/test/java/org/apache/helix/controller/changedetector/TestResourceChangeDetector.java
+++ b/helix-core/src/test/java/org/apache/helix/controller/changedetector/TestResourceChangeDetector.java
@@ -359,22 +359,21 @@ public void testIgnoreControllerGeneratedFields() {
         .addResource(CLUSTER_NAME, resourceName, NUM_PARTITIONS, STATE_MODEL);
     IdealState idealState = _dataAccessor.getProperty(_keyBuilder.idealStates(resourceName));
     idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
+    idealState.getRecord().getMapFields().put("Partition1", new HashMap<>());
     _dataAccessor.updateProperty(_keyBuilder.idealStates(resourceName), idealState);
-
     _dataProvider.notifyDataChange(ChangeType.CLUSTER_CONFIG);
     _dataProvider.notifyDataChange(ChangeType.IDEAL_STATE);
     _dataProvider.refresh(_dataAccessor);
 
+    // Test with ignore option to be true
     ResourceChangeDetector changeDetector = new ResourceChangeDetector(true);
     changeDetector.updateSnapshots(_dataProvider);
-
-    // Now, modify the field that is modifying by Helix logic
-    idealState.getRecord().getMapFields().put("Extra_Change", new HashMap<>());
-    _dataAccessor.updateProperty(_keyBuilder.idealStates(NEW_RESOURCE_NAME), idealState);
+    // Now, modify the field
+    idealState.getRecord().getMapFields().put("Partition1", Collections.singletonMap("foo", "bar"));
+    _dataAccessor.updateProperty(_keyBuilder.idealStates(resourceName), idealState);
     _dataProvider.notifyDataChange(ChangeType.IDEAL_STATE);
     _dataProvider.refresh(_dataAccessor);
     changeDetector.updateSnapshots(_dataProvider);
-
     Assert.assertEquals(changeDetector.getChangeTypes(),
         Collections.singleton(ChangeType.IDEAL_STATE));
     Assert.assertEquals(