Covering more recovery cases for derived source

Signed-off-by: Tanik Pansuriya <[email protected]>

Author: tanik98

server/src/internalClusterTest/java/org/opensearch/recovery/FullRollingRestartIT.java

@@ -573,21 +573,6 @@ public void testDerivedSourceWithConcurrentUpdatesRollingRestart() throws Except
 
         // Initial indexing
         int docCount = randomIntBetween(100, 200); // Reduced count for stability
-        bulkIndexDocumentsWithVersion(docCount);
-
-        refresh("test");
-        flush("test");
-        ensureGreen();
-
-        // Verify initial document count
-        assertHitCount(client().prepareSearch("test").setSize(0).get(), docCount);
-
-        // Start concurrent updates during rolling restart
-        logger.info("--> starting rolling restart with concurrent updates");
-        concurrentUpdatesRollingRestartWithVerification(docCount);
-    }
-
-    private void bulkIndexDocumentsWithVersion(int docCount) {
         BulkRequestBuilder bulkRequest = client().prepareBulk();
         for (int i = 0; i < docCount; i++) {
             bulkRequest.add(
@@ -606,6 +591,17 @@ private void bulkIndexDocumentsWithVersion(int docCount) {
             BulkResponse response = bulkRequest.execute().actionGet();
             assertFalse(response.hasFailures());
         }
+
+        refresh("test");
+        flush("test");
+        ensureGreen();
+
+        // Verify initial document count
+        assertHitCount(client().prepareSearch("test").setSize(0).get(), docCount);
+
+        // Start concurrent updates during rolling restart
+        logger.info("--> starting rolling restart with concurrent updates");
+        concurrentUpdatesRollingRestartWithVerification(docCount);
     }
 
     private void concurrentUpdatesRollingRestartWithVerification(int initialDocCount) throws Exception {
@@ -622,14 +618,7 @@ private void concurrentUpdatesRollingRestartWithVerification(int initialDocCount
                         int docId = randomIntBetween(0, initialDocCount - 1);
                         client().prepareUpdate("test", String.valueOf(docId))
                             .setRetryOnConflict(3)
-                            .setDoc(
-                                "counter",
-                                randomIntBetween(0, 1000),
-                                "last_updated",
-                                System.currentTimeMillis(),
-                                "version",
-                                System.currentTimeMillis()
-                            )
+                            .setDoc("counter", randomIntBetween(0, 1000), "last_updated", System.currentTimeMillis(), "version", 1)
                             .execute()
                             .actionGet();
                         totalUpdates.incrementAndGet();
@@ -686,7 +675,7 @@ private void verifyDerivedSourceWithUpdates(int expectedDocs) throws Exception {
         assertBusy(() -> {
             SearchResponse response = client().prepareSearch("test")
                 .setSize(expectedDocs)
-                .addSort("version", SortOrder.DESC) // Sort by version to ensure we see latest updates
+                .addSort("last_updated", SortOrder.DESC) // Sort by version to ensure we see latest updates
                 .get();
             assertHitCount(response, expectedDocs);
 
@@ -695,10 +684,10 @@ private void verifyDerivedSourceWithUpdates(int expectedDocs) throws Exception {
                 String id = hit.getId();
 
                 // Verify all required fields are present
-                assertNotNull("text_field missing for doc " + id, source.get("text_field"));
+                assertEquals("text value " + id, source.get("text_field"));
                 assertNotNull("counter missing for doc " + id, source.get("counter"));
-                assertNotNull("last_updated missing for doc " + id, source.get("last_updated"));
-                assertNotNull("version missing for doc " + id, source.get("version"));
+                assertFalse(((String) source.get("last_updated")).isEmpty());
+                assertEquals(1, source.get("version"));
 
                 // Verify text_field maintains original value
                 assertEquals("text value " + id, source.get("text_field"));

server/src/internalClusterTest/java/org/opensearch/recovery/RecoveryWhileUnderLoadIT.java

@@ -522,7 +522,7 @@ public void testRecoveryWithDerivedSourceEnabled() throws Exception {
         assertAcked(
             prepareCreate(
                 "test",
-                2,
+                1,
                 Settings.builder()
                     .put(SETTING_NUMBER_OF_SHARDS, numberOfShards)
                     .put(SETTING_NUMBER_OF_REPLICAS, 1)
@@ -541,7 +541,7 @@ public void testRecoveryWithDerivedSourceEnabled() throws Exception {
             }
         }
 
-        logger.info("--> allow 2 nodes for index [test] ...");
+        logger.info("--> allow 2 nodes for index [test] with replica ...");
         allowNodes("test", 2);
 
         logger.info("--> waiting for GREEN health status ...");
@@ -585,6 +585,7 @@ public void testReplicaRecoveryWithDerivedSourceBeforeRefresh() throws Exception
                     .put(SETTING_NUMBER_OF_REPLICAS, 0)
                     .put(IndexSettings.INDEX_DERIVED_SOURCE_SETTING.getKey(), true)
                     .put(IndexSettings.INDEX_TRANSLOG_DURABILITY_SETTING.getKey(), Translog.Durability.ASYNC)
+                    .put("index.refresh_interval", -1)
             ).setMapping(mapping)
         );
 
@@ -668,7 +669,9 @@ public void testReplicaRecoveryWithDerivedSourceFromTranslog() throws Exception
 
         // Kill replica node and index more documents
         final String replicaNode = ensureReplicaNode("test");
-        internalCluster().stopRandomNode(InternalTestCluster.nameFilter(replicaNode));
+        if (replicaNode != null) {
+            internalCluster().stopRandomNode(InternalTestCluster.nameFilter(replicaNode));
+        }
 
         int additionalDocs = randomIntBetween(50, 100);
         for (int i = docCount; i < docCount + additionalDocs; i++) {
@@ -722,7 +725,7 @@ public void testRecoverWhileUnderLoadWithDerivedSource() throws Exception {
         assertAcked(
             prepareCreate(
                 "test",
-                2,
+                1,
                 Settings.builder()
                     .put(SETTING_NUMBER_OF_SHARDS, numberOfShards)
                     .put(SETTING_NUMBER_OF_REPLICAS, 1)
@@ -781,18 +784,37 @@ protected XContentBuilder generateSource(long id, Random random) throws IOExcept
             logger.info("--> verifying indexed content");
 
             // Verify docs on primary
-            assertPrimaryDocCount(totalNumDocs);
+            SearchResponse primaryResponse = client().prepareSearch("test").setPreference("_primary").setTrackTotalHits(true).get();
+            assertHitCount(primaryResponse, totalNumDocs);
 
             // Verify docs and derived source on replica
-            assertReplicaDocsAndSource(totalNumDocs);
+            assertBusy(() -> {
+                SearchResponse replicaResponse = client().prepareSearch("test")
+                    .setPreference("_replica")
+                    .setTrackTotalHits(true)
+                    .setSize(totalNumDocs)
+                    .addSort("value", SortOrder.ASC)
+                    .get();
+
+                assertHitCount(replicaResponse, totalNumDocs);
+
+                // Verify source reconstruction on replica
+                for (SearchHit hit : replicaResponse.getHits()) {
+                    assertNotNull(hit.getSourceAsMap());
+                    assertEquals(3, hit.getSourceAsMap().size());
+                    int id = (Integer) hit.getSourceAsMap().get("value");
+                    assertEquals("name_" + id, hit.getSourceAsMap().get("name"));
+                    assertNotNull(hit.getSourceAsMap().get("timestamp"));
+                }
+            }, 30, TimeUnit.SECONDS);
 
             // Additional source verification with random sampling
             assertRandomDocsSource(50);
         }
     }
 
     public void testRecoverWithRelocationAndDerivedSource() throws Exception {
-        final int numShards = between(2, 5);
+        final int numShards = between(3, 5);
         logger.info("--> creating test index with derived source enabled...");
 
         String mapping = """
@@ -813,7 +835,7 @@ public void testRecoverWithRelocationAndDerivedSource() throws Exception {
         assertAcked(
             prepareCreate(
                 "test",
-                3,
+                1,
                 Settings.builder()
                     .put(SETTING_NUMBER_OF_SHARDS, numShards)
                     .put(SETTING_NUMBER_OF_REPLICAS, 0)
@@ -822,9 +844,8 @@ public void testRecoverWithRelocationAndDerivedSource() throws Exception {
                     .put(IndexService.RETENTION_LEASE_SYNC_INTERVAL_SETTING.getKey(), "100ms")
             ).setMapping(mapping)
         );
-
-        final int numDocs = scaledRandomIntBetween(2000, 10000);
-        int allowNodes = 2;
+        int numNodes = 1;
+        final int numDocs = scaledRandomIntBetween(1500, 2000);
 
         try (BackgroundIndexer indexer = new BackgroundIndexer("test", null, client(), numDocs) {
             @Override
@@ -836,14 +857,13 @@ protected XContentBuilder generateSource(long id, Random random) throws IOExcept
                     .endObject();
             }
         }) {
-            for (int i = 0; i < numDocs; i += scaledRandomIntBetween(100, Math.min(1000, numDocs))) {
+
+            for (int i = 0; i < numDocs; i += scaledRandomIntBetween(500, Math.min(1000, numDocs))) {
                 indexer.assertNoFailures();
                 logger.info("--> waiting for {} docs to be indexed ...", i);
                 waitForDocs(i, indexer);
-
-                // Alternate between 1 and 2 nodes to force relocation
-                allowNodes = 3 - allowNodes;  // Toggle between 1 and 2
-                allowNodes("test", allowNodes);
+                internalCluster().startDataOnlyNode();
+                numNodes++;
 
                 logger.info("--> waiting for GREEN health status ...");
                 ensureGreen(TimeValue.timeValueMinutes(2));
@@ -854,12 +874,11 @@ protected XContentBuilder generateSource(long id, Random random) throws IOExcept
 
             // Add replicas after stopping indexing
             logger.info("--> adding replicas ...");
-            allowNodes("test", 3);
             assertAcked(
                 client().admin()
                     .indices()
                     .prepareUpdateSettings("test")
-                    .setSettings(Settings.builder().put(SETTING_NUMBER_OF_REPLICAS, 1))
+                    .setSettings(Settings.builder().put(SETTING_NUMBER_OF_REPLICAS, numNodes - 1))
                     .get()
             );
             ensureGreen(TimeValue.timeValueMinutes(2));
@@ -868,37 +887,31 @@ protected XContentBuilder generateSource(long id, Random random) throws IOExcept
             client().admin().indices().prepareRefresh().get();
 
             // Verify final doc count and derived source reconstruction
-            assertPrimaryDocCount(numDocs);
-            assertReplicaDocsAndSource(numDocs);
-            assertRandomDocsSource(100);
-        }
-    }
-
-    private void assertPrimaryDocCount(int expectedCount) {
-        SearchResponse primaryResponse = client().prepareSearch("test").setPreference("_primary").setTrackTotalHits(true).get();
-        assertHitCount(primaryResponse, expectedCount);
-    }
+            SearchResponse primaryResponse = client().prepareSearch("test").setPreference("_primary").setTrackTotalHits(true).get();
+            assertHitCount(primaryResponse, numDocs);
 
-    private void assertReplicaDocsAndSource(int expectedCount) throws Exception {
-        assertBusy(() -> {
-            SearchResponse replicaResponse = client().prepareSearch("test")
-                .setPreference("_replica")
-                .setTrackTotalHits(true)
-                .setSize(expectedCount)
-                .addSort("value", SortOrder.ASC)
-                .get();
-
-            assertHitCount(replicaResponse, expectedCount);
+            assertBusy(() -> {
+                SearchResponse replicaResponse = client().prepareSearch("test")
+                    .setPreference("_replica")
+                    .setTrackTotalHits(true)
+                    .setSize(numDocs)
+                    .addSort("value", SortOrder.ASC)
+                    .get();
+
+                assertHitCount(replicaResponse, numDocs);
+
+                // Verify source reconstruction on replica
+                for (SearchHit hit : replicaResponse.getHits()) {
+                    assertNotNull(hit.getSourceAsMap());
+                    assertEquals(3, hit.getSourceAsMap().size());
+                    int id = (Integer) hit.getSourceAsMap().get("value");
+                    assertEquals("name_" + id, hit.getSourceAsMap().get("name"));
+                    assertNotNull(hit.getSourceAsMap().get("timestamp"));
+                }
+            }, 30, TimeUnit.SECONDS);
 
-            // Verify source reconstruction on replica
-            for (SearchHit hit : replicaResponse.getHits()) {
-                assertNotNull(hit.getSourceAsMap());
-                assertEquals(3, hit.getSourceAsMap().size());
-                int id = (Integer) hit.getSourceAsMap().get("value");
-                assertEquals("name_" + id, hit.getSourceAsMap().get("name"));
-                assertNotNull(hit.getSourceAsMap().get("timestamp"));
-            }
-        }, 30, TimeUnit.SECONDS);
+            assertRandomDocsSource(100);
+        }
     }
 
     private void assertRandomDocsSource(int sampleSize) {
@@ -921,7 +934,11 @@ private void assertRandomDocsSource(int sampleSize) {
     private String ensureReplicaNode(String index) {
         ClusterState state = client().admin().cluster().prepareState().get().getState();
         Index idx = state.metadata().index(index).getIndex();
-        String nodeId = state.routingTable().index(idx).shard(0).replicaShards().get(0).currentNodeId();
-        return state.nodes().get(nodeId).getName();
+        String replicaNode = state.routingTable().index(idx).shard(0).replicaShards().get(0).currentNodeId();
+        String clusterManagerNode = internalCluster().getClusterManagerName();
+        if (!replicaNode.equals(clusterManagerNode)) {
+            state.nodes().get(replicaNode).getName();
+        }
+        return null;
     }
 }