Avoid primary shard failure caused by merge segment warmer exceptions

Signed-off-by: guojialiang <guojialiang.2012@bytedance.com>

Author: guojialiang92

server/src/internalClusterTest/java/org/opensearch/indices/replication/MergedSegmentWarmerIT.java

@@ -14,6 +14,7 @@
 import org.opensearch.action.admin.indices.segments.IndexShardSegments;
 import org.opensearch.action.admin.indices.segments.IndicesSegmentResponse;
 import org.opensearch.action.admin.indices.segments.ShardSegments;
+import org.opensearch.action.index.IndexRequestBuilder;
 import org.opensearch.action.support.WriteRequest;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
@@ -38,6 +39,7 @@
 import java.util.stream.Collectors;
 
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
+import static org.hamcrest.Matchers.equalTo;
 
 /**
  * This class runs Segment Replication Integ test suite with merged segment warmer enabled.
@@ -59,6 +61,33 @@ protected Settings featureFlagSettings() {
         return featureSettings.build();
     }
 
+    public void testPrimaryNodeRestart() throws Exception {
+        logger.info("--> start nodes");
+        internalCluster().startNode();
+
+        logger.info("--> creating test index: {}", INDEX_NAME);
+        createIndex(INDEX_NAME, Settings.builder().put(indexSettings()).put("number_of_shards", 1).put("number_of_replicas", 0).build());
+
+        ensureGreen();
+
+        logger.info("--> indexing sample data");
+        final int numDocs = 100;
+        final IndexRequestBuilder[] docs = new IndexRequestBuilder[numDocs];
+
+        for (int i = 0; i < numDocs; i++) {
+            docs[i] = client().prepareIndex(INDEX_NAME)
+                .setSource("foo-int", randomInt(), "foo-string", randomAlphaOfLength(32), "foo-float", randomFloat());
+        }
+
+        indexRandom(true, docs);
+        flush();
+        assertThat(client().prepareSearch(INDEX_NAME).setSize(0).get().getHits().getTotalHits().value(), equalTo((long) numDocs));
+
+        logger.info("--> restarting cluster");
+        internalCluster().fullRestart();
+        ensureGreen();
+    }
+
     public void testMergeSegmentWarmer() throws Exception {
         final String primaryNode = internalCluster().startDataOnlyNode();
         final String replicaNode = internalCluster().startDataOnlyNode();

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

@@ -51,29 +51,39 @@ public MergedSegmentWarmer(
 
     @Override
     public void warm(LeafReader leafReader) throws IOException {
-        if (shouldWarm() == false) {
-            return;
-        }
-        // IndexWriter.IndexReaderWarmer#warm is called by IndexWriter#mergeMiddle. The type of leafReader should be SegmentReader.
-        assert leafReader instanceof SegmentReader;
-        assert indexShard.indexSettings().isSegRepLocalEnabled() || indexShard.indexSettings().isRemoteStoreEnabled();
+        try {
+            if (shouldWarm() == false) {
+                return;
+            }
+            // IndexWriter.IndexReaderWarmer#warm is called by IndexWriter#mergeMiddle. The type of leafReader should be SegmentReader.
+            assert leafReader instanceof SegmentReader;
+            assert indexShard.indexSettings().isSegRepLocalEnabled() || indexShard.indexSettings().isRemoteStoreEnabled();
 
-        long startTime = System.currentTimeMillis();
-        SegmentCommitInfo segmentCommitInfo = ((SegmentReader) leafReader).getSegmentInfo();
-        logger.trace(() -> new ParameterizedMessage("Warming segment: {}", segmentCommitInfo));
-        indexShard.publishMergedSegment(segmentCommitInfo);
-        logger.trace(() -> {
-            long segmentSize = -1;
-            try {
-                segmentSize = segmentCommitInfo.sizeInBytes();
-            } catch (IOException ignored) {}
-            return new ParameterizedMessage(
-                "Completed segment warming for {}. Size: {}B, Timing: {}ms",
-                segmentCommitInfo.info.name,
-                segmentSize,
-                (System.currentTimeMillis() - startTime)
+            long startTime = System.currentTimeMillis();
+            SegmentCommitInfo segmentCommitInfo = ((SegmentReader) leafReader).getSegmentInfo();
+            logger.info(() -> new ParameterizedMessage("Warming segment: {}", segmentCommitInfo));
+            indexShard.publishMergedSegment(segmentCommitInfo);
+            logger.trace(() -> {
+                long segmentSize = -1;
+                try {
+                    segmentSize = segmentCommitInfo.sizeInBytes();
+                } catch (IOException ignored) {}
+                return new ParameterizedMessage(
+                    "Completed segment warming for {}. Size: {}B, Timing: {}ms",
+                    segmentCommitInfo.info.name,
+                    segmentSize,
+                    (System.currentTimeMillis() - startTime)
+                );
+            });
+        } catch (Exception e) {
+            logger.warn(
+                () -> new ParameterizedMessage(
+                    "throw exception during merged segment warmer, skip merged segment {} warmer",
+                    ((SegmentReader) leafReader).getSegmentName()
+                ),
+                e
             );
-        });
+        }
     }
 
     // package-private for tests