[Tiered Caching] Additional ITs for cache stats

modules/cache-common/src/internalClusterTest/java/org.opensearch.cache.common.tier/TieredSpilloverCacheIT.java → modules/cache-common/src/internalClusterTest/java/org/opensearch/cache/common/tier/TieredSpilloverCacheIT.java

@@ -65,7 +65,7 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
         return Arrays.asList(TieredSpilloverCachePlugin.class, MockDiskCachePlugin.class);
     }
 
-    private Settings defaultSettings(String onHeapCacheSizeInBytesOrPecentage) {
+    static Settings defaultSettings(String onHeapCacheSizeInBytesOrPercentage) {
         return Settings.builder()
             .put(FeatureFlags.PLUGGABLE_CACHE, "true")
             .put(
@@ -88,7 +88,7 @@ private Settings defaultSettings(String onHeapCacheSizeInBytesOrPecentage) {
                 OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE)
                     .get(MAXIMUM_SIZE_IN_BYTES_KEY)
                     .getKey(),
-                onHeapCacheSizeInBytesOrPecentage
+                onHeapCacheSizeInBytesOrPercentage
             )
             .build();
     }

modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java

@@ -327,6 +327,7 @@ private Function<ICacheKey<K>, Tuple<V, String>> getValueFromTieredCache(boolean
     void handleRemovalFromHeapTier(RemovalNotification<ICacheKey<K>, V> notification) {
         ICacheKey<K> key = notification.getKey();
         boolean wasEvicted = SPILLOVER_REMOVAL_REASONS.contains(notification.getRemovalReason());
+        boolean countEvictionTowardsTotal = false; // Don't count this eviction towards the cache's total if it ends up in the disk tier
         if (caches.get(diskCache).isEnabled() && wasEvicted && evaluatePolicies(notification.getValue())) {
             try (ReleasableLock ignore = writeLock.acquire()) {
                 diskCache.put(key, notification.getValue()); // spill over to the disk tier and increment its stats
@@ -336,21 +337,28 @@ void handleRemovalFromHeapTier(RemovalNotification<ICacheKey<K>, V> notification
             // If the value is not going to the disk cache, send this notification to the TSC's removal listener
             // as the value is leaving the TSC entirely
             removalListener.onRemoval(notification);
+            countEvictionTowardsTotal = true;
         }
-        updateStatsOnRemoval(TIER_DIMENSION_VALUE_ON_HEAP, wasEvicted, key, notification.getValue());
+        updateStatsOnRemoval(TIER_DIMENSION_VALUE_ON_HEAP, wasEvicted, key, notification.getValue(), countEvictionTowardsTotal);
     }
 
     void handleRemovalFromDiskTier(RemovalNotification<ICacheKey<K>, V> notification) {
         // Values removed from the disk tier leave the TSC entirely
         removalListener.onRemoval(notification);
         boolean wasEvicted = SPILLOVER_REMOVAL_REASONS.contains(notification.getRemovalReason());
-        updateStatsOnRemoval(TIER_DIMENSION_VALUE_DISK, wasEvicted, notification.getKey(), notification.getValue());
+        updateStatsOnRemoval(TIER_DIMENSION_VALUE_DISK, wasEvicted, notification.getKey(), notification.getValue(), true);
     }
 
-    void updateStatsOnRemoval(String removedFromTierValue, boolean wasEvicted, ICacheKey<K> key, V value) {
+    void updateStatsOnRemoval(
+        String removedFromTierValue,
+        boolean wasEvicted,
+        ICacheKey<K> key,
+        V value,
+        boolean countEvictionTowardsTotal
+    ) {
         List<String> dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, removedFromTierValue);
         if (wasEvicted) {
-            statsHolder.incrementEvictions(dimensionValues);
+            statsHolder.incrementEvictions(dimensionValues, countEvictionTowardsTotal);
         }
         statsHolder.decrementItems(dimensionValues);
         statsHolder.decrementSizeInBytes(dimensionValues, weigher.applyAsLong(key, value));

modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCacheStatsHolder.java

@@ -105,20 +105,29 @@
         internalIncrement(dimensionValues, missIncrementer, true);
     }
 
+    /**
+     * This method shouldn't be used in this class. Instead, use incrementEvictions(dimensionValues, includeInTotal)
+     * which specifies whether the eviction should be included in the cache's total evictions, or if it should
+     * just count towards that tier's evictions.
+     * @param dimensionValues The dimension values
+     */
     @Override
     public void incrementEvictions(List<String> dimensionValues) {
-        final String tierValue = validateTierDimensionValue(dimensionValues);
+        throw new UnsupportedOperationException(
+            "TieredSpilloverCacheHolder must specify whether to include an eviction in the total cache stats. Use incrementEvictions(List<String> dimensionValues, boolean includeInTotal)"
+        );
+    }
 
-        // If the disk tier is present, only evictions from the disk tier should be included in total values.
+    /**
+     * Increment evictions for this set of dimension values.
+     * @param dimensionValues The dimension values
+     * @param includeInTotal Whether to include this eviction in the total for the whole cache's evictions
+     */
+    public void incrementEvictions(List<String> dimensionValues, boolean includeInTotal) {
+        validateTierDimensionValue(dimensionValues);
+        // If we count this eviction towards the total, we should increment all ancestor nodes. If not, only increment the leaf node.
         Consumer<DefaultCacheStatsHolder.Node> evictionsIncrementer = (node) -> {
-            if (tierValue.equals(TIER_DIMENSION_VALUE_ON_HEAP) && diskCacheEnabled) {
-                // If on-heap tier, increment only the leaf node corresponding to the on heap values; not the total values in its parent
-                // nodes
-                if (node.isAtLowestLevel()) {
-                    node.incrementEvictions();
-                }
-            } else {
-                // If disk tier, or on-heap tier with a disabled disk tier, increment the leaf node and its parents
+            if (includeInTotal || node.isAtLowestLevel()) {
                 node.incrementEvictions();
             }
         };

modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java

@@ -926,14 +926,14 @@ public void testDiskTierPolicies() throws Exception {
         MockCacheRemovalListener<String, String> removalListener = new MockCacheRemovalListener<>();
         TieredSpilloverCache<String, String> tieredSpilloverCache = intializeTieredSpilloverCache(
             keyValueSize,
-            100,
+            keyValueSize * 100,
             removalListener,
             Settings.builder()
                 .put(
                     OpenSearchOnHeapCacheSettings.getSettingListForCacheType(CacheType.INDICES_REQUEST_CACHE)
                         .get(MAXIMUM_SIZE_IN_BYTES_KEY)
                         .getKey(),
-                    onHeapCacheSize * 50 + "b"
+                    onHeapCacheSize * keyValueSize + "b"
                 )
                 .build(),
             0,
@@ -955,6 +955,7 @@ public void testDiskTierPolicies() throws Exception {
 
         LoadAwareCacheLoader<ICacheKey<String>, String> loader = getLoadAwareCacheLoader(keyValuePairs);
 
+        int expectedEvictions = 0;
         for (String key : keyValuePairs.keySet()) {
             ICacheKey<String> iCacheKey = getICacheKey(key);
             Boolean expectedOutput = expectedOutputs.get(key);
@@ -967,8 +968,15 @@ public void testDiskTierPolicies() throws Exception {
             } else {
                 // Should miss as heap tier size = 0 and the policy rejected it
                 assertNull(result);
+                expectedEvictions++;
             }
         }
+
+        // We expect values that were evicted from the heap tier and not allowed into the disk tier by the policy
+        // to count towards total evictions
+        assertEquals(keyValuePairs.size(), getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_ON_HEAP));
+        assertEquals(0, getEvictionsForTier(tieredSpilloverCache, TIER_DIMENSION_VALUE_DISK)); // Disk tier is large enough for no evictions
+        assertEquals(expectedEvictions, getTotalStatsSnapshot(tieredSpilloverCache).getEvictions());
     }
 
     public void testTookTimePolicyFromFactory() throws Exception {
@@ -1493,6 +1501,11 @@ private ImmutableCacheStats getStatsSnapshotForTier(TieredSpilloverCache<?, ?> t
         return snapshot;
     }
 
+    private ImmutableCacheStats getTotalStatsSnapshot(TieredSpilloverCache<?, ?> tsc) throws IOException {
+        ImmutableCacheStatsHolder cacheStats = tsc.stats(new String[0]);
+        return cacheStats.getStatsForDimensionValues(List.of());
+    }
+
     // Duplicated here from EhcacheDiskCacheTests.java, we can't add a dependency on that plugin
     static class StringSerializer implements Serializer<String, byte[]> {
         private final Charset charset = StandardCharsets.UTF_8;