diff --git a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java index a6ea36dc501e7..be777d901d719 100644 --- a/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java +++ b/modules/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java @@ -374,14 +374,14 @@ private V compute(ICacheKey key, LoadAwareCacheLoader, V> loader @Override public void invalidate(ICacheKey key) { - for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { - if (key.getDropStatsForDimensions()) { - List dimensionValues = statsHolder.getDimensionsWithTierValue(key.dimensions, cacheEntry.getValue().tierName); - statsHolder.removeDimensions(dimensionValues); // TODO: fix!! - } - if (key.key != null) { - try (ReleasableLock ignore = writeLock.acquire()) { - cacheEntry.getKey().invalidate(key); + if (key.getDropStatsForDimensions()) { + statsHolder.removeDimensions(key.dimensions); + } else { + for (Map.Entry, TierInfo> cacheEntry : caches.entrySet()) { + if (key.key != null) { + try (ReleasableLock ignore = writeLock.acquire()) { + cacheEntry.getKey().invalidate(key); + } } } } diff --git a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java index 051aa0354babe..33d8e41718790 100644 --- a/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java +++ b/modules/cache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java @@ -2112,6 +2112,60 @@ public void testTieredCacheDefaultSegmentCount() { assertTrue(VALID_SEGMENT_COUNT_VALUES.contains(tieredSpilloverCache.getNumberOfSegments())); } + public void testDropStatsForDimensions() throws Exception { + int onHeapCacheSize = randomIntBetween(300, 600); + int diskCacheSize = randomIntBetween(700, 1200); + int numberOfSegments = getNumberOfSegments(); + int keyValueSize = 50; + MockCacheRemovalListener removalListener = new MockCacheRemovalListener<>(); + TieredSpilloverCache tieredSpilloverCache = initializeTieredSpilloverCache( + keyValueSize, + diskCacheSize, + removalListener, + Settings.builder() + .put( + TieredSpilloverCacheSettings.TIERED_SPILLOVER_ONHEAP_STORE_SIZE.getConcreteSettingForNamespace( + CacheType.INDICES_REQUEST_CACHE.getSettingPrefix() + ).getKey(), + onHeapCacheSize * keyValueSize + "b" + ) + .build(), + 0, + numberOfSegments + ); + + List> usedKeys = new ArrayList<>(); + // Fill the cache, getting some entries + evictions for both tiers + int minMisses = (diskCacheSize + onHeapCacheSize) / keyValueSize + 10; + int numMisses = onHeapCacheSize + diskCacheSize + randomIntBetween(minMisses, minMisses + 50); + for (int iter = 0; iter < numMisses; iter++) { + ICacheKey key = getICacheKey(UUID.randomUUID().toString()); + usedKeys.add(key); + LoadAwareCacheLoader, String> tieredCacheLoader = getLoadAwareCacheLoader(); + tieredSpilloverCache.computeIfAbsent(key, tieredCacheLoader); + } + // Also do some random hits + Random rand = Randomness.get(); + int approxNumHits = 30; + for (int i = 0; i < approxNumHits; i++) { + LoadAwareCacheLoader, String> tieredCacheLoader = getLoadAwareCacheLoader(); + ICacheKey key = usedKeys.get(rand.nextInt(usedKeys.size())); + tieredSpilloverCache.computeIfAbsent(key, tieredCacheLoader); + } + + ImmutableCacheStats totalStats = tieredSpilloverCache.stats().getTotalStats(); + assertTrue(totalStats.getHits() > 0); + assertTrue(totalStats.getMisses() > 0); + assertTrue(totalStats.getEvictions() > 0); + + // Since all the keys have the same dimension values, except tiers, we only need to remove that one, and we expect all stats values + // should be 0 after that. + ICacheKey dropDimensionsKey = new ICacheKey<>(null, getMockDimensions()); + dropDimensionsKey.setDropStatsForDimensions(true); + tieredSpilloverCache.invalidate(dropDimensionsKey); + assertEquals(new ImmutableCacheStats(0, 0, 0, 0, 0), tieredSpilloverCache.stats().getTotalStats()); + } + private List getMockDimensions() { List dims = new ArrayList<>(); for (String dimensionName : dimensionNames) {