Add hardening, so that the caller of setPartitionCapacityMap cannot m…

…odify the mapped passed as argument with the reference they still hold.

helix-core/src/main/java/org/apache/helix/model/ResourceConfig.java

@@ -392,7 +392,7 @@ public void setPreferenceLists(Map<String, List<String>> instanceLists) {
   public Map<String, Map<String, Integer>> getPartitionCapacityMap() throws IOException {
     // It is very expensive to deserialize the partition capacity map every time this is called.
     // Cache the deserialized map to avoid the overhead.
-    if (_deserializedPartitionCapacityMap != null) {
+    if (_deserializedPartitionCapacityMap != null && !_deserializedPartitionCapacityMap.isEmpty()) {
       return _deserializedPartitionCapacityMap;
     }
 
@@ -410,7 +410,7 @@ public Map<String, Map<String, Integer>> getPartitionCapacityMap() throws IOExce
     }
 
     // Only set the deserialized map when the deserialization succeeds, so we don't have the potential
-    // of having a partially contracted map.
+    // of having a partially populated map.
     _deserializedPartitionCapacityMap = partitionCapacityMap;
     return _deserializedPartitionCapacityMap;
   }
@@ -434,6 +434,9 @@ public void setPartitionCapacityMap(Map<String, Map<String, Integer>> partitionC
     }
 
     Map<String, String> newCapacityRecord = new HashMap<>();
+    // We want a copy of the partitionCapacityMap, so that the caller can no longer modify the
+    // _deserializedPartitionCapacityMap after this call through their reference to partitionCapacityMap.
+    Map<String, Map<String, Integer>> newDeserializedPartitionCapacityMap = new HashMap<>();
     for (String partition : partitionCapacityMap.keySet()) {
       Map<String, Integer> capacities = partitionCapacityMap.get(partition);
       // Verify the input is valid
@@ -445,11 +448,12 @@ public void setPartitionCapacityMap(Map<String, Map<String, Integer>> partitionC
             String.format("Capacity Data contains a negative value:%s", capacities.toString()));
       }
       newCapacityRecord.put(partition, _objectMapper.writeValueAsString(capacities));
+      newDeserializedPartitionCapacityMap.put(partition, Map.copyOf(capacities));
     }
 
     _record.setMapField(ResourceConfigProperty.PARTITION_CAPACITY_MAP.name(), newCapacityRecord);
     // Set deserialize map after we have successfully added it to the record.
-    _deserializedPartitionCapacityMap = partitionCapacityMap;
+    _deserializedPartitionCapacityMap = newDeserializedPartitionCapacityMap;
   }
 
   /**

helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestAssignableReplica.java

@@ -71,8 +71,14 @@ public void testConstructReplicaWithResourceConfig() throws IOException {
     Map<String, Integer> capacityDataMapResource2 = new HashMap<>();
     capacityDataMapResource2.put("item1", 5);
     capacityDataMapResource2.put("item2", 10);
+
+    // We should not directly modify the contents returned by getPartitionCapacityMap()
+    // This will not guard against the modification of the kv pairs in the inner maps as this
+    // is not creating a deepCopy but will ensure we don't override top level kv pairs in
+    // testResourceConfigResource.
     Map<String, Map<String, Integer>> capacityMap =
-        testResourceConfigResource.getPartitionCapacityMap();
+        new HashMap<>(testResourceConfigResource.getPartitionCapacityMap());
+
     String partitionName2 = partitionNamePrefix + 2;
     capacityMap.put(partitionName2, capacityDataMapResource2);
     testResourceConfigResource.setPartitionCapacityMap(capacityMap);