Use point-in-time key snapshot for fielddata cache cleanup

server/src/main/java/org/opensearch/common/cache/Cache.java

@@ -740,6 +740,26 @@ public void remove() {
         };
     }
 
+    /**
+     * A point-in-time copy of the keys in the cache, in no particular order. Unlike {@link #keys()}, the returned
+     * list is not backed by the cache and is safe to iterate under any concurrent mutation. The copy is weakly
+     * consistent: it contains every key that was present in the cache for the entire duration of the call, and may
+     * or may not reflect concurrent insertions and removals. Keys are copied one segment at a time while holding
+     * only that segment's read lock, so this method never blocks cache reads, briefly blocks writes to a single
+     * segment at a time, and does not acquire the LRU lock.
+     *
+     * @return a copy of the keys in the cache
+     */
+    public List<K> keysSnapshot() {
+        List<K> keys = new ArrayList<>(count);
+        for (CacheSegment<K, V> segment : segments) {
+            try (ReleasableLock ignored = segment.readLock.acquire()) {
+                keys.addAll(segment.map.keySet());
+            }
+        }
+        return keys;
+    }
+
     private class CacheIterator implements Iterator<Entry<K, V>> {
         private Entry<K, V> current;
         private Entry<K, V> next;

server/src/main/java/org/opensearch/indices/fielddata/cache/IndicesFieldDataCache.java

@@ -67,7 +67,6 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -219,7 +218,7 @@ public void clear(Index index, String field) {
         fieldsOfIndex.add(field);
     }
 
-    // The synchronized block is to avoid having multiple simultaneous cache iterators removing keys.
+    // The synchronized block is to avoid redundant concurrent sweeps; correctness does not depend on it.
     public void clear() {
         if (!(indicesToClear.isEmpty() && fieldsToClear.isEmpty())) {
             // Copy marked indices/fields before iteration, and only remove keys matching the copies
@@ -239,30 +238,39 @@ public void clear() {
             }
 
             synchronized (this) {
-                for (Iterator<Key> iterator = getCache().keys().iterator(); iterator.hasNext();) {
-                    Key key = iterator.next();
+                // Iterate a point-in-time copy of the keys rather than Cache.keys(): the live iterator walks the
+                // LRU list without holding the LRU lock, so a concurrent cache hit relinking the current entry to
+                // the head can silently skip entries. Since marks are consumed above before scanning, an entry
+                // skipped here would stay unreclaimed indefinitely. The snapshot is immune to LRU reordering, and
+                // invalidating a key that was concurrently removed is a no-op.
+                for (Key key : getCache().keysSnapshot()) {
                     if (indicesToClearCopy.contains(key.indexCache.index)) {
-                        removeKey(iterator);
+                        removeKey(key);
                         continue;
                     }
                     Set<String> fieldsOfIndexToClear = fieldsToClearCopy.get(key.indexCache.index);
                     if (fieldsOfIndexToClear != null && fieldsOfIndexToClear.contains(key.indexCache.fieldName)) {
-                        removeKey(iterator);
+                        removeKey(key);
                     }
                 }
             }
         }
         cache.refresh();
     }
 
-    private void removeKey(Iterator<Key> iterator) {
+    private void removeKey(Key key) {
         try {
-            iterator.remove();
+            invalidateKey(key);
         } catch (Exception e) {
             logger.warn("Exception occurred while removing key from cache", e);
         }
     }
 
+    // Package-private so tests can inject removal failures.
+    void invalidateKey(Key key) {
+        getCache().invalidate(key);
+    }
+
     /**
      * Computes a weight based on ramBytesUsed
      *
@@ -384,9 +392,8 @@ private void notifyOnCache(ShardId shardId, Accountable accountable) {
         @Override
         public void onClose(CacheKey key) throws IOException {
             // Invalidate the exact key synchronously rather than deferring to the periodic
-            // cache-scan cleanup: a reader closes only once, and the scan can miss entries under
-            // concurrent cache mutation, permanently retaining stale fielddata. Exact-key
-            // invalidation is O(1) so there is no need to batch it.
+            // cache-scan cleanup: a reader closes only once, and exact-key invalidation is O(1),
+            // so there is no need to batch it.
             // Don't call cache.refresh() here as it would have bad performance implications.
             nodeLevelCache.getCache().invalidate(new Key(this, key, null));
         }

server/src/test/java/org/opensearch/common/cache/CacheTests.java

@@ -956,6 +956,71 @@ public void testRemoveUsingValuesIterator() {
         }
     }
 
+    public void testKeysSnapshotContainsAllKeys() {
+        Cache<Integer, String> cache = CacheBuilder.<Integer, String>builder().build();
+        Set<Integer> expected = new HashSet<>();
+        for (int i = 0; i < numberOfEntries; i++) {
+            cache.put(i, Integer.toString(i));
+            expected.add(i);
+        }
+        List<Integer> snapshot = cache.keysSnapshot();
+        assertEquals(expected.size(), snapshot.size());
+        assertEquals(expected, new HashSet<>(snapshot));
+
+        // the snapshot is a copy: invalidating entries does not affect it
+        for (int i = 0; i < numberOfEntries / 2; i++) {
+            cache.invalidate(i);
+        }
+        assertEquals(expected, new HashSet<>(snapshot));
+    }
+
+    // Unlike Cache.keys(), whose iterator can silently skip entries when a concurrent cache hit relinks the
+    // current entry to the head of the LRU list, keysSnapshot must observe every key that is present for the
+    // whole duration of the call, regardless of concurrent reads promoting entries.
+    public void testKeysSnapshotUnderConcurrentPromotion() throws BrokenBarrierException, InterruptedException {
+        final Cache<Integer, String> cache = CacheBuilder.<Integer, String>builder().build();
+        final int entries = randomIntBetween(1000, 5000);
+        final Set<Integer> expected = new HashSet<>();
+        for (int i = 0; i < entries; i++) {
+            cache.put(i, Integer.toString(i));
+            expected.add(i);
+        }
+
+        final int numberOfReaders = randomIntBetween(2, 4);
+        final AtomicBoolean done = new AtomicBoolean(false);
+        final CyclicBarrier barrier = new CyclicBarrier(1 + numberOfReaders);
+        final List<Thread> threads = new ArrayList<>();
+        for (int t = 0; t < numberOfReaders; t++) {
+            Thread thread = new Thread(() -> {
+                try {
+                    barrier.await();
+                    Random random = new Random(randomLong());
+                    while (done.get() == false) {
+                        cache.get(random.nextInt(entries));
+                    }
+                } catch (BrokenBarrierException | InterruptedException e) {
+                    throw new AssertionError(e);
+                }
+            });
+            threads.add(thread);
+            thread.start();
+        }
+
+        barrier.await();
+        try {
+            for (int iteration = 0; iteration < 100; iteration++) {
+                List<Integer> snapshot = cache.keysSnapshot();
+                assertEquals(expected.size(), snapshot.size());
+                assertEquals(expected, new HashSet<>(snapshot));
+            }
+        } finally {
+            done.set(true);
+            for (Thread thread : threads) {
+                thread.join();
+            }
+        }
+    }
+
     public void testWithInvalidSegmentNumber() {
         assertThrows(
             "Number of segments for cache should be a power of two up-to 256",

server/src/test/java/org/opensearch/index/fielddata/IndexFieldDataServiceTests.java

@@ -45,7 +45,6 @@
 import org.apache.lucene.util.Accountable;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.SetOnce;
-import org.opensearch.common.cache.Cache;
 import org.opensearch.common.lucene.index.OpenSearchDirectoryReader;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
@@ -77,17 +76,14 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
 
 import org.mockito.Mockito;
 
 import static org.hamcrest.Matchers.containsString;
-import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
 public class IndexFieldDataServiceTests extends OpenSearchSingleNodeTestCase {
@@ -551,86 +547,6 @@ public void testRequireDocValuesOnBools() {
         doTestRequireDocValues(new BooleanFieldMapper.BooleanFieldType("field", true, false, false, null, Collections.emptyMap()));
     }
 
-    public void testExceptionWhileRemovingKey() throws Exception {
-        // Even if one key throws an error on removal, others should be cleared.
-        final IndexService indexService = createIndex("test");
-        final IndicesService indicesService = getInstanceFromNode(IndicesService.class);
-        IndicesFieldDataCache fdCache = indicesService.getIndicesFieldDataCache();
-        IndicesFieldDataCache fdCacheSpy = spy(fdCache);
-        final boolean[] hasThrownException = { false }; // Throw a runtime exception on the first remove
-        Cache<IndicesFieldDataCache.Key, Accountable> internalCache = fdCacheSpy.getCache();
-        Cache<IndicesFieldDataCache.Key, Accountable> internalCacheSpy = spy(internalCache);
-        doReturn(internalCacheSpy).when(fdCacheSpy).getCache();
-
-        // Add values to the (spy) cache before constructing our mocking iterator
-        final IndexFieldDataService ifdService = new IndexFieldDataService(
-            indexService.getIndexSettings(),
-            fdCacheSpy,
-            indicesService.getCircuitBreakerService(),
-            indexService.mapperService(),
-            indexService.getThreadPool()
-        );
-        final BuilderContext ctx = new BuilderContext(indexService.getIndexSettings().getSettings(), new ContentPath(1));
-        final MappedFieldType mapper1 = new TextFieldMapper.Builder("field_1", createDefaultIndexAnalyzers()).fielddata(true)
-            .build(ctx)
-            .fieldType();
-        final MappedFieldType mapper2 = new TextFieldMapper.Builder("field_2", createDefaultIndexAnalyzers()).fielddata(true)
-            .build(ctx)
-            .fieldType();
-        final IndexWriter writer = new IndexWriter(new ByteBuffersDirectory(), new IndexWriterConfig(new KeywordAnalyzer()));
-        Document doc = new Document();
-        doc.add(new StringField("field_1", "thisisastring", Store.NO));
-        doc.add(new StringField("field_2", "thisisanotherstring", Store.NO));
-        writer.addDocument(doc);
-        final IndexReader reader = DirectoryReader.open(writer);
-        IndexFieldData<?> ifd1 = ifdService.getForField(mapper1, "test", () -> { throw new UnsupportedOperationException(); });
-        IndexFieldData<?> ifd2 = ifdService.getForField(mapper2, "test", () -> { throw new UnsupportedOperationException(); });
-        LeafReaderContext leafReaderContext = reader.getContext().leaves().get(0);
-        LeafFieldData loadField1 = ifd1.load(leafReaderContext);
-        LeafFieldData loadField2 = ifd2.load(leafReaderContext);
-        assertBusy(() -> assertEquals(2, internalCacheSpy.count()));
-
-        Iterator<IndicesFieldDataCache.Key> realIterator = internalCacheSpy.keys().iterator();
-        Iterator<IndicesFieldDataCache.Key> erroringIterator = new Iterator<IndicesFieldDataCache.Key>() {
-            @Override
-            public boolean hasNext() {
-                return realIterator.hasNext();
-            }
-
-            @Override
-            public IndicesFieldDataCache.Key next() {
-                return realIterator.next();
-            }
-
-            @Override
-            public void remove() {
-                if (!hasThrownException[0]) {
-                    hasThrownException[0] = true;
-                    throw new UnsupportedOperationException("uh oh!");
-                }
-                realIterator.remove();
-            }
-        };
-        Iterable<IndicesFieldDataCache.Key> erroringIterable = () -> erroringIterator;
-        doReturn(erroringIterable).when(internalCacheSpy).keys();
-
-        // Clear the cache for both fields of the index, we expect 1 key will remain afterwards
-        ifdService.clear();
-        // Manually run the cache's clear keys loop
-        fdCacheSpy.clear();
-        assertEquals(1, internalCacheSpy.count());
-
-        reader.close();
-        loadField1.close();
-        loadField2.close();
-        try {
-            // Writer shutdown has some issue when trying to invalidate from the FD cache with the mock
-            writer.close();
-        } catch (AssertionError ignored) {}
-        // Ensure cache fully cleared before other tests in the suite begin
-        fdCacheSpy.close();
-    }
-
     public void testSetShardIdentityResolverRejectsNull() {
         ThreadPool threadPool = new TestThreadPool("test_set_resolver_null");
         try {