Merge branch 'main' into unified-allocator-followup

Author: gaurav-amz

modules/store-subdirectory/src/internalClusterTest/java/org/opensearch/plugin/store/subdirectory/SubdirectoryAwareRecoveryTests.java

@@ -30,6 +30,8 @@
 import org.opensearch.index.engine.EngineException;
 import org.opensearch.index.engine.EngineFactory;
 import org.opensearch.index.engine.InternalEngine;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.SegmentInfosCatalogSnapshot;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.indices.IndicesService;
 import org.opensearch.plugins.EnginePlugin;
@@ -306,6 +308,42 @@ public GatedCloseable<IndexCommit> acquireSafeIndexCommit() throws EngineExcepti
                 throw new EngineException(shardId, "Failed to acquire safe index commit", e);
             }
         }
+
+        @Override
+        public GatedCloseable<CatalogSnapshot> acquireSafeCatalogSnapshot() throws EngineException {
+            final GatedCloseable<CatalogSnapshot> base = super.acquireSafeCatalogSnapshot();
+            try {
+                // Collect subdirectory file paths (prefixed with subdirectory name).
+                final SegmentInfos subInfos = SegmentInfos.readLatestCommit(subdirectoryDirectory);
+                final Set<String> extra = new HashSet<>();
+                for (String f : subInfos.files(true)) {
+                    extra.add(Path.of(SUBDIRECTORY_NAME, f).toString());
+                }
+                for (String f : subdirectoryDirectory.listAll()) {
+                    if (f.startsWith(NON_SEGMENT_FILE_PREFIX)) {
+                        extra.add(Path.of(SUBDIRECTORY_NAME, f).toString());
+                    }
+                }
+                // Reuse the SegmentInfos from the base snapshot — no redundant disk read.
+                SegmentInfos mainInfos = ((SegmentInfosCatalogSnapshot) base.get()).getSegmentInfos();
+                CatalogSnapshot wrapped = new SegmentInfosCatalogSnapshot(mainInfos) {
+                    @Override
+                    public Collection<String> getFiles(boolean includeSegmentsFile) throws IOException {
+                        Set<String> files = new HashSet<>(super.getFiles(includeSegmentsFile));
+                        files.addAll(extra);
+                        return files;
+                    }
+                };
+                return new GatedCloseable<>(wrapped, base::close);
+            } catch (Exception e) {
+                try {
+                    base.close();
+                } catch (IOException ex) {
+                    e.addSuppressed(ex);
+                }
+                throw new EngineException(shardId, "Failed to acquire safe catalog snapshot", e);
+            }
+        }
     }
 
     /**

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionService.java

@@ -94,6 +94,7 @@ protected void doStop() {
     @Override
     protected void doClose() throws IOException {
         releaseRuntime();
+        NativeBridge.shutdownTokioRuntimeManager();
     }
 
     /**

sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReaderManager.java

@@ -37,7 +37,7 @@ public class DatafusionReaderManager implements EngineReaderManager<DatafusionRe
 
     private static final Logger logger = LogManager.getLogger(DatafusionReaderManager.class);
 
-    private final Map<CatalogSnapshot, DatafusionReader> readers = new HashMap<>();
+    private final Map<Long, DatafusionReader> readers = new HashMap<>();
     private final DataFormat dataFormat;
     private final String directoryPath;
     private final DataFusionService dataFusionService;
@@ -66,15 +66,19 @@ public DatafusionReaderManager(
 
     @Override
     public DatafusionReader getReader(CatalogSnapshot catalogSnapshot) throws IOException {
-        if (readers.containsKey(catalogSnapshot)) {
-            return readers.get(catalogSnapshot);
+        if (catalogSnapshot == null) {
+            throw new IllegalArgumentException("catalogSnapshot must not be null");
         }
-        throw new IOException("No DataFusion reader available");
+        DatafusionReader reader = readers.get(catalogSnapshot.getId());
+        if (reader == null) {
+            throw new IOException("No DataFusion reader available for catalog snapshot [version=" + catalogSnapshot.getId() + "]");
+        }
+        return reader;
     }
 
     @Override
     public void onDeleted(CatalogSnapshot catalogSnapshot) throws IOException {
-        DatafusionReader removed = readers.remove(catalogSnapshot);
+        DatafusionReader removed = readers.remove(catalogSnapshot.getId());
         if (removed != null) {
             removed.close();
         }
@@ -98,13 +102,13 @@ public void beforeRefresh() throws IOException {}
     @Override
     public void afterRefresh(boolean didRefresh, CatalogSnapshot catalogSnapshot) throws IOException {
         if (didRefresh == false) return;
-        if (readers.containsKey(catalogSnapshot)) return;
+        if (readers.containsKey(catalogSnapshot.getId())) return;
         DatafusionReader reader = new DatafusionReader(
             directoryPath,
             catalogSnapshot.getSearchableFiles(dataFormat.name()),
             dataformatAwareStoreHandle
         );
-        readers.put(catalogSnapshot, reader);
+        readers.put(catalogSnapshot.getId(), reader);
     }
 
     private Collection<String> toAbsolutePaths(Collection<String> fileNames) {

sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReaderManagerTests.java

@@ -219,7 +219,7 @@ public void testGetReaderWithNoRefreshThrows() throws IOException {
         DataFusionService mockService = mock(DataFusionService.class);
 
         DatafusionReaderManager manager = new DatafusionReaderManager(TEST_FORMAT, shardPath, mockService, null);
-        expectThrows(IOException.class, () -> manager.getReader(null));
+        expectThrows(IllegalArgumentException.class, () -> manager.getReader(null));
         manager.close();
     }
 }

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

@@ -15,21 +15,26 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.IndexSettings;
 import org.opensearch.index.engine.dataformat.DataFormat;
+import org.opensearch.index.engine.dataformat.DataFormatDescriptor;
 import org.opensearch.index.engine.dataformat.DataFormatPlugin;
+import org.opensearch.index.engine.dataformat.DataFormatRegistry;
 import org.opensearch.index.engine.dataformat.DeleteExecutionEngine;
 import org.opensearch.index.engine.dataformat.IndexingEngineConfig;
 import org.opensearch.index.engine.dataformat.IndexingExecutionEngine;
 import org.opensearch.index.engine.dataformat.ReaderManagerConfig;
 import org.opensearch.index.engine.exec.EngineReaderManager;
 import org.opensearch.index.engine.exec.commit.Committer;
 import org.opensearch.index.engine.exec.commit.CommitterFactory;
+import org.opensearch.index.store.checksum.LuceneChecksumHandler;
 import org.opensearch.plugins.EnginePlugin;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.plugins.SearchBackEndPlugin;
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
+import java.util.function.Supplier;
 
 /**
  * Plugin providing Lucene as a data format, search back-end, and committer
@@ -86,6 +91,14 @@ public DataFormat getDataFormat() {
         );
     }
 
+    @Override
+    public Map<String, Supplier<DataFormatDescriptor>> getFormatDescriptors(
+        IndexSettings indexSettings,
+        DataFormatRegistry dataFormatRegistry
+    ) {
+        return Map.of(DATA_FORMAT.name(), () -> new DataFormatDescriptor(DATA_FORMAT.name(), new LuceneChecksumHandler()));
+    }
+
     // --- SearchBackEndPlugin ---
 
     /** {@inheritDoc} Returns {@code "lucene"}. */

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

@@ -11,7 +11,11 @@
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SegmentCommitInfo;
+import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.index.SegmentReader;
+import org.opensearch.be.lucene.index.LuceneReplicaCommitter;
+import org.opensearch.common.CheckedBiFunction;
+import org.opensearch.common.SuppressForbidden;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.DataFormat;
 import org.opensearch.index.engine.exec.EngineReaderManager;
@@ -44,30 +48,42 @@
  * @opensearch.experimental
  */
 @ExperimentalApi
+@SuppressForbidden(reason = "reference counting is required here")
 public class LuceneReaderManager implements EngineReaderManager<LuceneReader> {
 
     private final DataFormat dataFormat;
-    private final Map<CatalogSnapshot, LuceneReader> readers = new HashMap<>();
+    private final Map<Long, LuceneReader> readers;
     private volatile DirectoryReader currentReader;
+    private final CheckedBiFunction<DirectoryReader, SegmentInfos, DirectoryReader, IOException> readerRefresher;
 
     /**
      * Creates a new LuceneReaderManager.
      *
-     * @param dataFormat the data format this reader manager serves
-     * @param initialReader the initial DirectoryReader, must not be null
+     * @param dataFormat      the data format this reader manager serves
+     * @param initialReader   the initial DirectoryReader, must not be null
+     * @param readers         shared map of generation to DirectoryReader for segment-level reader reuse
+     * @param readerRefresher function that opens a refreshed reader given the current reader and new
+     *                        {@link SegmentInfos}; returns {@code null} if no refresh is needed
      * @throws NullPointerException if initialReader is null
      */
-    public LuceneReaderManager(DataFormat dataFormat, DirectoryReader initialReader) {
+    public LuceneReaderManager(
+        DataFormat dataFormat,
+        DirectoryReader initialReader,
+        Map<Long, LuceneReader> readers,
+        CheckedBiFunction<DirectoryReader, SegmentInfos, DirectoryReader, IOException> readerRefresher
+    ) {
         this.dataFormat = dataFormat;
         Objects.requireNonNull(initialReader, "initialReader must not be null");
         this.currentReader = initialReader;
+        this.readers = readers;
+        this.readerRefresher = readerRefresher;
     }
 
     @Override
     public LuceneReader getReader(CatalogSnapshot catalogSnapshot) throws IOException {
-        LuceneReader reader = readers.get(catalogSnapshot);
+        LuceneReader reader = readers.get(catalogSnapshot.getId());
         if (reader == null) {
-            throw new IllegalStateException("No reader available for catalog snapshot [gen=" + catalogSnapshot.getGeneration() + "]");
+            throw new IllegalStateException("No reader available for catalog snapshot [version=" + catalogSnapshot.getId() + "]");
         }
         return reader;
     }
@@ -79,17 +95,25 @@ public void beforeRefresh() throws IOException {
 
     @Override
     public void afterRefresh(boolean didRefresh, CatalogSnapshot catalogSnapshot) throws IOException {
-        if (didRefresh == false || readers.containsKey(catalogSnapshot)) {
+        if (didRefresh == false || readers.containsKey(catalogSnapshot.getId())) {
             return;
         }
-        DirectoryReader refreshed = DirectoryReader.openIfChanged(currentReader);
+        DirectoryReader refreshed = readerRefresher.apply(currentReader, LuceneReplicaCommitter.getSegmentInfos(catalogSnapshot));
         if (refreshed != null) {
-            assert readersAreSame(catalogSnapshot, refreshed);
+            // 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.
             currentReader = refreshed;
+        } else {
+            // If same reader is used, assert that calalog snapshot is same.
+            currentReader.incRef();
         }
+        assert readersAreSame(catalogSnapshot, currentReader);
 
         Map<Long, String> generationToSegmentName = buildGenerationToSegmentName(catalogSnapshot, currentReader.leaves());
-        readers.put(catalogSnapshot, new LuceneReader(currentReader, generationToSegmentName));
+        readers.put(catalogSnapshot.getId(), new LuceneReader(currentReader, generationToSegmentName));
     }
 
     private static Map<Long, String> buildGenerationToSegmentName(CatalogSnapshot catalogSnapshot, List<LeafReaderContext> leaves) {
@@ -135,12 +159,12 @@ private static Map<Long, String> buildGenerationToSegmentName(CatalogSnapshot ca
      * 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
+     * @param reader         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) {
+    private boolean readersAreSame(CatalogSnapshot catalogSnapshot, DirectoryReader reader) {
         Collection<Long> generationsReferenced = catalogSnapshot.getSegments().stream().map(Segment::generation).sorted().toList();
-        return generationsReferenced.equals(collectReferencedGenerations(readers));
+        return generationsReferenced.equals(collectReferencedGenerations(reader));
     }
 
     /**
@@ -167,9 +191,9 @@ private Collection<Long> collectReferencedGenerations(DirectoryReader reader) {
 
     @Override
     public void onDeleted(CatalogSnapshot catalogSnapshot) throws IOException {
-        LuceneReader reader = readers.remove(catalogSnapshot);
+        LuceneReader reader = readers.remove(catalogSnapshot.getId());
         if (reader != null) {
-            reader.directoryReader().close();
+            reader.directoryReader().decRef();
         }
     }
 
@@ -186,7 +210,7 @@ public void onFilesAdded(Collection<String> files) throws IOException {
     @Override
     public void close() throws IOException {
         for (LuceneReader reader : readers.values()) {
-            reader.directoryReader().close();
+            reader.directoryReader().decRef();
         }
         readers.clear();
     }

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

@@ -11,14 +11,22 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.index.StandardDirectoryReader;
 import org.opensearch.be.lucene.index.LuceneIndexingExecutionEngine;
+import org.opensearch.common.CheckedBiFunction;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.ReaderManagerConfig;
 import org.opensearch.index.engine.exec.EngineReaderManager;
 import org.opensearch.index.engine.exec.commit.IndexStoreProvider;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Static helpers for creating Lucene-based {@link EngineReaderManager} instances.
@@ -51,12 +59,30 @@ static EngineReaderManager<LuceneReader> createReaderManager(ReaderManagerConfig
         IndexStoreProvider provider = settings.indexStoreProvider()
             .orElseThrow(() -> new IllegalStateException("IndexStoreProvider is required to create LuceneReaderManager"));
         DirectoryReader directoryReader;
+        Map<Long, LuceneReader> readers = new ConcurrentHashMap<>();
+        CheckedBiFunction<DirectoryReader, SegmentInfos, DirectoryReader, IOException> readerRefresher = null;
         if (provider.getStore(settings.format()) instanceof LuceneIndexingExecutionEngine.LuceneFormatStore luceneProvider) {
             directoryReader = DirectoryReader.open(luceneProvider.writer());
+            readers = luceneProvider.readers();
+            readerRefresher = (dr, sis) -> DirectoryReader.openIfChanged(dr);
         } else {
             logger.warn("Initialising it with a DirectorReader instead of a writer");
             directoryReader = StandardDirectoryReader.open(provider.getStore(settings.format()).store().directory());
+            readerRefresher = LuceneSearchBackEnd::buildReader;
         }
-        return new LuceneReaderManager(settings.format(), directoryReader);
+        return new LuceneReaderManager(settings.format(), directoryReader, readers, readerRefresher);
+    }
+
+    private static DirectoryReader buildReader(DirectoryReader oldReader, SegmentInfos newSis) throws IOException {
+        if (newSis == null || ((StandardDirectoryReader) oldReader).getSegmentInfos().version == newSis.version) {
+            return null;
+        }
+        final List<LeafReader> subs = new ArrayList<>();
+        for (LeafReaderContext ctx : oldReader.leaves()) {
+            subs.add(ctx.reader());
+        }
+        // Segment_n here is ignored because it is either already committed on disk as part of previous commit point or
+        // does not yet exist on store (not yet committed)
+        return StandardDirectoryReader.open(oldReader.directory(), newSis, subs, null, null);
     }
 }

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

@@ -49,6 +49,7 @@ public void onInit(List<? extends IndexCommit> commits) throws IOException {
         // which will delete it once a CS commit exists.
         for (IndexCommit commit : commits) {
             if (commit.getUserData().get(CatalogSnapshot.CATALOG_SNAPSHOT_ID) == null) {
+                assert nonCatalogSnapshotCommit == null;
                 nonCatalogSnapshotCommit = commit;
             }
         }
@@ -74,6 +75,7 @@ public void onCommit(List<? extends IndexCommit> commits) throws IOException {
         // since it is no longer needed for recovery.
         if (hasCSCommit && nonCatalogSnapshotCommit != null) {
             nonCatalogSnapshotCommit.delete();
+            pendingDeletes.remove(0L);
             nonCatalogSnapshotCommit = null;
         }
     }
@@ -86,7 +88,7 @@ public void onCommit(List<? extends IndexCommit> commits) throws IOException {
      * @param snapshotId the CatalogSnapshot ID to purge
      */
     void purgeCommit(long snapshotId) {
-        assert trackedCommits.containsKey(snapshotId);
+        assert (snapshotId == 0L) || trackedCommits.containsKey(snapshotId);
         pendingDeletes.add(snapshotId);
     }
 }