[Segment Replication] Update getSegmentInfosSnapshot() logic to return SegmentInfos with highest generation number

CHANGELOG.md

@@ -17,6 +17,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - PR reference to checkout code for changelog verifier ([#4296](https://github.com/opensearch-project/OpenSearch/pull/4296))
 - Restore using the class ClusterInfoRequest and ClusterInfoRequestBuilder from package 'org.opensearch.action.support.master.info' for subclasses ([#4307](https://github.com/opensearch-project/OpenSearch/pull/4307))
 - Do not fail replica shard due to primary closure ([#4133](https://github.com/opensearch-project/OpenSearch/pull/4133))
+- [Segment Replication] Update getSegmentInfosSnapshot() logic to return SegmentInfos with the highest generation number
 
 ### Security
 

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

@@ -2131,19 +2131,30 @@ protected SegmentInfos getLatestSegmentInfos() {
     }
 
     /**
-     * Fetch the latest {@link SegmentInfos} object via {@link #getLatestSegmentInfos()}
-     * but also increment the ref-count to ensure that these segment files are retained
-     * until the reference is closed. On close, the ref-count is decremented.
+     * Fetch the latest {@link SegmentInfos} object. The method uses reads in-memory SegmentInfos via {@link #getLatestSegmentInfos()}
+     * and on-disk SegmentInfos and returns instance which has the highest segment generation. This method also
+     * increment the ref-count to ensure that these segment files are retained until the reference is closed. On close,
+     * the ref-count is decremented.
      */
     @Override
     public GatedCloseable<SegmentInfos> getSegmentInfosSnapshot() {
-        final SegmentInfos segmentInfos = getLatestSegmentInfos();
+        final SegmentInfos inMemorySegInfos = getLatestSegmentInfos();
+        SegmentInfos result;
         try {
-            indexWriter.incRefDeleter(segmentInfos);
+            final long lastCommitGeneration = store.getLastCommitGeneration();
+            final long generation = inMemorySegInfos.getGeneration();
+            if (generation < lastCommitGeneration) {
+                // the latest in memory infos is behind disk, read the latest SegmentInfos from disk and return that.
+                final SegmentInfos latestCommit = store.readLastCommittedSegmentsInfo();
+                result = latestCommit;
+            } else {
+                result = inMemorySegInfos;
+            }
+            indexWriter.incRefDeleter(result);
         } catch (IOException e) {
             throw new EngineException(shardId, e.getMessage(), e);
         }
-        return new GatedCloseable<>(segmentInfos, () -> indexWriter.decRefDeleter(segmentInfos));
+        return new GatedCloseable<>(result, () -> indexWriter.decRefDeleter(result));
     }
 
     @Override

server/src/main/java/org/opensearch/index/store/Store.java

@@ -334,6 +334,13 @@ public MetadataSnapshot getMetadata(SegmentInfos segmentInfos) throws IOExceptio
         return new MetadataSnapshot(segmentInfos, directory, logger);
     }
 
+    /**
+     * Returns generation number (_N in segment_N) from last commit on disk.
+     */
+    public long getLastCommitGeneration() throws IOException {
+        return SegmentInfos.getLastCommitGeneration(directory());
+    }
+
     /**
      * Renames all the given files from the key of the map to the
      * value of the map. All successfully renamed files are removed from the map in-place.

server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java

@@ -7546,6 +7546,61 @@ public void testGetSegmentInfosSnapshot() throws IOException {
         engine.close();
     }
 
+    // This method simulates behaviour of on-disk SegmentInfos having higher segment generation number compared to
+    // actual in-memory SegmentInfos and verified that engine.getSegmentInfosSnapshot returns on-disk SegmentInfos
+    public void testGetSegmentInfosWithHighestGenOnDisk() throws IOException {
+        IOUtils.close(store, engine);
+        Store store = spy(createStore());
+        InternalEngine engine = createEngine(store, createTempDir());
+        final int numDocs = randomIntBetween(10, 100);
+        for (int docId = 0; docId < numDocs; docId++) {
+            index(engine, docId);
+            if (randomBoolean()) {
+                engine.refresh("test");
+            }
+        }
+        engine.flush(true, true);
+
+        SegmentInfos sisDisk = store.readLastCommittedSegmentsInfo();
+        // Increment generation number of on-disk SegmentInfos
+        sisDisk.setNextWriteGeneration(sisDisk.getGeneration() + 1);
+
+        when(store.getLastCommitGeneration()).thenReturn(sisDisk.getGeneration());
+        when(store.readLastCommittedSegmentsInfo()).thenReturn(sisDisk);
+
+        GatedCloseable<SegmentInfos> segmentInfosSnapshot = engine.getSegmentInfosSnapshot();
+        assertEquals(segmentInfosSnapshot.get(), sisDisk);
+        assertEquals(segmentInfosSnapshot.get().getGeneration(), sisDisk.getGeneration());
+        segmentInfosSnapshot.close();
+        store.close();
+        engine.close();
+    }
+
+    // This method verifies that when on-disk segment generation is not higher compared to memory copy, then
+    // engine.getSegmentInfosSnapsho returns in-memory SegmentInfos
+    public void testGetSegmentInfosWithHighestGenInMemory() throws IOException {
+        IOUtils.close(store, engine);
+        Store store = spy(createStore());
+        InternalEngine engine = createEngine(store, createTempDir());
+        final int numDocs = randomIntBetween(10, 100);
+        for (int docId = 0; docId < numDocs; docId++) {
+            index(engine, docId);
+            if (randomBoolean()) {
+                engine.refresh("test");
+            }
+        }
+        engine.flush(true, true);
+
+        SegmentInfos sisInMemory = engine.getLatestSegmentInfos();
+
+        GatedCloseable<SegmentInfos> segmentInfosSnapshot = engine.getSegmentInfosSnapshot();
+        assertEquals(segmentInfosSnapshot.get(), sisInMemory);
+        assertEquals(segmentInfosSnapshot.get().getGeneration(), sisInMemory.getGeneration());
+        segmentInfosSnapshot.close();
+        store.close();
+        engine.close();
+    }
+
     public void testGetProcessedLocalCheckpoint() throws IOException {
         final long expectedLocalCheckpoint = 1L;
         IOUtils.close(store, engine);