Make ICache keys() iteration safe under concurrent mutation

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

@@ -89,7 +89,7 @@ public V computeIfAbsent(ICacheKey<K> key, LoadAwareCacheLoader<ICacheKey<K>, V>
     public void invalidate(ICacheKey<K> key) {
         V value = this.cache.remove(key);
         if (value != null) {
-            removalListener.onRemoval(new RemovalNotification<>(key, cache.get(key), RemovalReason.INVALIDATED));
+            removalListener.onRemoval(new RemovalNotification<>(key, value, RemovalReason.INVALIDATED));
         }
     }
 

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

@@ -31,6 +31,7 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.common.unit.ByteSizeValue;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -132,9 +133,16 @@ public void invalidateAll() {
         cacheStatsHolder.reset();
     }
 
+    /**
+     * Returns a point-in-time copy of the keys in the cache rather than the live {@link Cache#keys()} view. The
+     * live iterator walks the LRU linked list without holding the LRU lock, so a concurrent cache hit relinking
+     * a not-yet-visited entry to the head of the list can cause iteration to silently skip it. The copy is
+     * unmodifiable, so {@link java.util.Iterator#remove()} — which would be a silent no-op against a copy —
+     * fails fast; use {@link #invalidate(ICacheKey)} to remove entries.
+     */
     @Override
     public Iterable<ICacheKey<K>> keys() {
-        return cache.keys();
+        return Collections.unmodifiableList(cache.keysSnapshot());
     }
 
     @Override

server/src/main/java/org/opensearch/indices/IndicesRequestCache.java

@@ -785,23 +785,26 @@ private synchronized void cleanCache(double stalenessThreshold) {
 
             Set<List<String>> dimensionListsToDrop = new HashSet<>();
 
-            for (Iterator<ICacheKey<Key>> iterator = cache.keys().iterator(); iterator.hasNext();) {
-                ICacheKey<Key> key = iterator.next();
+            // The cleanup marks are consumed above before this scan, so a key skipped here would stay in the
+            // cache until size-based eviction reaches it. Remove matches with exact-key invalidate() rather
+            // than through the keys() iterator, which may be walking a point-in-time copy that does not support
+            // removal (see OpenSearchOnHeapCache#keys()); invalidating a concurrently removed key is a no-op.
+            for (ICacheKey<Key> key : cache.keys()) {
                 Key delegatingKey = key.key;
                 Tuple<ShardId, Integer> shardIdInfo = new Tuple<>(delegatingKey.shardId, delegatingKey.indexShardHashCode);
                 if (cleanupKeysFromFullClean.contains(shardIdInfo) || cleanupKeysFromClosedShards.contains(shardIdInfo)) {
-                    iterator.remove();
+                    cache.invalidate(key);
                 } else {
                     CacheEntity cacheEntity = cacheEntityLookup.apply(delegatingKey.shardId).orElse(null);
                     if (cacheEntity == null) {
                         // If cache entity is null, it means that index or shard got deleted/closed meanwhile.
                         // So we will delete this key.
                         dimensionListsToDrop.add(key.dimensions);
-                        iterator.remove();
+                        cache.invalidate(key);
                     } else {
                         CleanupKey cleanupKey = new CleanupKey(cacheEntity, delegatingKey.readerCacheKeyId);
                         if (cleanupKeysFromOutdatedReaders.contains(cleanupKey)) {
-                            iterator.remove();
+                            cache.invalidate(key);
                         }
                     }
                 }

server/src/test/java/org/opensearch/common/cache/store/OpenSearchOnHeapCacheTests.java

@@ -25,8 +25,11 @@
 import org.opensearch.test.OpenSearchTestCase;
 
 import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Random;
+import java.util.Set;
 import java.util.UUID;
 
 import static org.opensearch.common.cache.store.settings.OpenSearchOnHeapCacheSettings.MAXIMUM_SIZE_IN_BYTES_KEY;
@@ -251,6 +254,48 @@ public void onRemoval(RemovalNotification<ICacheKey<K>, V> notification) {
         }
     }
 
+    /**
+     * keys() must be a point-in-time view that is safe to iterate under concurrent mutation: a cache hit
+     * mid-iteration relinks the accessed entry to the head of the underlying LRU list, and the live
+     * Cache#keys() iterator silently skips an entry promoted from behind the not-yet-visited portion of the
+     * list. This test fails if keys() exposes the live LRU iteration and passes with a point-in-time copy.
+     */
+    public void testKeysIsSafeToIterateUnderConcurrentPromotion() throws Exception {
+        MockRemovalListener<String, String> listener = new MockRemovalListener<>();
+        OpenSearchOnHeapCache<String, String> cache = getCache(100, listener, true);
+        List<ICacheKey<String>> insertedKeys = new ArrayList<>();
+        for (int i = 0; i < 5; i++) {
+            ICacheKey<String> key = getICacheKey("key" + i);
+            cache.computeIfAbsent(key, getLoadAwareCacheLoader());
+            insertedKeys.add(key);
+        }
+
+        Iterator<ICacheKey<String>> iterator = cache.keys().iterator();
+        Set<ICacheKey<String>> observedKeys = new HashSet<>();
+        observedKeys.add(iterator.next());
+        // A hit on the least-recently-used key relinks it at the head of the LRU list, behind a live
+        // iterator's cursor
+        assertNotNull(cache.get(insertedKeys.get(0)));
+        while (iterator.hasNext()) {
+            observedKeys.add(iterator.next());
+        }
+        assertEquals(new HashSet<>(insertedKeys), observedKeys);
+    }
+
+    public void testKeysIteratorDoesNotSupportRemoval() throws Exception {
+        MockRemovalListener<String, String> listener = new MockRemovalListener<>();
+        OpenSearchOnHeapCache<String, String> cache = getCache(100, listener, true);
+        ICacheKey<String> key = getICacheKey("key");
+        cache.computeIfAbsent(key, getLoadAwareCacheLoader());
+
+        Iterator<ICacheKey<String>> iterator = cache.keys().iterator();
+        iterator.next();
+        // Removing through the keys() iterator would be a silent no-op on the point-in-time copy; it must
+        // fail fast so callers use invalidate(key) instead
+        expectThrows(UnsupportedOperationException.class, iterator::remove);
+        assertEquals(1, cache.count());
+    }
+
     private ICacheKey<String> getICacheKey(String key) {
         List<String> dims = new ArrayList<>();
         for (String dimName : dimensionNames) {

server/src/test/java/org/opensearch/indices/IndicesRequestCacheTests.java

@@ -115,6 +115,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.emptySet;
@@ -412,6 +413,72 @@ public void testCacheCleanupBasedOnStaleThreshold_StalenessEqualToThreshold() th
         IOUtils.close(secondReader);
     }
 
+    /**
+     * Reproduces stale-entry retention in the cleanup sweep: a cache hit during the sweep relinks the
+     * accessed entry to the head of the on-heap cache's LRU list, and a live LRU iteration silently skips
+     * a not-yet-visited entry promoted behind its cursor. Because the cleanup marks are consumed before
+     * the scan, the skipped entry is never revisited and survives until size-based eviction. The hit is
+     * triggered deterministically from the entity lookup of the first swept entry, mimicking a concurrent
+     * search request. This test fails if the sweep iterates the live LRU view of the on-heap cache's keys
+     * and passes with the point-in-time snapshot.
+     */
+    public void testCleanCacheIsNotDefeatedByCacheHitPromotingEntryMidSweep() throws Exception {
+        threadPool = getThreadPool();
+        IndexShard secondShard = createIndex("second-test").getShard(0);
+        IndicesService indicesService = getInstanceFromNode(IndicesService.class);
+        AtomicReference<Runnable> onFirstLookup = new AtomicReference<>();
+        try (NodeEnvironment env = newNodeEnvironment(Settings.EMPTY)) {
+            cache = new IndicesRequestCache(Settings.EMPTY, shardId -> {
+                Runnable hook = onFirstLookup.getAndSet(null);
+                if (hook != null) {
+                    hook.run();
+                }
+                return indicesService.indicesRequestCache.cacheEntityLookup.apply(shardId);
+            },
+                new CacheModule(new ArrayList<>(), Settings.EMPTY).getCacheService(),
+                threadPool,
+                ClusterServiceUtils.createClusterService(threadPool),
+                env
+            );
+        }
+        writer.addDocument(newDoc(0, "foo"));
+        DirectoryReader readerB1 = getReader(writer, secondShard.shardId());
+        DirectoryReader readerA1 = getReader(writer, indexShard.shardId());
+        DirectoryReader readerA2 = getReader(writer, indexShard.shardId());
+        DirectoryReader readerB2 = getReader(writer, secondShard.shardId());
+
+        // LRU order after insertion, head to tail: B2, A2, A1, B1
+        cache.getOrCompute(getEntity(secondShard), getLoader(readerB1), readerB1, getTermBytes());
+        cache.getOrCompute(getEntity(indexShard), getLoader(readerA1), readerA1, getTermBytes());
+        cache.getOrCompute(getEntity(indexShard), getLoader(readerA2), readerA2, getTermBytes());
+        cache.getOrCompute(getEntity(secondShard), getLoader(readerB2), readerB2, getTermBytes());
+        assertEquals(4, cache.count());
+
+        // Mark all of the second shard's entries for cleanup; the marks are consumed by the next sweep
+        cache.clear(getEntity(secondShard));
+
+        // While the sweep processes its first entry, a concurrent request hits the second shard's
+        // not-yet-visited entry, relinking it at the head of the LRU list. The loader throws so the hit
+        // cannot re-insert the entry if the sweep already removed it (snapshot order is arbitrary).
+        onFirstLookup.set(() -> {
+            try {
+                cache.getOrCompute(
+                    getEntity(secondShard),
+                    () -> { throw new IOException("no reload during sweep"); },
+                    readerB1,
+                    getTermBytes()
+                );
+            } catch (Exception ignored) {
+                // a miss means the sweep already removed the entry; nothing to promote
+            }
+        });
+        cache.cacheCleanupManager.cleanCache();
+
+        // Both second-shard entries must be removed; only the two live first-shard entries remain
+        assertEquals(2, cache.count());
+        IOUtils.close(readerB1, readerA1, readerA2, readerB2);
+    }
+
     // when a cache entry that is Stale is evicted for any reason, we have to deduct the count from our staleness count
     public void testStaleCount_OnRemovalNotificationOfStaleKey_DecrementsStaleCount() throws Exception {
         threadPool = getThreadPool();