Add unit tests

Signed-off-by: Sachin Kale <kalsac@amazon.com>
opensearch-project · Nov 16, 2022 · f7d9b2e · f7d9b2e
1 parent cd6bf70
commit f7d9b2e
Show file tree

Hide file tree

Showing 5 changed files with 66 additions and 17 deletions.
diff --git a/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java b/server/src/main/java/org/opensearch/index/shard/RemoteStoreRefreshListener.java
@@ -93,10 +93,11 @@ public void afterRefresh(boolean didRefresh) {
                     }
                     try {
                         String lastCommittedLocalSegmentFileName = SegmentInfos.getLastCommitSegmentsFileName(storeDirectory);
-                        if (!remoteDirectory.containsFile(
-                            lastCommittedLocalSegmentFileName,
-                            getChecksumOfLocalFile(lastCommittedLocalSegmentFileName)
-                        )) {
+                        if (lastCommittedLocalSegmentFileName != null
+                            && !remoteDirectory.containsFile(
+                                lastCommittedLocalSegmentFileName,
+                                getChecksumOfLocalFile(lastCommittedLocalSegmentFileName)
+                            )) {
                             deleteStaleCommits();
                         }
                         String segment_info_snapshot_filename = null;
@@ -120,7 +121,7 @@ public void afterRefresh(boolean didRefresh) {
 
                                 boolean uploadStatus = uploadNewSegments(refreshedLocalFiles);
                                 if (uploadStatus) {
-                                    if(segmentFilesFromSnapshot.equals(new HashSet<>(refreshedLocalFiles))) {
+                                    if (segmentFilesFromSnapshot.equals(new HashSet<>(refreshedLocalFiles))) {
                                         segment_info_snapshot_filename = uploadSegmentInfosSnapshot(latestSegmentInfos.get(), segmentInfos);
                                         refreshedLocalFiles.add(segment_info_snapshot_filename);
                                     }

diff --git a/server/src/main/java/org/opensearch/index/shard/StoreRecovery.java b/server/src/main/java/org/opensearch/index/shard/StoreRecovery.java
@@ -39,6 +39,7 @@
 import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.store.BufferedChecksumIndexInput;
+import org.apache.lucene.store.ChecksumIndexInput;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.FilterDirectory;
 import org.apache.lucene.store.IOContext;
@@ -473,18 +474,22 @@ private void recoverFromRemoteStore(IndexShard indexShard) throws IndexShardReco
                 }
             }
 
-            if(segmentInfosSnapshotFilename != null) {
+            if (segmentInfosSnapshotFilename != null) {
                 String[] filenameTokens = segmentInfosSnapshotFilename.split("__");
-                SegmentInfos infos_snapshot = SegmentInfos.readCommit(
-                    store.directory(),
-                    new BufferedChecksumIndexInput(storeDirectory.openInput(segmentInfosSnapshotFilename, IOContext.DEFAULT)),
-                    Integer.parseInt(filenameTokens[1])
-                );
-                store.commitSegmentInfos(
-                    infos_snapshot,
-                    Long.parseLong(filenameTokens[2]),
-                    Long.parseLong(filenameTokens[2])
-                );
+                try (
+                    ChecksumIndexInput indexInput = new BufferedChecksumIndexInput(
+                        storeDirectory.openInput(segmentInfosSnapshotFilename, IOContext.DEFAULT)
+                    )
+                ) {
+                    SegmentInfos infos_snapshot = SegmentInfos.readCommit(
+                        store.directory(),
+                        indexInput,
+                        Integer.parseInt(filenameTokens[1])
+                    );
+                    store.commitSegmentInfos(infos_snapshot, Long.parseLong(filenameTokens[2]), Long.parseLong(filenameTokens[2]));
+                } catch(IOException e) {
+                    logger.info("Exception while reading {}, falling back to commit level restore", segmentInfosSnapshotFilename);
+                }
             }
 
             // This creates empty trans-log for now

diff --git a/server/src/test/java/org/opensearch/index/shard/IndexShardTests.java b/server/src/test/java/org/opensearch/index/shard/IndexShardTests.java
@@ -2739,7 +2739,10 @@ public void restoreShard(
     public void testRestoreShardFromRemoteStore() throws IOException {
         IndexShard target = newStartedShard(
             true,
-            Settings.builder().put(IndexMetadata.SETTING_REMOTE_STORE_ENABLED, true).build(),
+            Settings.builder()
+                .put(IndexMetadata.SETTING_REPLICATION_TYPE, ReplicationType.SEGMENT)
+                .put(IndexMetadata.SETTING_REMOTE_STORE_ENABLED, true)
+                .build(),
             new InternalEngineFactory()
         );
 

diff --git a/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java b/server/src/test/java/org/opensearch/index/shard/RemoteStoreRefreshListenerTests.java
@@ -8,6 +8,7 @@
 
 package org.opensearch.index.shard;
 
+import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.FilterDirectory;
@@ -30,6 +31,8 @@
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 
+import static org.opensearch.index.shard.RemoteStoreRefreshListener.SEGMENT_INFO_SNAPSHOT_FILENAME;
+
 public class RemoteStoreRefreshListenerTests extends IndexShardTestCase {
     private IndexShard indexShard;
     private RemoteStoreRefreshListener remoteStoreRefreshListener;
@@ -204,13 +207,23 @@ public void onFailure(Exception e) {
     private void verifyUploadedSegments(RemoteSegmentStoreDirectory remoteSegmentStoreDirectory) throws IOException {
         Map<String, RemoteSegmentStoreDirectory.UploadedSegmentMetadata> uploadedSegments = remoteSegmentStoreDirectory
             .getSegmentsUploadedToRemoteStore();
+        String segmentsNFilename = null;
         try (GatedCloseable<SegmentInfos> segmentInfosGatedCloseable = indexShard.getSegmentInfosSnapshot()) {
             SegmentInfos segmentInfos = segmentInfosGatedCloseable.get();
             for (String file : segmentInfos.files(true)) {
                 if (!RemoteStoreRefreshListener.EXCLUDE_FILES.contains(file)) {
                     assertTrue(uploadedSegments.containsKey(file));
                 }
+                if (file.startsWith(IndexFileNames.SEGMENTS)) {
+                    segmentsNFilename = file;
+                }
             }
         }
+        if (segmentsNFilename != null) {
+            String commitGeneration = segmentsNFilename.substring((IndexFileNames.SEGMENTS + "_").length());
+            assertTrue(
+                uploadedSegments.keySet().stream().anyMatch(s -> s.startsWith(SEGMENT_INFO_SNAPSHOT_FILENAME + "__" + commitGeneration))
+            );
+        }
     }
 }
diff --git a/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryTests.java b/server/src/test/java/org/opensearch/index/store/RemoteSegmentStoreDirectoryTests.java
@@ -317,6 +317,33 @@ public void testCopyFromException() throws IOException {
         storeDirectory.close();
     }
 
+    public void testCpoyFromOverride() throws IOException {
+        String filename = "_100.si";
+        populateMetadata();
+        remoteSegmentStoreDirectory.init();
+
+        Directory storeDirectory = LuceneTestCase.newDirectory();
+        IndexOutput indexOutput = storeDirectory.createOutput(filename, IOContext.DEFAULT);
+        indexOutput.writeString("Hello World!");
+        CodecUtil.writeFooter(indexOutput);
+        indexOutput.close();
+        storeDirectory.sync(List.of(filename));
+
+        assertFalse(remoteSegmentStoreDirectory.getSegmentsUploadedToRemoteStore().containsKey(filename));
+        remoteSegmentStoreDirectory.copyFrom(storeDirectory, filename, filename, IOContext.DEFAULT, true);
+        RemoteSegmentStoreDirectory.UploadedSegmentMetadata uploadedSegmentMetadata = remoteSegmentStoreDirectory
+            .getSegmentsUploadedToRemoteStore()
+            .get(filename);
+        assertNotNull(uploadedSegmentMetadata);
+        remoteSegmentStoreDirectory.copyFrom(storeDirectory, filename, filename, IOContext.DEFAULT, true);
+        assertEquals(
+            uploadedSegmentMetadata.toString(),
+            remoteSegmentStoreDirectory.getSegmentsUploadedToRemoteStore().get(filename).toString()
+        );
+
+        storeDirectory.close();
+    }
+
     public void testContainsFile() throws IOException {
         List<String> metadataFiles = List.of("metadata__1__5__abc");
         when(remoteMetadataDirectory.listFilesByPrefix(RemoteSegmentStoreDirectory.MetadataFilenameUtils.METADATA_PREFIX)).thenReturn(