Add catalogSnapshotManager for dataformat aware engine (#20982)

* Initial commit for CatalogSnapshot

Signed-off-by: Arpit Bandejiya <abandeji@amazon.com>

* Address comments

Signed-off-by: Arpit Bandejiya <abandeji@amazon.com>

* Address comment part2

Signed-off-by: Arpit Bandejiya <abandeji@amazon.com>

---------

Signed-off-by: Arpit Bandejiya <abandeji@amazon.com>
Co-authored-by: Mohit Godwani <81609427+mgodwan@users.noreply.github.com>

Author: alchemist51

.gitignore

@@ -70,3 +70,4 @@ testfixtures_shared/
 # build files generated
 doc-tools/missing-doclet/bin/
 /sandbox/plugins/engine-datafusion/target/
+**/Cargo.lock

CHANGELOG.md

@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add support for enabling pluggable data formats, starting with phase-1 of decoupling shard from engine, and introducing basic abstractions ([#20675](https://github.com/opensearch-project/OpenSearch/pull/20675))
 - Add concurrent queue in libs and composite engine sandbox plugin ([#20909](https://github.com/opensearch-project/OpenSearch/pull/20909))
 - Add interface for the Multi format merge flow ([#20908](https://github.com/opensearch-project/OpenSearch/pull/20908))
+- Add CatalogSnapshotManager lifecycle management with reference-counted snapshot tracking and serialization support for Segment and WriterFileSet ([#20982](https://github.com/opensearch-project/OpenSearch/pull/20982))
 
 - Add warmup phase to wait for lag to catch up in pull-based ingestion before serving ([#20526](https://github.com/opensearch-project/OpenSearch/pull/20526))
 - Add a new static method to IndicesOptions API to expose `STRICT_EXPAND_OPEN_HIDDEN_FORBID_CLOSED` index option ([#20980](https://github.com/opensearch-project/OpenSearch/pull/20980))

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

@@ -10,8 +10,8 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.DataFormat;
-import org.opensearch.index.engine.exec.CatalogSnapshot;
 import org.opensearch.index.engine.exec.EngineReaderManager;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 import org.opensearch.index.shard.ShardPath;
 
 import java.io.IOException;

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

@@ -12,8 +12,8 @@
 import org.apache.lucene.search.ReferenceManager;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.DataFormat;
-import org.opensearch.index.engine.exec.CatalogSnapshot;
 import org.opensearch.index.engine.exec.EngineReaderManager;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 
 import java.io.IOException;
 import java.util.Collection;

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/DefaultPlanExecutorTests.java

@@ -30,15 +30,17 @@
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.concurrent.GatedCloseable;
 import org.opensearch.core.index.Index;
 import org.opensearch.index.IndexService;
 import org.opensearch.index.engine.DataFormatAwareEngine;
 import org.opensearch.index.engine.dataformat.DataFormat;
 import org.opensearch.index.engine.dataformat.FieldTypeCapabilities;
-import org.opensearch.index.engine.exec.CatalogSnapshot;
 import org.opensearch.index.engine.exec.EngineReaderManager;
 import org.opensearch.index.engine.exec.Segment;
 import org.opensearch.index.engine.exec.WriterFileSet;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.indices.IndicesService;
 import org.opensearch.test.OpenSearchTestCase;
@@ -104,13 +106,15 @@ public void testEndToEndExecuteWithMockBackend() throws IOException {
 
         Segment seg1 = Segment.builder(0L).addSearchableFiles(format, fs1).build();
         Segment seg2 = Segment.builder(1L).addSearchableFiles(format, fs2).build();
-        MockCatalogSnapshot snapshot = new MockCatalogSnapshot(1L, List.of(seg1, seg2), format);
+
+        CatalogSnapshotManager snapshotManager = new CatalogSnapshotManager(1L, 1L, 0L, List.of(seg1, seg2), 2L, Map.of());
 
         MockReaderManager readerManager = new MockReaderManager(format.name());
-        readerManager.afterRefresh(true, snapshot);
+        try (GatedCloseable<CatalogSnapshot> ref = snapshotManager.acquireSnapshot()) {
+            readerManager.afterRefresh(true, ref.get());
+        }
 
-        DataFormatAwareEngine engine = new DataFormatAwareEngine(Map.of(format, readerManager));
-        engine.setLatestSnapshot(snapshot);
+        DataFormatAwareEngine engine = new DataFormatAwareEngine(Map.of(format, readerManager), snapshotManager);
 
         // Mock shard + cluster wiring
         IndexShard shard = mock(IndexShard.class);
@@ -288,16 +292,18 @@ public String serializeToString() {
         }
 
         @Override
-        public void setCatalogSnapshotMap(Map<Long, ? extends CatalogSnapshot> map) {}
-
-        @Override
-        public void setUserData(Map<String, String> userData, boolean b) {}
+        public void setUserData(Map<String, String> userData) {}
 
         @Override
         public Object getReader(DataFormat dataFormat) {
             return null;
         }
 
+        @Override
+        public MockCatalogSnapshot clone() {
+            return new MockCatalogSnapshot(generation, segments, format);
+        }
+
         @Override
         protected void closeInternal() {}
     }

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

@@ -11,10 +11,11 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.concurrent.GatedCloseable;
 import org.opensearch.index.engine.dataformat.DataFormat;
-import org.opensearch.index.engine.exec.CatalogSnapshot;
 import org.opensearch.index.engine.exec.DataFormatAwareEngineFactory;
 import org.opensearch.index.engine.exec.EngineReaderManager;
 import org.opensearch.index.engine.exec.IndexReaderProvider;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshotManager;
 
 import java.io.Closeable;
 import java.io.IOException;
@@ -35,62 +36,56 @@
 public class DataFormatAwareEngine implements IndexReaderProvider, Closeable {
 
     private final Map<DataFormat, EngineReaderManager<?>> readerManagers;
-    private volatile CatalogSnapshot latestSnapshot;
+    private volatile CatalogSnapshotManager catalogSnapshotManager;
 
     /**
-     * Constructs a new DataFormatAwareEngine with pre-built maps.
+     * Constructs a new DataFormatAwareEngine.
      * Prefer using {@link DataFormatAwareEngineFactory#create()}.
      */
+    public DataFormatAwareEngine(Map<DataFormat, EngineReaderManager<?>> readerManagers, CatalogSnapshotManager catalogSnapshotManager) {
+        this.readerManagers = readerManagers;
+        this.catalogSnapshotManager = catalogSnapshotManager;
+    }
+
+    /**
+     * Constructs a new DataFormatAwareEngine without a snapshot manager.
+     * The manager must be set via {@link #setCatalogSnapshotManager} before acquiring readers.
+     */
     public DataFormatAwareEngine(Map<DataFormat, EngineReaderManager<?>> readerManagers) {
         this.readerManagers = readerManagers;
     }
 
-    public EngineReaderManager<?> getReaderManager(DataFormat format) {
-        return readerManagers.get(format);
+    public void setCatalogSnapshotManager(CatalogSnapshotManager catalogSnapshotManager) {
+        this.catalogSnapshotManager = catalogSnapshotManager;
     }
 
-    /**
-     * Called by the catalog snapshot lifecycle listener after a refresh
-     * to update the latest searchable snapshot.
-     */
-    public void setLatestSnapshot(CatalogSnapshot snapshot) {
-        CatalogSnapshot prev = this.latestSnapshot;
-        this.latestSnapshot = snapshot;
-        if (prev != null) {
-            prev.decRef();
-        }
+    public EngineReaderManager<?> getReaderManager(DataFormat format) {
+        return readerManagers.get(format);
     }
 
     /**
      * Acquires a DataFormatAwareReader on the latest catalog snapshot.
-     * The snapshot is incRef'd; the caller MUST close the returned
-     * {@link DataFormatAwareReader} when done, which decRef's the snapshot.
+     * The caller MUST close the returned {@link DataFormatAwareReader} when done,
+     * which releases the snapshot reference.
      */
     public GatedCloseable<Reader> acquireReader() throws IOException {
-        CatalogSnapshot snapshot = latestSnapshot;
-        if (snapshot == null) {
-            throw new IllegalStateException("No catalog snapshot available");
+        if (catalogSnapshotManager == null) {
+            throw new IllegalStateException("CatalogSnapshotManager not set");
         }
-        return acquireReader(snapshot);
-    }
-
-    /**
-     * Acquires a dataFormatAwareReader on a specific catalog snapshot.
-     */
-    private GatedCloseable<Reader> acquireReader(CatalogSnapshot catalogSnapshot) throws IOException {
-        catalogSnapshot.incRef();
+        GatedCloseable<CatalogSnapshot> snapshotRef = catalogSnapshotManager.acquireSnapshot();
         try {
+            CatalogSnapshot catalogSnapshot = snapshotRef.get();
             Map<DataFormat, Object> readers = new HashMap<>();
             for (Map.Entry<DataFormat, EngineReaderManager<?>> entry : readerManagers.entrySet()) {
                 Object reader = entry.getValue().getReader(catalogSnapshot);
                 if (reader != null) {
                     readers.put(entry.getKey(), reader);
                 }
             }
-            DataFormatAwareReader reader = new DataFormatAwareReader(catalogSnapshot, readers);
+            DataFormatAwareReader reader = new DataFormatAwareReader(catalogSnapshot, snapshotRef, readers);
             return new GatedCloseable<>(reader, reader::close);
         } catch (Exception e) {
-            catalogSnapshot.decRef();
+            snapshotRef.close();
             throw e;
         }
     }
@@ -102,10 +97,16 @@ private GatedCloseable<Reader> acquireReader(CatalogSnapshot catalogSnapshot) th
     @ExperimentalApi
     public static class DataFormatAwareReader implements IndexReaderProvider.Reader {
         private final CatalogSnapshot catalogSnapshot;
+        private final GatedCloseable<CatalogSnapshot> snapshotRef;
         private final Map<DataFormat, Object> readers;
 
-        DataFormatAwareReader(CatalogSnapshot catalogSnapshot, Map<DataFormat, Object> readers) {
+        DataFormatAwareReader(
+            CatalogSnapshot catalogSnapshot,
+            GatedCloseable<CatalogSnapshot> snapshotRef,
+            Map<DataFormat, Object> readers
+        ) {
             this.catalogSnapshot = catalogSnapshot;
+            this.snapshotRef = snapshotRef;
             this.readers = readers;
         }
 
@@ -121,7 +122,11 @@ public CatalogSnapshot catalogSnapshot() {
 
         @Override
         public void close() {
-            catalogSnapshot.decRef();
+            try {
+                snapshotRef.close();
+            } catch (IOException e) {
+                throw new RuntimeException("Failed to release catalog snapshot reference", e);
+            }
         }
     }
 

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

@@ -14,8 +14,8 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.index.VersionType;
-import org.opensearch.index.engine.exec.CatalogSnapshot;
 import org.opensearch.index.engine.exec.Indexer;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 import org.opensearch.index.mapper.DocumentMapperForType;
 import org.opensearch.index.mapper.SourceToParse;
 import org.opensearch.index.merge.MergeStats;

server/src/main/java/org/opensearch/index/engine/dataformat/merge/MergeHandler.java

@@ -15,9 +15,9 @@
 import org.opensearch.common.logging.Loggers;
 import org.opensearch.core.index.shard.ShardId;
 import org.opensearch.index.engine.dataformat.MergeResult;
-import org.opensearch.index.engine.exec.CatalogSnapshot;
 import org.opensearch.index.engine.exec.Indexer;
 import org.opensearch.index.engine.exec.Segment;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 
 import java.util.ArrayDeque;
 import java.util.Collection;

server/src/main/java/org/opensearch/index/engine/exec/CatalogSnapshotLifecycleListener.java

@@ -9,6 +9,7 @@
 package org.opensearch.index.engine.exec;
 
 import org.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 
 import java.io.IOException;
 

server/src/main/java/org/opensearch/index/engine/exec/DataFormatEngineCatalogSnapshotListener.java

@@ -10,6 +10,7 @@
 
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.index.engine.dataformat.DataFormat;
+import org.opensearch.index.engine.exec.coord.CatalogSnapshot;
 
 import java.io.IOException;
 import java.util.Collection;