diff --git a/server/src/main/java/org/opensearch/common/cache/stats/CacheStatsHolder.java b/server/src/main/java/org/opensearch/common/cache/stats/CacheStatsHolder.java index 214b14e9fefdd..a8b7c27ef9e79 100644 --- a/server/src/main/java/org/opensearch/common/cache/stats/CacheStatsHolder.java +++ b/server/src/main/java/org/opensearch/common/cache/stats/CacheStatsHolder.java @@ -12,6 +12,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.TreeMap; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; @@ -156,28 +157,7 @@ private boolean internalIncrementHelper( * Produce an immutable version of these stats. */ public ImmutableCacheStatsHolder getImmutableCacheStatsHolder() { - ImmutableCacheStatsHolder.Node snapshot = new ImmutableCacheStatsHolder.Node("", true, statsRoot.getImmutableStats()); - // Traverse the tree and build a corresponding tree of MDCSDimensionNode, to pass to MultiDimensionCacheStats. - if (statsRoot.getChildren() != null) { - for (Node child : statsRoot.getChildren().values()) { - getImmutableCacheStatsHelper(child, snapshot); - } - } - return new ImmutableCacheStatsHolder(snapshot, dimensionNames); - } - - private void getImmutableCacheStatsHelper(Node currentNodeInOriginalTree, ImmutableCacheStatsHolder.Node parentInNewTree) { - ImmutableCacheStatsHolder.Node newNode = createMatchingImmutableCacheStatsHolderNode(currentNodeInOriginalTree); - parentInNewTree.getChildren().put(newNode.getDimensionValue(), newNode); - for (Node child : currentNodeInOriginalTree.children.values()) { - getImmutableCacheStatsHelper(child, newNode); - } - } - - private ImmutableCacheStatsHolder.Node createMatchingImmutableCacheStatsHolderNode(Node node) { - ImmutableCacheStats immutableCacheStats = node.getImmutableStats(); - boolean isLeafNode = node.getChildren().isEmpty(); - return new ImmutableCacheStatsHolder.Node(node.getDimensionValue(), !isLeafNode, immutableCacheStats); + return new ImmutableCacheStatsHolder(statsRoot.snapshot(), dimensionNames); } public void removeDimensions(List dimensionValues) { @@ -300,5 +280,16 @@ Node getChild(String dimensionValue) { Node createChild(String dimensionValue, boolean createMapInChild) { return children.computeIfAbsent(dimensionValue, (key) -> new Node(dimensionValue, createMapInChild)); } + + ImmutableCacheStatsHolder.Node snapshot() { + TreeMap snapshotChildren = null; + if (!children.isEmpty()) { + snapshotChildren = new TreeMap<>(); + for (Node child : children.values()) { + snapshotChildren.put(child.getDimensionValue(), child.snapshot()); + } + } + return new ImmutableCacheStatsHolder.Node(dimensionValue, snapshotChildren, getImmutableStats()); + } } } diff --git a/server/src/main/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolder.java b/server/src/main/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolder.java index 117ee06819c76..12e325046d83b 100644 --- a/server/src/main/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolder.java +++ b/server/src/main/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolder.java @@ -10,6 +10,7 @@ import org.opensearch.common.annotation.ExperimentalApi; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -75,17 +76,17 @@ static class Node { // The stats for this node. If a leaf node, corresponds to the stats for this combination of dimensions; if not, // contains the sum of its children's stats. - private ImmutableCacheStats stats; + private final ImmutableCacheStats stats; private static final Map EMPTY_CHILDREN_MAP = new HashMap<>(); - Node(String dimensionValue, boolean createChildrenMap, ImmutableCacheStats stats) { + Node(String dimensionValue, TreeMap snapshotChildren, ImmutableCacheStats stats) { this.dimensionValue = dimensionValue; - if (createChildrenMap) { - this.children = new TreeMap<>(); // This map should be ordered to enforce a consistent order in API response - } else { + this.stats = stats; + if (snapshotChildren == null) { this.children = EMPTY_CHILDREN_MAP; + } else { + this.children = Collections.unmodifiableMap(snapshotChildren); } - this.stats = stats; } Map getChildren() { @@ -96,10 +97,6 @@ public ImmutableCacheStats getStats() { return stats; } - public void setStats(ImmutableCacheStats stats) { - this.stats = stats; - } - public String getDimensionValue() { return dimensionValue; } diff --git a/server/src/test/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolderTests.java b/server/src/test/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolderTests.java index 9972bb5378860..933b8abd6e392 100644 --- a/server/src/test/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolderTests.java +++ b/server/src/test/java/org/opensearch/common/cache/stats/ImmutableCacheStatsHolderTests.java @@ -26,12 +26,12 @@ public void testGet() throws Exception { for (List dimensionValues : expected.keySet()) { CacheStats expectedCounter = expected.get(dimensionValues); - ImmutableCacheStats actualStatsHolder = CacheStatsHolderTests.getNode(dimensionValues, cacheStatsHolder.getStatsRoot()) + ImmutableCacheStats actualCacheStatsHolder = CacheStatsHolderTests.getNode(dimensionValues, cacheStatsHolder.getStatsRoot()) .getImmutableStats(); - ImmutableCacheStats actualCacheStats = getNode(dimensionValues, stats.getStatsRoot()).getStats(); + ImmutableCacheStats actualImmutableCacheStatsHolder = getNode(dimensionValues, stats.getStatsRoot()).getStats(); - assertEquals(expectedCounter.immutableSnapshot(), actualStatsHolder); - assertEquals(expectedCounter.immutableSnapshot(), actualCacheStats); + assertEquals(expectedCounter.immutableSnapshot(), actualCacheStatsHolder); + assertEquals(expectedCounter.immutableSnapshot(), actualImmutableCacheStatsHolder); } // test gets for total (this also checks sum-of-children logic)