Fixing test failures

Signed-off-by: Kamal Nayan <askkamal@amazon.com>

Author: Kamal Nayan

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

@@ -53,6 +53,7 @@ public class LuceneReaderManager implements EngineReaderManager<LuceneReader> {
 
     private final DataFormat dataFormat;
     private final Map<Long, LuceneReader> readers;
+    private volatile DirectoryReader initialReader;
     private volatile DirectoryReader currentReader;
     private final CheckedBiFunction<DirectoryReader, SegmentInfos, DirectoryReader, IOException> readerRefresher;
 
@@ -74,6 +75,7 @@ public LuceneReaderManager(
     ) {
         this.dataFormat = dataFormat;
         Objects.requireNonNull(initialReader, "initialReader must not be null");
+        this.initialReader = initialReader;
         this.currentReader = initialReader;
         this.readers = readers;
         this.readerRefresher = readerRefresher;
@@ -195,6 +197,7 @@ public void onDeleted(CatalogSnapshot catalogSnapshot) throws IOException {
         if (reader != null) {
             reader.directoryReader().decRef();
         }
+        releaseInitialReader();
     }
 
     @Override
@@ -213,5 +216,24 @@ public void close() throws IOException {
             reader.directoryReader().decRef();
         }
         readers.clear();
+        releaseInitialReader();
+    }
+
+    /**
+     * Releases the initial {@link DirectoryReader} opened by
+     * {@link LuceneSearchBackEnd#createReaderManager} and nulls out the field. The check
+     * makes the operation idempotent so it is safe to invoke from both {@link #close()}
+     * and {@link #onDeleted(CatalogSnapshot)} without double-freeing.
+     * <p>
+     * The initial reader is owned by the manager (not the {@link #readers} map), so it is
+     * not covered by the per-snapshot decRef loop in {@code close()}. Without this release,
+     * the {@link java.nio.channels.FileChannel}s held by its {@link SegmentReader}s would
+     * leak until JVM exit (detected by Lucene's {@code LeakFS} at suite teardown).
+     */
+    private void releaseInitialReader() throws IOException {
+        if (initialReader != null) {
+            initialReader.decRef();
+            initialReader = null;
+        }
     }
 }

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

@@ -512,15 +512,16 @@ public void testMultipleSnapshotsShareReaderWhenRefresherReturnsNull() throws IO
         // RefCount should be initial + 3 (one incRef per afterRefresh)
         assertEquals(initialRefCount + 3, initialReader.getRefCount());
 
-        // Deleting each snapshot should decRef once
+        // First onDeleted: -1 for map entry, -1 for releaseInitialReader (idempotent thereafter).
         rm.onDeleted(snap1);
-        assertEquals(initialRefCount + 2, initialReader.getRefCount());
+        assertEquals(initialRefCount + 1, initialReader.getRefCount());
 
+        // Subsequent onDeleted: -1 for map entry only.
         rm.onDeleted(snap2);
-        assertEquals(initialRefCount + 1, initialReader.getRefCount());
+        assertEquals(initialRefCount, initialReader.getRefCount());
 
         rm.onDeleted(snap3);
-        assertEquals(initialRefCount, initialReader.getRefCount());
+        assertEquals(initialRefCount - 1, initialReader.getRefCount());
     }
 
     public void testCloseDecRefsAllAccumulatedReaders() throws IOException {
@@ -540,9 +541,9 @@ public void testCloseDecRefsAllAccumulatedReaders() throws IOException {
         rm.afterRefresh(true, stubSnapshot(3));
         assertEquals(initialRefCount + 3, initialReader.getRefCount());
 
-        // close() should decRef all 3
+        // close() decRefs the 3 map entries and the initial open(writer) reference: refCount → 0.
         rm.close();
-        assertEquals(initialRefCount, initialReader.getRefCount());
+        assertEquals(initialRefCount - 1, initialReader.getRefCount());
     }
 
     public void testMixedRefreshSomeNullSomeNew() throws IOException {

server/src/main/java/org/opensearch/index/engine/Engine.java

@@ -1292,6 +1292,24 @@ public GatedCloseable<CatalogSnapshot> acquireSafeCatalogSnapshot() throws Engin
         }
     }
 
+    /**
+     * Acquires a {@link CatalogSnapshot} of the engine's current in-memory state — i.e. the
+     * latest visible-to-readers segments, including any uncommitted operations that have been
+     * refreshed. Used by replication-checkpoint listeners and any caller that wants the live
+     * search-visible state rather than the on-disk safe commit.
+     *
+     * <p>The default implementation bridges via {@link #getSegmentInfosSnapshot()} and wraps the
+     * resulting {@link SegmentInfos} as a {@link SegmentInfosCatalogSnapshot}. Engines whose
+     * native state isn't a single {@link SegmentInfos} (e.g. multi-format engines that aggregate
+     * Lucene + other formats) should override this to return their native catalog snapshot.
+     */
+    @ExperimentalApi
+    public GatedCloseable<CatalogSnapshot> acquireSnapshot() {
+        final GatedCloseable<SegmentInfos> segmentInfosRef = getSegmentInfosSnapshot();
+        final CatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfosRef.get());
+        return new GatedCloseable<>(snapshot, segmentInfosRef::close);
+    }
+
     /**
      * Acquires a {@link CatalogSnapshot} pinned to the most recent commit on disk,
      * regardless of retention policy. Default wraps {@link #acquireLastIndexCommit(boolean)}.

server/src/main/java/org/opensearch/index/engine/EngineBackedIndexer.java

@@ -9,7 +9,6 @@
 package org.opensearch.index.engine;
 
 import org.apache.lucene.index.IndexCommit;
-import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.store.ByteBuffersDataOutput;
 import org.apache.lucene.store.ByteBuffersIndexOutput;
 import org.opensearch.common.annotation.ExperimentalApi;
@@ -418,15 +417,16 @@ public void finalizeReplication(CatalogSnapshot catalogSnapshot) throws IOExcept
     }
 
     /**
-     * Returns a snapshot of the catalog of segments in this engine. This snapshot is
-     * guaranteed to be consistent and can be used for recovery purposes.
+     * Returns a snapshot of the catalog of segments in this engine. Delegates to
+     * {@link Engine#acquireSnapshot()} so subclasses (e.g. the read-only wrapper used during
+     * engine reset) can route to a different snapshot source without going through the
+     * {@link Engine#getSegmentInfosSnapshot()} bridge — which is required when the underlying
+     * source is a non-Lucene indexer (e.g. {@link DataFormatAwareEngine}).
      */
     @ExperimentalApi
     @Override
     public GatedCloseable<CatalogSnapshot> acquireSnapshot() {
-        GatedCloseable<SegmentInfos> segmentInfosRef = engine.getSegmentInfosSnapshot();
-        SegmentInfosCatalogSnapshot snapshot = new SegmentInfosCatalogSnapshot(segmentInfosRef.get());
-        return new GatedCloseable<>(snapshot, segmentInfosRef::close);
+        return engine.acquireSnapshot();
     }
 
     @Override

server/src/main/java/org/opensearch/index/shard/IndexShard.java

@@ -5581,11 +5581,12 @@ public GatedCloseable<CatalogSnapshot> acquireLastCommittedSnapshot(boolean flus
                 }
 
                 @Override
-                public GatedCloseable<SegmentInfos> getSegmentInfosSnapshot() {
-                    if (newEngineReference.get() == null) {
+                public GatedCloseable<CatalogSnapshot> acquireSnapshot() {
+                    final Indexer ref = newEngineReference.get();
+                    if (ref == null) {
                         throw new AlreadyClosedException("engine was closed");
                     }
-                    return applyOnEngine(newEngineReference.get(), Engine::getSegmentInfosSnapshot);
+                    return ref.acquireSnapshot();
                 }
 
                 @Override

server/src/test/java/org/opensearch/index/shard/IndexShardTests.java

@@ -4861,7 +4861,7 @@ public void onBeginTranslogRecovery() {
      * {@code newEngineReference} is still null -- the window between ReadOnlyEngine installation
      * and {@code newEngineReference.set(newEngine)} inside {@code resetEngineToGlobalCheckpoint}.
      * Covers the defensive null-check branches in {@code acquireLastIndexCommit},
-     * {@code acquireSafeIndexCommit}, and {@code getSegmentInfosSnapshot}.
+     * {@code acquireSafeIndexCommit}, and {@code acquireSnapshot}.
      */
     public void testDelegateThrowsAlreadyClosedBeforeNewEngineSet() throws Exception {
         CountDownLatch creatingEngineLatch = new CountDownLatch(1);
@@ -4906,7 +4906,7 @@ public void testDelegateThrowsAlreadyClosedBeforeNewEngineSet() throws Exception
         // The ReadOnlyEngine is now the current engine, but newEngineReference is still null.
         expectThrows(AlreadyClosedException.class, () -> shard.acquireLastIndexCommit(false));
         expectThrows(AlreadyClosedException.class, shard::acquireSafeIndexCommit);
-        expectThrows(AlreadyClosedException.class, shard::getSegmentInfosSnapshot);
+        expectThrows(AlreadyClosedException.class, shard::getCatalogSnapshot);
 
         proceedWithCreationLatch.countDown();
         assertTrue("engine reset should complete", engineResetLatch.await(30, TimeUnit.SECONDS));