diff --git a/helix-core/src/main/java/org/apache/helix/model/ResourceConfig.java b/helix-core/src/main/java/org/apache/helix/model/ResourceConfig.java index e4dd0ab8a2..69aa533fdc 100644 --- a/helix-core/src/main/java/org/apache/helix/model/ResourceConfig.java +++ b/helix-core/src/main/java/org/apache/helix/model/ResourceConfig.java @@ -392,7 +392,7 @@ public void setPreferenceLists(Map> instanceLists) { public Map> 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> 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> partitionC } Map 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> newDeserializedPartitionCapacityMap = new HashMap<>(); for (String partition : partitionCapacityMap.keySet()) { Map capacities = partitionCapacityMap.get(partition); // Verify the input is valid @@ -445,11 +448,12 @@ public void setPartitionCapacityMap(Map> 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; } /** diff --git a/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestAssignableReplica.java b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestAssignableReplica.java index de2ce7e560..55abfdc490 100644 --- a/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestAssignableReplica.java +++ b/helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestAssignableReplica.java @@ -71,8 +71,14 @@ public void testConstructReplicaWithResourceConfig() throws IOException { Map 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> capacityMap = - testResourceConfigResource.getPartitionCapacityMap(); + new HashMap<>(testResourceConfigResource.getPartitionCapacityMap()); + String partitionName2 = partitionNamePrefix + 2; capacityMap.put(partitionName2, capacityDataMapResource2); testResourceConfigResource.setPartitionCapacityMap(capacityMap);