Made SingleDimensionCacheStats also take in tier dimensions

Signed-off-by: Peter Alfonsi <petealft@amazon.com>

plugins/cache-ehcache/src/main/java/org/opensearch/cache/store/disk/EhcacheDiskCache.java

@@ -24,6 +24,7 @@
 import org.opensearch.common.cache.RemovalReason;
 import org.opensearch.common.cache.stats.CacheStats;
 import org.opensearch.common.cache.ICacheKey;
+import org.opensearch.common.cache.stats.CacheStatsDimension;
 import org.opensearch.common.cache.stats.SingleDimensionCacheStats;
 import org.opensearch.common.cache.store.builders.ICacheBuilder;
 import org.opensearch.common.cache.store.enums.CacheStoreType;
@@ -148,7 +149,7 @@ private EhcacheDiskCache(Builder<K, V> builder) {
             this.valueSerializer);
         this.cache = buildCache(Duration.ofMillis(expireAfterAccess.getMillis()), builder);
         this.shardIdDimensionName = Objects.requireNonNull(builder.shardIdDimensionName, "Dimension name can't be null");
-        this.stats = new SingleDimensionCacheStats(shardIdDimensionName);
+        this.stats = new SingleDimensionCacheStats(shardIdDimensionName, CacheStatsDimension.TIER_DIMENSION_VALUE_DISK);
     }
 
     private Cache<ICacheKey, byte[]> buildCache(Duration expireAfterAccess, Builder<K, V> builder) {

plugins/cache-ehcache/src/test/java/org/opensearch/cache/store/disk/EhCacheDiskCacheTests.java

@@ -607,6 +607,36 @@ public void testMemoryTracking() throws Exception {
         }
     }
 
+    public void testGetStatsByTierName() throws Exception {
+        Settings settings = Settings.builder().build();
+        MockRemovalListener<String, String> mockRemovalListener = new MockRemovalListener<>();
+        ToLongBiFunction<ICacheKey<String>, String> weigher = getWeigher();
+        try (NodeEnvironment env = newNodeEnvironment(settings)) {
+            ICache<String, String> ehcacheTest = new EhcacheDiskCache.Builder<String, String>().setThreadPoolAlias("ehcacheTest")
+                .setStoragePath(env.nodePaths()[0].indicesPath.toString() + "/request_cache")
+                .setKeyType(String.class)
+                .setValueType(String.class)
+                .setKeySerializer(new StringSerializer())
+                .setValueSerializer(new StringSerializer())
+                .setShardIdDimensionName(dimensionName)
+                .setCacheType(CacheType.INDICES_REQUEST_CACHE)
+                .setSettings(settings)
+                .setExpireAfterAccess(TimeValue.MAX_VALUE)
+                .setMaximumWeightInBytes(CACHE_SIZE_IN_BYTES)
+                .setRemovalListener(mockRemovalListener)
+                .setWeigher(weigher)
+                .build();
+            int randomKeys = randomIntBetween(10, 100);
+            for (int i = 0; i < randomKeys; i++) {
+                ehcacheTest.put(getICacheKey(UUID.randomUUID().toString()), UUID.randomUUID().toString());
+            }
+            assertEquals(randomKeys, ehcacheTest.stats().getEntriesByDimensions(List.of(new CacheStatsDimension(CacheStatsDimension.TIER_DIMENSION_NAME, CacheStatsDimension.TIER_DIMENSION_VALUE_DISK))));
+            assertEquals(0, ehcacheTest.stats().getEntriesByDimensions(List.of(new CacheStatsDimension(CacheStatsDimension.TIER_DIMENSION_NAME, CacheStatsDimension.TIER_DIMENSION_VALUE_ON_HEAP))));
+
+            ehcacheTest.close();
+        }
+    }
+
     private static String generateRandomString(int length) {
         String characters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
         StringBuilder randomString = new StringBuilder(length);

server/src/main/java/org/opensearch/common/cache/stats/CacheStatsDimension.java

@@ -16,6 +16,10 @@
 import java.util.Objects;
 
 public class CacheStatsDimension implements Writeable {
+    // Values for tier dimensions, that are reused across CacheStats implementations
+    public static final String TIER_DIMENSION_NAME = "tier";
+    public static final String TIER_DIMENSION_VALUE_ON_HEAP = "on_heap";
+    public static final String TIER_DIMENSION_VALUE_DISK = "disk";
     public final String dimensionName;
     public final String dimensionValue;
     public CacheStatsDimension(String dimensionName, String dimensionValue) {

server/src/main/java/org/opensearch/common/cache/stats/SingleDimensionCacheStats.java

@@ -14,14 +14,15 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
 /**
- * A CacheStats implementation for caches that aggregate over a single dimension.
+ * A CacheStats implementation for caches that aggregate over a single dimension, as well as holding a tier dimension.
  * For example, caches in the IndicesRequestCache only aggregate over ShardId value.
  */
 public class SingleDimensionCacheStats implements CacheStats {
@@ -39,10 +40,12 @@ public class SingleDimensionCacheStats implements CacheStats {
     private final CounterMetric totalMemorySize;
     private final CounterMetric totalEntries;
 
-    // The allowed dimension name. This stats only allows a single dimension name
-    private final String allowedDimensionName;
+    // The allowed dimension name. This stats only allows a single dimension name. Package-private for testing.
+    final String allowedDimensionName;
+    // The value of the tier dimension for entries in this Stats object. Package-private for testing.
+    final String tierDimensionValue;
 
-    public SingleDimensionCacheStats(String allowedDimensionName) {
+    public SingleDimensionCacheStats(String allowedDimensionName, String tierDimensionValue) {
         this.hitsMap = new ConcurrentHashMap<>();
         this.missesMap = new ConcurrentHashMap<>();
         this.evictionsMap = new ConcurrentHashMap<>();
@@ -56,6 +59,7 @@ public SingleDimensionCacheStats(String allowedDimensionName) {
         this.totalEntries = new CounterMetric();
 
         this.allowedDimensionName = allowedDimensionName;
+        this.tierDimensionValue = tierDimensionValue;
     }
 
     public SingleDimensionCacheStats(StreamInput in) throws IOException {
@@ -77,6 +81,7 @@ public SingleDimensionCacheStats(StreamInput in) throws IOException {
         totalEntries.inc(in.readVLong());
 
         this.allowedDimensionName = in.readString();
+        this.tierDimensionValue = in.readString();
     }
 
     @Override
@@ -104,7 +109,34 @@ public long getTotalEntries() {
         return this.totalEntries.count();
     }
 
-    private long internalGetByDimension(List<CacheStatsDimension> dimensions, Map<String, CounterMetric> metricsMap) {
+    private long internalGetByDimension(List<CacheStatsDimension> dimensions, Map<String, CounterMetric> metricsMap, CounterMetric totalMetric) {
+        CacheStatsDimension tierDimension = getTierDimensionIfPresent(dimensions);
+        if (tierDimension != null) {
+            // This get request includes a tier dimension. Return values only if the tier dimension value
+            // matches the one for this stats object, otherwise return 0
+            assert dimensions.size() == 1 || dimensions.size() == 2; // There can be at most one non-tier dimension value
+            if (tierDimension.dimensionValue.equals(tierDimensionValue)) {
+                // The list passed in may not be mutable; create a mutable copy to remove the tier dimension
+                ArrayList<CacheStatsDimension> modifiedDimensions = new ArrayList<>(dimensions);
+                modifiedDimensions.remove(tierDimension);
+
+                if (modifiedDimensions.size() == 1){
+                    return internalGetHelper(modifiedDimensions, metricsMap);
+                } else {
+                    return totalMetric.count();
+                }
+
+            } else {
+                // Return 0 for incorrect tier value
+                return 0;
+            }
+        } else {
+            // This get request doesn't include a tier dimension. Return the appropriate values.
+            return internalGetHelper(dimensions, metricsMap);
+        }
+    }
+
+    private long internalGetHelper(List<CacheStatsDimension> dimensions, Map<String, CounterMetric> metricsMap) {
         assert dimensions.size() == 1;
         CounterMetric counter = metricsMap.get(dimensions.get(0).dimensionValue);
         if (counter == null) {
@@ -113,29 +145,41 @@ private long internalGetByDimension(List<CacheStatsDimension> dimensions, Map<St
         return counter.count();
     }
 
+    /**
+     * Returns the dimension that represents a tier value, if one is present. Otherwise return null.
+     */
+    private CacheStatsDimension getTierDimensionIfPresent(List<CacheStatsDimension> dimensions) {
+        for (CacheStatsDimension dim : dimensions) {
+            if (dim.dimensionName.equals(CacheStatsDimension.TIER_DIMENSION_NAME)) {
+                return dim;
+            }
+        }
+        return null;
+    }
+
     @Override
     public long getHitsByDimensions(List<CacheStatsDimension> dimensions) {
-        return internalGetByDimension(dimensions, hitsMap);
+        return internalGetByDimension(dimensions, hitsMap, totalHits);
     }
 
     @Override
     public long getMissesByDimensions(List<CacheStatsDimension> dimensions) {
-        return internalGetByDimension(dimensions, missesMap);
+        return internalGetByDimension(dimensions, missesMap, totalMisses);
     }
 
     @Override
     public long getEvictionsByDimensions(List<CacheStatsDimension> dimensions) {
-        return internalGetByDimension(dimensions, evictionsMap);
+        return internalGetByDimension(dimensions, evictionsMap, totalEvictions);
     }
 
     @Override
     public long getMemorySizeByDimensions(List<CacheStatsDimension> dimensions) {
-        return internalGetByDimension(dimensions, memorySizeMap);
+        return internalGetByDimension(dimensions, memorySizeMap, totalMemorySize);
     }
 
     @Override
     public long getEntriesByDimensions(List<CacheStatsDimension> dimensions) {
-        return internalGetByDimension(dimensions, entriesMap);
+        return internalGetByDimension(dimensions, entriesMap, totalEntries);
     }
 
     private boolean checkDimensionList(List<CacheStatsDimension> dimensions) {
@@ -199,10 +243,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeVLong(totalEntries.count());
 
         out.writeString(allowedDimensionName);
-    }
-
-    public String getAllowedDimensionName() {
-        return allowedDimensionName;
+        out.writeString(tierDimensionValue);
     }
 
     // For converting to StreamOutput/StreamInput, write maps of longs rather than CounterMetrics which don't support writing

server/src/main/java/org/opensearch/common/cache/store/OpenSearchOnHeapCache.java

@@ -17,6 +17,7 @@
 import org.opensearch.common.cache.RemovalReason;
 import org.opensearch.common.cache.stats.CacheStats;
 import org.opensearch.common.cache.ICacheKey;
+import org.opensearch.common.cache.stats.CacheStatsDimension;
 import org.opensearch.common.cache.stats.SingleDimensionCacheStats;
 import org.opensearch.common.cache.store.builders.ICacheBuilder;
 import org.opensearch.common.settings.Settings;
@@ -49,7 +50,7 @@ public OpenSearchOnHeapCache(Builder<K, V> builder) {
         }
         cache = cacheBuilder.build();
         String dimensionName = Objects.requireNonNull(builder.shardIdDimensionName, "Shard id dimension name can't be null");
-        this.stats = new SingleDimensionCacheStats(dimensionName);
+        this.stats = new SingleDimensionCacheStats(dimensionName, CacheStatsDimension.TIER_DIMENSION_VALUE_ON_HEAP);
         this.removalListener = builder.getRemovalListener();
     }
 

server/src/test/java/org/opensearch/common/cache/stats/SingleDimensionCacheStatsTests.java

@@ -22,8 +22,9 @@
 
 public class SingleDimensionCacheStatsTests extends OpenSearchTestCase {
     private final String dimensionName = "shardId";
+    private final String tierName = "test_tier";
     public void testAddAndGet() throws Exception {
-        StatsAndExpectedResults statsAndExpectedResults = getPopulatedStats();
+        StatsAndExpectedResults statsAndExpectedResults = getPopulatedStats(tierName);
         SingleDimensionCacheStats stats = statsAndExpectedResults.stats;
 
         checkShardResults(statsAndExpectedResults);
@@ -32,10 +33,34 @@ public void testAddAndGet() throws Exception {
         // Check values returned for a nonexistent dimension value or name return 0
         assertEquals(0, stats.getHitsByDimensions(List.of(new CacheStatsDimension(dimensionName, "nonexistent"))));
         assertEquals(0, stats.getHitsByDimensions(List.of(new CacheStatsDimension("nonexistentName", "nonexistentValue"))));
+
+        // Check sending too many values causes an assertion error
+        assertThrows(AssertionError.class, () -> stats.getHitsByDimensions(List.of(getDim(0), new CacheStatsDimension("test", "value"))));
+    }
+
+    public void testTierFiltering() throws Exception {
+        StatsAndExpectedResults statsAndExpectedResults = getPopulatedStats(tierName);
+        SingleDimensionCacheStats stats = statsAndExpectedResults.stats;
+
+        // Values should be returned if the tier dimension value matches the one passed to SingleDimensionCacheStats. Otherwise we should get 0.
+        CacheStatsDimension matchingTierDim = new CacheStatsDimension(CacheStatsDimension.TIER_DIMENSION_NAME, tierName);
+        CacheStatsDimension nonMatchingTierDim = new CacheStatsDimension(CacheStatsDimension.TIER_DIMENSION_NAME, "another_tier");
+
+        assertEquals(stats.getTotalHits(), stats.getHitsByDimensions(List.of(matchingTierDim)));
+        assertEquals(0, stats.getHitsByDimensions(List.of(nonMatchingTierDim)));
+        for (int i = 0; i < statsAndExpectedResults.numShardIds; i++) {
+            assertEquals(stats.getHitsByDimensions(List.of(getDim(i))), stats.getHitsByDimensions(List.of(getDim(i), matchingTierDim)));
+            assertEquals(stats.getHitsByDimensions(List.of(getDim(i))), stats.getHitsByDimensions(List.of(matchingTierDim, getDim(i))));
+            assertEquals(0, stats.getHitsByDimensions(List.of(getDim(i), nonMatchingTierDim)));
+            assertEquals(0, stats.getHitsByDimensions(List.of(nonMatchingTierDim, getDim(i))));
+
+        }
+        // Check sending too many values causes an assertion error
+        assertThrows(AssertionError.class, () -> stats.getHitsByDimensions(List.of(getDim(0), matchingTierDim, new CacheStatsDimension("test", "value"))));
     }
 
     public void testSerialization() throws Exception {
-        StatsAndExpectedResults statsAndExpectedResults = getPopulatedStats();
+        StatsAndExpectedResults statsAndExpectedResults = getPopulatedStats(tierName);
         SingleDimensionCacheStats stats = statsAndExpectedResults.stats;
         Map<String, Map<String, Long>> expectedResults = statsAndExpectedResults.expectedShardResults;
 
@@ -47,7 +72,8 @@ public void testSerialization() throws Exception {
         StatsAndExpectedResults deserializedStatsAndExpectedResults = new StatsAndExpectedResults(deserialized, expectedResults, statsAndExpectedResults.numShardIds);
         checkShardResults(deserializedStatsAndExpectedResults);
         checkTotalResults(deserializedStatsAndExpectedResults);
-        assertEquals(deserialized.getAllowedDimensionName(), stats.getAllowedDimensionName());
+        assertEquals(deserialized.allowedDimensionName, stats.allowedDimensionName);
+        assertEquals(deserialized.tierDimensionValue, stats.tierDimensionValue);
     }
 
     private CacheStatsDimension getDim(int i) {
@@ -68,8 +94,8 @@ private long sumMap(Map<String, Long> inputMap) {
         return result;
     }
 
-    private StatsAndExpectedResults getPopulatedStats() {
-        SingleDimensionCacheStats stats = new SingleDimensionCacheStats(dimensionName);
+    private StatsAndExpectedResults getPopulatedStats(String tierName) {
+        SingleDimensionCacheStats stats = new SingleDimensionCacheStats(dimensionName, tierName);
 
         int numShardIds = 10;
         Map<String, Long> expectedHits = new HashMap<>();