Integrate fix with TSC

Signed-off-by: Peter Alfonsi <petealft@amazon.com>
opensearch-project · Nov 4, 2024 · afd22e8 · afd22e8
1 parent 50707b0
commit afd22e8
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 8 deletions.
diff --git a/...les/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java b/...les/cache-common/src/main/java/org/opensearch/cache/common/tier/TieredSpilloverCache.java
@@ -374,14 +374,14 @@ private V compute(ICacheKey<K> key, LoadAwareCacheLoader<ICacheKey<K>, V> loader
 
         @Override
         public void invalidate(ICacheKey<K> key) {
-            for (Map.Entry<ICache<K, V>, TierInfo> cacheEntry : caches.entrySet()) {
-                if (key.getDropStatsForDimensions()) {
-                    List<String> 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<ICache<K, V>, TierInfo> cacheEntry : caches.entrySet()) {
+                    if (key.key != null) {
+                        try (ReleasableLock ignore = writeLock.acquire()) {
+                            cacheEntry.getKey().invalidate(key);
+                        }
                     }
                 }
             }

diff --git a/...ache-common/src/test/java/org/opensearch/cache/common/tier/TieredSpilloverCacheTests.java b/...ache-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<String, String> removalListener = new MockCacheRemovalListener<>();
+        TieredSpilloverCache<String, String> 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<ICacheKey<String>> 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<String> key = getICacheKey(UUID.randomUUID().toString());
+            usedKeys.add(key);
+            LoadAwareCacheLoader<ICacheKey<String>, 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<ICacheKey<String>, String> tieredCacheLoader = getLoadAwareCacheLoader();
+            ICacheKey<String> 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<String> dropDimensionsKey = new ICacheKey<>(null, getMockDimensions());
+        dropDimensionsKey.setDropStatsForDimensions(true);
+        tieredSpilloverCache.invalidate(dropDimensionsKey);
+        assertEquals(new ImmutableCacheStats(0, 0, 0, 0, 0), tieredSpilloverCache.stats().getTotalStats());
+    }
+
     private List<String> getMockDimensions() {
         List<String> dims = new ArrayList<>();
         for (String dimensionName : dimensionNames) {