ixing the race condition between Merge and Refresh flow

Signed-off-by: Sagar Darji <darsaga@amazon.com>

Author: Sagar Darji

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneReaderManager.java

@@ -9,9 +9,12 @@
 package org.opensearch.be.lucene;
 
 import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentReader;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.DataFormat;
 import org.opensearch.index.engine.exec.EngineReaderManager;
+import org.opensearch.index.engine.exec.Segment;
 import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 
 import java.io.IOException;
@@ -20,6 +23,8 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.opensearch.be.lucene.index.LuceneWriter.WRITER_GENERATION_ATTRIBUTE;
+
 /**
  * Lucene implementation of {@link EngineReaderManager}.
  * <p>
@@ -72,11 +77,60 @@ public void afterRefresh(boolean didRefresh, CatalogSnapshot catalogSnapshot) th
         }
         DirectoryReader refreshed = DirectoryReader.openIfChanged(currentReader);
         if (refreshed != null) {
+            // Guard against refresh/merge-apply races: a prior IT regression surfaced when
+            // overlapping threads produced a refreshed reader whose leaves disagreed with the
+            // catalog snapshot being registered, effectively pairing the snapshot with a stale
+            // reader. This assert catches that drift in test builds before the mismatched pair
+            // is published to readers.
+            assert readersAreSame(catalogSnapshot, refreshed);
             currentReader = refreshed;
         }
         readers.put(catalogSnapshot, currentReader);
     }
 
+    /**
+     * Consistency check: verifies that the refreshed {@link DirectoryReader} reflects exactly
+     * the set of segments the given {@link CatalogSnapshot} references. Compares the sorted
+     * list of writer generations drawn from the snapshot's {@link Segment Segments} against
+     * the sorted list of writer generations read off each leaf of the reader (via the
+     * {@link org.opensearch.be.lucene.index.LuceneWriter#WRITER_GENERATION_ATTRIBUTE} stamped
+     * onto every Lucene segment at write time).
+     *
+     * <p>Used only in an {@code assert} to catch refresh/catalog drift in test builds — if
+     * this ever returns {@code false} in production, it means a Lucene reader has been paired
+     * with the wrong catalog snapshot.
+     *
+     * @param catalogSnapshot catalog snapshot whose referenced generations are the expected set
+     * @param readers         DirectoryReader whose leaves' generations are the actual set
+     * @return {@code true} iff both lists contain the same generations in the same (sorted) order
+     */
+    private boolean readersAreSame(CatalogSnapshot catalogSnapshot, DirectoryReader readers) {
+        Collection<Long> generationsReferenced = catalogSnapshot.getSegments().stream().map(Segment::generation).sorted().toList();
+        return generationsReferenced.equals(collectReferencedGenerations(readers));
+    }
+
+    /**
+     * Extracts the writer generation from each leaf of the given {@link DirectoryReader} and
+     * returns them as a sorted list. Each leaf's {@link SegmentReader} carries a
+     * {@link SegmentCommitInfo} whose {@code SegmentInfo} is stamped with the
+     * {@link org.opensearch.be.lucene.index.LuceneWriter#WRITER_GENERATION_ATTRIBUTE} when the
+     * segment is written; parsing that attribute yields the generation that produced the leaf.
+     *
+     * @param reader the DirectoryReader to inspect
+     * @return generations of all leaves, sorted ascending
+     * @throws NumberFormatException if a leaf is missing the writer-generation attribute or
+     *                               its value is not parseable as a long (indicates a segment
+     *                               not produced by {@link org.opensearch.be.lucene.index.LuceneWriter})
+     * @throws ClassCastException    if any leaf reader is not a {@link SegmentReader}
+     */
+    private Collection<Long> collectReferencedGenerations(DirectoryReader reader) {
+        return reader.leaves().stream().map(lrc -> {
+            SegmentReader segmentReader = (SegmentReader) lrc.reader();
+            SegmentCommitInfo sci = segmentReader.getSegmentInfo();
+            return Long.parseLong(sci.info.getAttribute(WRITER_GENERATION_ATTRIBUTE));
+        }).sorted().toList();
+    }
+
     @Override
     public void onDeleted(CatalogSnapshot catalogSnapshot) throws IOException {
         DirectoryReader reader = readers.remove(catalogSnapshot);

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneCommitter.java

@@ -65,6 +65,19 @@
  * The store reference is incremented on construction and decremented on {@link #close()}.
  * Closing the committer also closes the underlying IndexWriter.
  *
+ * <h2>Refresh-lock coordination</h2>
+ *
+ * <p>The engine passes a {@code preMergeCommitHook} via {@link CommitterConfig}. We wire it
+ * into Lucene as a {@code MergedSegmentWarmer} on the {@link IndexWriterConfig}. The warmer
+ * runs between {@code mergeMiddle} and {@code commitMerge} while the {@link IndexWriter}
+ * monitor is <em>not</em> held, so invoking the hook there establishes the ordering
+ * {@code refreshLock → IW monitor} on the merge thread — matching the refresh path and
+ * avoiding the lock inversion that would occur if coordination happened inside
+ * {@code commitMerge}. Ownership of whatever the hook acquires (currently the engine's
+ * refresh lock) is transferred to the engine's {@code applyMergeChanges} callback, which
+ * releases it after the catalog is updated. This committer never touches the refresh lock
+ * directly.
+ *
  * @opensearch.experimental
  */
 @ExperimentalApi
@@ -73,7 +86,7 @@ public class LuceneCommitter extends SafeBootstrapCommitter {
     private static final Logger logger = LogManager.getLogger(LuceneCommitter.class);
 
     private final Store store;
-    private final IndexWriter indexWriter;
+    private final MergeIndexWriter indexWriter;
     private final LuceneCommitDeletionPolicy deletionPolicy;
     private final AtomicBoolean isClosed = new AtomicBoolean();
 
@@ -90,7 +103,7 @@ public LuceneCommitter(CommitterConfig committerConfig) throws IOException {
         this.store.incRef();
         try {
             this.deletionPolicy = new LuceneCommitDeletionPolicy();
-            IndexWriterConfig iwc = createIndexWriterConfig(committerConfig.engineConfig());
+            IndexWriterConfig iwc = createIndexWriterConfig(committerConfig);
             this.indexWriter = new MergeIndexWriter(store.directory(), iwc);
         } catch (Exception e) {
             store.decRef();
@@ -203,14 +216,15 @@ public boolean isCommitManagedFile(String fileName) {
      *
      * @return the index writer, or null if closed
      */
-    IndexWriter getIndexWriter() {
+    MergeIndexWriter getIndexWriter() {
         ensureOpen();
         return indexWriter;
     }
 
     // --- Internal ---
 
-    private IndexWriterConfig createIndexWriterConfig(EngineConfig engineConfig) {
+    private IndexWriterConfig createIndexWriterConfig(CommitterConfig committerConfig) {
+        EngineConfig engineConfig = committerConfig.engineConfig();
         if (engineConfig == null) {
             IndexWriterConfig iwc = new IndexWriterConfig();
             iwc.setIndexDeletionPolicy(deletionPolicy);
@@ -226,6 +240,14 @@ private IndexWriterConfig createIndexWriterConfig(EngineConfig engineConfig) {
         }
         iwc.setRAMBufferSizeMB(engineConfig.getIndexingBufferSize().getMbFrac());
         iwc.setUseCompoundFile(engineConfig.useCompoundFile());
+        // Refresh-lock hand-off: the MergedSegmentWarmer fires on the merge thread between
+        // mergeMiddle and commitMerge, while the IndexWriter monitor is NOT held. Invoking
+        // the engine-provided preMergeCommitHook here gives the merge path the ordering
+        // refreshLock → IW monitor, which matches the refresh path (DataFormatAwareEngine#refresh
+        // takes refreshLock before calling IndexWriter#addIndexes). Ownership of whatever the
+        // hook acquires is transferred to applyMergeChanges, which releases it after the
+        // catalog is updated. See the class Javadoc.
+        iwc.setMergedSegmentWarmer(_ -> committerConfig.preMergeCommitHook().run());
 
         // Determine if Lucene is a secondary format in a composite setup.
         // When secondary, use a SortedNumericSortField on the row ID so MultiSorter can reorder

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/index/LuceneIndexingExecutionEngine.java

@@ -15,6 +15,7 @@
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.MergeIndexWriter;
 import org.apache.lucene.index.NoMergePolicy;
 import org.apache.lucene.index.SegmentCommitInfo;
 import org.apache.lucene.index.SegmentReader;
@@ -74,7 +75,7 @@ public class LuceneIndexingExecutionEngine implements IndexingExecutionEngine<Lu
     private static final Logger logger = LogManager.getLogger(LuceneIndexingExecutionEngine.class);
 
     private final LuceneDataFormat dataFormat;
-    private final IndexWriter sharedWriter;
+    private final MergeIndexWriter sharedWriter;
     private final Store store;
     private final Path baseDirectory;
     private final Analyzer analyzer;
@@ -123,7 +124,7 @@ public LuceneIndexingExecutionEngine(
      *
      * @return the index writer
      */
-    public IndexWriter getWriter() {
+    public MergeIndexWriter getWriter() {
         return sharedWriter;
     }
 

sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/merge/LuceneMerger.java

@@ -80,14 +80,11 @@ private static Field initSegmentInfosField() {
     private final Path storeDirectory;
     private final LuceneMergeStrategy strategy;
 
-    public LuceneMerger(IndexWriter indexWriter, DataFormat dataFormat, Path storeDirectory) {
+    public LuceneMerger(MergeIndexWriter indexWriter, DataFormat dataFormat, Path storeDirectory) {
         if (indexWriter == null) {
             throw new IllegalArgumentException("IndexWriter must not be null");
         }
-        if (indexWriter instanceof MergeIndexWriter == false) {
-            throw new IllegalArgumentException("IndexWriter must be a MergeIndexWriter, got " + indexWriter.getClass().getName());
-        }
-        this.indexWriter = (MergeIndexWriter) indexWriter;
+        this.indexWriter = indexWriter;
         this.dataFormat = dataFormat;
         this.storeDirectory = storeDirectory;
         // TODO implement primary and integrate the same here
@@ -137,8 +134,21 @@ public MergeResult merge(MergeInput mergeInput) throws IOException {
         MergePolicy.OneMerge oneMerge = strategy.createOneMerge(matchingSegments, rowIdMapping);
         indexWriter.executeMerge(oneMerge, mergeInput.newWriterGeneration());
 
+        // Stamp the merged segment with its writer generation so downstream lookups
+        // (e.g. findMatchingSegments on a subsequent merge) can correlate it.
+        //
+        // This mutation is in-memory only: Lucene writes the .si file exactly once at
+        // segment creation via SegmentInfoFormat.write(...) and does not rewrite it on
+        // later commits, so this attribute will not survive a writer reopen. That is
+        // acceptable here because the attribute is only consumed within the lifetime
+        // of the live IndexWriter's SegmentInfos.
+        SegmentCommitInfo mergedInfo = oneMerge.getMergeInfo();
+        if (mergedInfo != null) {
+            mergedInfo.info.putAttribute(WRITER_GENERATION_ATTRIBUTE, String.valueOf(mergeInput.newWriterGeneration()));
+        }
+
         // Build the merged WriterFileSet from the output segment info
-        WriterFileSet mergedFileSet = buildMergedFileSet(oneMerge.getMergeInfo(), mergeInput.newWriterGeneration());
+        WriterFileSet mergedFileSet = buildMergedFileSet(mergedInfo, mergeInput.newWriterGeneration());
 
         // Delegate RowIdMapping production to the strategy
         RowIdMapping outputMapping = strategy.buildRowIdMapping(oneMerge, mergeInput);

sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneReaderManagerTests.java

@@ -15,13 +15,17 @@
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.MMapDirectory;
 import org.apache.lucene.store.NIOFSDirectory;
 import org.opensearch.be.lucene.index.LuceneCommitter;
 import org.opensearch.be.lucene.index.LuceneIndexingExecutionEngine;
+import org.opensearch.be.lucene.index.LuceneWriter;
+import org.opensearch.common.SuppressForbidden;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.index.IndexSettings;
@@ -99,6 +103,17 @@ private DirectoryReader openReader() throws IOException {
     }
 
     private CatalogSnapshot stubSnapshot(long generation) {
+        return stubSnapshot(generation, List.of());
+    }
+
+    /**
+     * Builds a stub snapshot whose segment list contains the given writer generations.
+     * This is required by {@link LuceneReaderManager#afterRefresh}'s assertion, which
+     * compares the snapshot's segment generations against the writer-generation attribute
+     * on each leaf in the refreshed {@link DirectoryReader}.
+     */
+    private CatalogSnapshot stubSnapshot(long generation, List<Long> segmentGenerations) {
+        List<Segment> segs = segmentGenerations.stream().map(g -> Segment.builder(g).build()).toList();
         return new CatalogSnapshot("test", generation, 1) {
             @Override
             protected void closeInternal() {}
@@ -115,7 +130,7 @@ public long getId() {
 
             @Override
             public List<Segment> getSegments() {
-                return List.of();
+                return segs;
             }
 
             @Override
@@ -168,11 +183,36 @@ public Collection<String> getFiles(boolean includeSegmentsFile) {
         };
     }
 
-    private void addDoc(String id) throws IOException {
+    private void addDoc(String id, long generation) throws IOException {
         Document doc = new Document();
         doc.add(new StringField("id", id, Field.Store.YES));
         indexWriter.addDocument(doc);
         indexWriter.commit();
+        stampLatestSegmentGeneration(generation);
+    }
+
+    /**
+     * Stamps the most recently written segment with the {@code writer_generation} attribute
+     * that {@link LuceneReaderManager#afterRefresh}'s assertion expects. In production this
+     * is done by {@code LuceneWriterCodec}; tests that write directly through a plain
+     * {@link IndexWriter} must stamp it themselves.
+     */
+    @SuppressForbidden(reason = "Need reflection to stamp writer_generation on segments for testing")
+    private void stampLatestSegmentGeneration(long generation) throws IOException {
+        try {
+            java.lang.reflect.Field segInfosField = IndexWriter.class.getDeclaredField("segmentInfos");
+            segInfosField.setAccessible(true);
+            SegmentInfos segInfos = (SegmentInfos) segInfosField.get(indexWriter);
+            if (segInfos.size() == 0) {
+                return;
+            }
+            SegmentCommitInfo last = segInfos.asList().get(segInfos.size() - 1);
+            if (last.info.getAttribute(LuceneWriter.WRITER_GENERATION_ATTRIBUTE) == null) {
+                last.info.putAttribute(LuceneWriter.WRITER_GENERATION_ATTRIBUTE, String.valueOf(generation));
+            }
+        } catch (ReflectiveOperationException e) {
+            throw new IOException("Failed to stamp writer_generation via reflection", e);
+        }
     }
 
     public void testAfterRefreshCreatesReader() throws IOException {
@@ -195,21 +235,24 @@ public void testAfterRefreshNoOpWhenDidRefreshFalse() throws IOException {
     public void testMultipleRefreshesWithIndexing() throws IOException {
         LuceneReaderManager rm = new LuceneReaderManager(dataFormat, openReader());
 
+        // Empty initial reader — no segments yet.
         CatalogSnapshot snap1 = stubSnapshot(1);
         rm.afterRefresh(true, snap1);
         DirectoryReader reader1 = rm.getReader(snap1);
         assertEquals(0, new IndexSearcher(reader1).count(new MatchAllDocsQuery()));
 
-        addDoc("doc1");
-        CatalogSnapshot snap2 = stubSnapshot(2);
+        // Add doc1 in generation 10, refresh. Reader now has one leaf stamped with gen=10.
+        addDoc("doc1", 10L);
+        CatalogSnapshot snap2 = stubSnapshot(2, List.of(10L));
         rm.afterRefresh(true, snap2);
         DirectoryReader reader2 = rm.getReader(snap2);
         assertEquals(1, new IndexSearcher(reader2).count(new MatchAllDocsQuery()));
 
         assertEquals(0, new IndexSearcher(reader1).count(new MatchAllDocsQuery()));
 
-        addDoc("doc2");
-        CatalogSnapshot snap3 = stubSnapshot(3);
+        // Add doc2 in generation 20. Reader now has two leaves stamped with gens {10, 20}.
+        addDoc("doc2", 20L);
+        CatalogSnapshot snap3 = stubSnapshot(3, List.of(10L, 20L));
         rm.afterRefresh(true, snap3);
         DirectoryReader reader3 = rm.getReader(snap3);
         assertEquals(2, new IndexSearcher(reader3).count(new MatchAllDocsQuery()));
@@ -286,7 +329,7 @@ public void testCreateReaderManagerWithLuceneIndexingEngine() throws IOException
             )
             .retentionLeasesSupplier(() -> new RetentionLeases(0, 0, java.util.Collections.emptyList()))
             .build();
-        CommitterConfig cs = new CommitterConfig(engineConfig);
+        CommitterConfig cs = new CommitterConfig(engineConfig, () -> {});
         LuceneCommitter committer = new LuceneCommitter(cs);
 
         try {