From 93270f1ffeaecd700aae9be24406879847edddc1 Mon Sep 17 00:00:00 2001
From: Peter Alfonsi <petealft@amazon.com>
Date: Wed, 10 Apr 2024 15:40:41 -0700
Subject: [PATCH] Addressed Ankit's comments

Signed-off-by: Peter Alfonsi <petealft@amazon.com>
---
 .../cache/stats/MultiDimensionCacheStats.java |  4 +--
 .../common/cache/stats/StatsHolder.java       | 27 +++++++++++--------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/server/src/main/java/org/opensearch/common/cache/stats/MultiDimensionCacheStats.java b/server/src/main/java/org/opensearch/common/cache/stats/MultiDimensionCacheStats.java
index 10b3ca8570d19..f60029381e0e9 100644
--- a/server/src/main/java/org/opensearch/common/cache/stats/MultiDimensionCacheStats.java
+++ b/server/src/main/java/org/opensearch/common/cache/stats/MultiDimensionCacheStats.java
@@ -40,7 +40,7 @@ public MultiDimensionCacheStats(StreamInput in) throws IOException {
         // of the last node we read. This allows us to avoid ambiguity if nodes have the same dimension value, without
         // having to serialize the whole path to each node.
         this.dimensionNames = List.of(in.readStringArray());
-        this.statsRoot = new MDCSDimensionNode(null, true);
+        this.statsRoot = new MDCSDimensionNode("", true);
         List<MDCSDimensionNode> ancestorsOfLastRead = List.of(statsRoot);
         while (ancestorsOfLastRead != null) {
             ancestorsOfLastRead = readAndAttachDimensionNode(in, ancestorsOfLastRead);
@@ -142,7 +142,7 @@ public long getTotalEntries() {
      */
     MDCSDimensionNode aggregateByLevels(List<String> levels) {
         List<String> filteredLevels = filterLevels(levels);
-        MDCSDimensionNode newRoot = new MDCSDimensionNode(null, true, statsRoot.getStats());
+        MDCSDimensionNode newRoot = new MDCSDimensionNode("", true, statsRoot.getStats());
         for (MDCSDimensionNode child : statsRoot.children.values()) {
             aggregateByLevelsHelper(newRoot, child, filteredLevels, 0);
         }
diff --git a/server/src/main/java/org/opensearch/common/cache/stats/StatsHolder.java b/server/src/main/java/org/opensearch/common/cache/stats/StatsHolder.java
index 271264ce7b71c..09174055770da 100644
--- a/server/src/main/java/org/opensearch/common/cache/stats/StatsHolder.java
+++ b/server/src/main/java/org/opensearch/common/cache/stats/StatsHolder.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.common.cache.stats;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
@@ -40,8 +41,8 @@ public class StatsHolder {
     private final Lock lock = new ReentrantLock();
 
     public StatsHolder(List<String> dimensionNames) {
-        this.dimensionNames = dimensionNames;
-        this.statsRoot = new DimensionNode(null, true); // The root node has no dimension value associated with it, only children
+        this.dimensionNames = Collections.unmodifiableList(dimensionNames);
+        this.statsRoot = new DimensionNode("", true); // The root node has the empty string as its dimension value
     }
 
     public List<String> getDimensionNames() {
@@ -102,7 +103,17 @@ public long count() {
 
     private void internalIncrement(List<String> dimensionValues, Consumer<DimensionNode> adder, boolean createNodesIfAbsent) {
         assert dimensionValues.size() == dimensionNames.size();
-        internalIncrementHelper(dimensionValues, statsRoot, 0, adder, createNodesIfAbsent);
+        // First try to increment without creating nodes
+        boolean didIncrement = internalIncrementHelper(dimensionValues, statsRoot, 0, adder, false);
+        // If we failed to increment, because nodes had to be created, obtain the lock and run again while creating nodes if needed
+        if (!didIncrement) {
+            try {
+                lock.lock();
+                internalIncrementHelper(dimensionValues, statsRoot, 0, adder, createNodesIfAbsent);
+            } finally {
+                lock.unlock();
+            }
+        }
     }
 
     /**
@@ -126,14 +137,8 @@ private boolean internalIncrementHelper(
         DimensionNode child = node.getChild(dimensionValues.get(depth));
         if (child == null) {
             if (createNodesIfAbsent) {
-                // If we have to create a new node, obtain the lock first
                 boolean createMapInChild = depth < dimensionValues.size() - 1;
-                lock.lock();
-                try {
-                    child = node.createChild(dimensionValues.get(depth), createMapInChild);
-                } finally {
-                    lock.unlock();
-                }
+                child = node.createChild(dimensionValues.get(depth), createMapInChild);
             } else {
                 return false;
             }
@@ -150,7 +155,7 @@ private boolean internalIncrementHelper(
      * Produce an immutable CacheStats representation of these stats.
      */
     public CacheStats getCacheStats() {
-        MDCSDimensionNode snapshot = new MDCSDimensionNode(null, true, statsRoot.getStatsSnapshot());
+        MDCSDimensionNode snapshot = new MDCSDimensionNode("", true, statsRoot.getStatsSnapshot());
         // Traverse the tree and build a corresponding tree of MDCSDimensionNode, to pass to MultiDimensionCacheStats.
         if (statsRoot.getChildren() != null) {
             for (DimensionNode child : statsRoot.getChildren().values()) {