Add _deserializedPartitionCapacityMap to ResourceConfig model to ensu…

…re that the partition capacity map is only deserialized one time (#2657)

Add _deserializedPartitionCapacityMap to ResourceConfig model to ensure that the partition capacity map is only deserialized one time. This will lead to major performance gains as this was previously called in several parts of BestPossibleStateCalculation and CurrentStateComputation.

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

@@ -69,6 +69,7 @@ public enum ResourceConfigConstants {
   private static final ObjectMapper _objectMapper = new ObjectMapper();
 
   public static final String DEFAULT_PARTITION_KEY = "DEFAULT";
+  private Map<String, Map<String, Integer>> _deserializedPartitionCapacityMap;
 
   /**
    * Instantiate for a specific instance
@@ -389,6 +390,12 @@ public void setPreferenceLists(Map<String, List<String>> instanceLists) {
    * @throws IOException - when JSON conversion fails
    */
   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 && !_deserializedPartitionCapacityMap.isEmpty()) {
+      return _deserializedPartitionCapacityMap;
+    }
+
     Map<String, String> partitionCapacityData =
         _record.getMapField(ResourceConfigProperty.PARTITION_CAPACITY_MAP.name());
     Map<String, Map<String, Integer>> partitionCapacityMap = new HashMap<>();
@@ -401,7 +408,11 @@ public Map<String, Map<String, Integer>> getPartitionCapacityMap() throws IOExce
         partitionCapacityMap.put(partition, capacities);
       }
     }
-    return partitionCapacityMap;
+
+    // Only set the deserialized map when the deserialization succeeds, so we don't have the potential
+    // of having a partially populated map.
+    _deserializedPartitionCapacityMap = partitionCapacityMap;
+    return _deserializedPartitionCapacityMap;
   }
 
   /**
@@ -423,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
@@ -434,9 +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 = 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);