Merge branch 'tiramisu-stats-tiers' into tiramisu-stats-api

peteralfonsi · Apr 10, 2024 · 97fe6f3 · 97fe6f3
2 parents 92dbb73 + 93270f1
commit 97fe6f3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 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
@@ -52,7 +52,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);
@@ -167,7 +167,7 @@ public void writeToWithClassName(StreamOutput out) throws IOException {
      */
     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
@@ -10,6 +10,7 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
@@ -46,9 +47,9 @@ public class StatsHolder {
     private final String storeName;
 
     public StatsHolder(List<String> dimensionNames, String storeName) {
-        this.dimensionNames = dimensionNames;
+        this.dimensionNames = Collections.unmodifiableList(dimensionNames);
         this.storeName = storeName;
-        this.statsRoot = new DimensionNode(null, true); // The root node has no dimension value associated with it, only children
+        this.statsRoot = new DimensionNode("", true); // The root node has the empty string as its dimension value
     }
 
     public List<String> getDimensionNames() {
@@ -110,7 +111,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();
+            }
+        }
     }
 
     /**
@@ -134,14 +145,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;
             }
@@ -158,7 +163,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()) {