address sagar's comments

Signed-off-by: Peter Alfonsi <petealft@amazon.com>

Author: Peter Alfonsi

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

@@ -1877,7 +1877,7 @@ public void run() {
                 logger.trace("running periodic field data cache cleanup");
             }
             try {
-                this.cache.clearMarkedKeys();
+                this.cache.clear();
             } catch (Exception e) {
                 logger.warn("Exception during periodic field data cache cleanup:", e);
             }

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

@@ -63,6 +63,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -154,23 +155,37 @@ public void clear(Index index, String field) {
         fieldsOfIndex.add(field);
     }
 
-    // This method can stay synchronized to avoid having multiple simultaneous cache iterators removing keys.
-    // This shouldn't cause performance impact as it's called only from the scheduled cache clear thread
-    // and is not blocking search traffic.
-    public void clearMarkedKeys() {
+    // The synchronized block is to avoid having multiple simultaneous cache iterators removing keys.
+    public void clear() {
         if (!(indicesToClear.isEmpty() && fieldsToClear.isEmpty())) {
+            // Copy marked indices/fields before iteration, and only remove keys matching the copies
+            // in case new entries are marked for cleanup mid-iteration
+            Set<Index> indicesToClearCopy = Set.copyOf(indicesToClear);
+            Map<Index, Set<String>> fieldsToClearCopy = new HashMap<>();
+            for (Map.Entry<Index, Set<String>> entry : fieldsToClear.entrySet()) {
+                fieldsToClearCopy.put(entry.getKey(), Set.copyOf(entry.getValue()));
+            }
+            // remove this way instead of clearing all, in case a new entry has been marked since copying
+            indicesToClear.removeAll(indicesToClearCopy);
+            for (Map.Entry<Index, Set<String>> entry : fieldsToClearCopy.entrySet()) {
+                Set<String> fieldsForIndex = fieldsToClear.get(entry.getKey());
+                if (fieldsForIndex != null) {
+                    fieldsForIndex.removeAll(entry.getValue());
+                }
+            }
+
             synchronized (this) {
                 for (Iterator<Key> iterator = getCache().keys().iterator(); iterator.hasNext();) {
                     Key key = iterator.next();
-                    if (indicesToClear.contains(key.indexCache.index)) {
+                    if (indicesToClearCopy.contains(key.indexCache.index)) {
                         try {
                             iterator.remove();
                         } catch (Exception e) {
                             logger.warn("Exception occurred while removing key from cache", e);
                         }
                         continue;
                     }
-                    Set<String> fieldsOfIndexToClear = fieldsToClear.get(key.indexCache.index);
+                    Set<String> fieldsOfIndexToClear = fieldsToClearCopy.get(key.indexCache.index);
                     if (fieldsOfIndexToClear != null && fieldsOfIndexToClear.contains(key.indexCache.fieldName)) {
                         try {
                             iterator.remove();
@@ -180,8 +195,6 @@ public void clearMarkedKeys() {
                     }
                 }
             }
-            indicesToClear.clear();
-            fieldsToClear.clear();
         }
         cache.refresh();
     }

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

@@ -544,7 +544,7 @@ public void remove() {
         // 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.clearMarkedKeys();
+        fdCacheSpy.clear();
         assertEquals(1, internalCacheSpy.count());
 
         reader.close();