From 44d34619335e37b9a6bc8dbbbac097a290e2a6b6 Mon Sep 17 00:00:00 2001
From: Zachary Pinto <zapinto@linkedin.com>
Date: Wed, 18 Oct 2023 11:22:15 -0700
Subject: [PATCH] Add _deserializedPartitionCapacityMap to ResourceConfig model
 to ensure 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.
---
 .../apache/helix/model/ResourceConfig.java    | 19 ++++++++++++++++++-
 .../waged/model/TestAssignableReplica.java    |  8 +++++++-
 2 files changed, 25 insertions(+), 2 deletions(-)

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 13e96cf99e..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
@@ -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;
   }
 
   /**
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<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);