Modifying derived source related ITs in RelocationIT to make it more deterministic

Signed-off-by: Tanik Pansuriya <panbhai@amazon.com>

Author: tanik98

CHANGELOG.md

@@ -45,43 +45,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added approximation support for range queries with now in date field ([#18511](https://github.com/opensearch-project/OpenSearch/pull/18511))
 - Upgrade to protobufs 0.6.0 and clean up deprecated TermQueryProtoUtils code ([#18880](https://github.com/opensearch-project/OpenSearch/pull/18880))
 - Prevent shard initialization failure due to streaming consumer errors ([#18877](https://github.com/opensearch-project/OpenSearch/pull/18877))
-- Add support for linux riscv64 platform ([#18156](https://github.com/opensearch-project/OpenSearch/pull/18156))
-- [Rule based auto-tagging] Add get rule API ([#17336](https://github.com/opensearch-project/OpenSearch/pull/17336))
-- [Rule based auto-tagging] Add Delete Rule API ([#18184](https://github.com/opensearch-project/OpenSearch/pull/18184))
-- Add paginated wlm/stats API ([#17638](https://github.com/opensearch-project/OpenSearch/pull/17638))
-- [Rule based auto-tagging] Add Create rule API ([#17792](https://github.com/opensearch-project/OpenSearch/pull/17792))
-- Implement parallel shard refresh behind cluster settings ([#17782](https://github.com/opensearch-project/OpenSearch/pull/17782))
-- Bump OpenSearch Core main branch to 3.0.0 ([#18039](https://github.com/opensearch-project/OpenSearch/pull/18039))
-- [Rule based Auto-tagging] Add wlm `ActionFilter` ([#17791](https://github.com/opensearch-project/OpenSearch/pull/17791))
-- Update API of Message in index to add the timestamp for lag calculation in ingestion polling ([#17977](https://github.com/opensearch-project/OpenSearch/pull/17977/))
-- Add Warm Disk Threshold Allocation Decider for Warm shards ([#18082](https://github.com/opensearch-project/OpenSearch/pull/18082))
-- Add composite directory factory ([#17988](https://github.com/opensearch-project/OpenSearch/pull/17988))
-- [Rule based auto-tagging] Add refresh based synchronization service for `Rule`s ([#18128](https://github.com/opensearch-project/OpenSearch/pull/18128))
-- Add pull-based ingestion error metrics and make internal queue size configurable ([#18088](https://github.com/opensearch-project/OpenSearch/pull/18088))
-- [Derive Source] Adding support for derive source feature and implementing it for various type of field mappers ([#17759](https://github.com/opensearch-project/OpenSearch/pull/17759))
-- [Derive Source] Adding integration of derived source feature across diff paths ([#18054](https://github.com/opensearch-project/OpenSearch/pull/18054))
-- [Security Manager Replacement] Enhance Java Agent to intercept newByteChannel ([#17989](https://github.com/opensearch-project/OpenSearch/pull/17989))
-- Enabled Async Shard Batch Fetch by default ([#18139](https://github.com/opensearch-project/OpenSearch/pull/18139))
-- Allow to get the search request from the QueryCoordinatorContext ([#17818](https://github.com/opensearch-project/OpenSearch/pull/17818))
-- Reject close index requests, while remote store migration is in progress.([#18327](https://github.com/opensearch-project/OpenSearch/pull/18327))
-- Improve sort-query performance by retaining the default `totalHitsThreshold` for approximated `match_all` queries ([#18189](https://github.com/opensearch-project/OpenSearch/pull/18189))
-- Enable testing for ExtensiblePlugins using classpath plugins ([#16908](https://github.com/opensearch-project/OpenSearch/pull/16908))
-- Introduce system generated ingest pipeline ([#17817](https://github.com/opensearch-project/OpenSearch/pull/17817)))
-- Apply cluster state metadata and routing table diff when building cluster state from remote([#18256](https://github.com/opensearch-project/OpenSearch/pull/18256))
-- Support create mode in pull-based ingestion and add retries for transient failures ([#18250](https://github.com/opensearch-project/OpenSearch/pull/18250)))
-- Decouple the init of Crypto Plugin and KeyProvider in CryptoRegistry ([18270](https://github.com/opensearch-project/OpenSearch/pull18270)))
-- Support cluster write block in pull-based ingestion ([#18280](https://github.com/opensearch-project/OpenSearch/pull/18280)))
-- Use QueryCoordinatorContext for the rewrite in validate API. ([#18272](https://github.com/opensearch-project/OpenSearch/pull/18272))
-- Upgrade crypto kms plugin dependencies for AWS SDK v2.x. ([#18268](https://github.com/opensearch-project/OpenSearch/pull/18268))
-- Add support for `matched_fields` with the unified highlighter ([#18164](https://github.com/opensearch-project/OpenSearch/issues/18164))
-- Add BooleanQuery rewrite for must_not RangeQuery clauses ([#17655](https://github.com/opensearch-project/OpenSearch/pull/17655))
-- [repository-s3] Add support for SSE-KMS and S3 bucket owner verification ([#18312](https://github.com/opensearch-project/OpenSearch/pull/18312))
-- Optimize gRPC perf by passing by reference ([#18303](https://github.com/opensearch-project/OpenSearch/pull/18303))
-- Added File Cache Stats - Involves Block level as well as full file level stats ([#17538](https://github.com/opensearch-project/OpenSearch/issues/17479))
-- Added time_in_execution attribute to /_cluster/pending_tasks response ([#17780](https://github.com/opensearch-project/OpenSearch/pull/17780))
-- Added File Cache Pinning ([#17617](https://github.com/opensearch-project/OpenSearch/issues/13648))
-- Support consumer reset in Resume API for pull-based ingestion. This PR includes a breaking change for the experimental pull-based ingestion feature. ([#18332](https://github.com/opensearch-project/OpenSearch/pull/18332))
-- Add FIPS build tooling ([#4254](https://github.com/opensearch-project/security/issues/4254))
 - [Derived Source] Add integration of derived source feature across various paths like get/search/recovery ([#18565](https://github.com/opensearch-project/OpenSearch/pull/18565))
 
 ### Changed

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

@@ -97,14 +97,14 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -916,24 +916,23 @@ public void testRelocationWithDerivedSourceAndConcurrentIndexing() throws Except
             ).setMapping(mapping)
         );
 
-        // Start background indexing
-        AtomicBoolean stopIndexing = new AtomicBoolean(false);
-        AtomicInteger docCount = new AtomicInteger(0);
+        // Start indexing
+        int docCount = randomIntBetween(10, 100);
+        CountDownLatch docCountLatch = new CountDownLatch(docCount);
         Thread indexingThread = new Thread(() -> {
-            while (stopIndexing.get() == false) {
+            while (docCountLatch.getCount() > 0) {
                 try {
-                    int id = docCount.incrementAndGet();
+                    long id = docCountLatch.getCount();
                     client().prepareIndex("test").setId(String.valueOf(id)).setSource("name", "test" + id, "value", id).get();
+                    docCountLatch.countDown();
                     Thread.sleep(10); // Small delay to prevent overwhelming
                 } catch (Exception e) {
                     logger.error("Error in background indexing", e);
                 }
             }
         });
-        indexingThread.start();
 
-        // Let it index some documents
-        Thread.sleep(2000);
+        indexingThread.start();
 
         // Start relocation
         String node2 = internalCluster().startNode();
@@ -943,20 +942,18 @@ public void testRelocationWithDerivedSourceAndConcurrentIndexing() throws Except
         client().admin().cluster().prepareReroute().add(new MoveAllocationCommand("test", 0, node1, node2)).get();
         ensureGreen(TimeValue.timeValueMinutes(2));
 
-        // Stop indexing
-        stopIndexing.set(true);
+        // Let it index docCount documents
+        docCountLatch.await();
         indexingThread.join();
 
-        // Verify all documents
-        int finalDocCount = docCount.get();
         assertBusy(() -> {
             refresh();
             SearchResponse response = client().prepareSearch("test")
                 .setQuery(matchAllQuery())
-                .setSize(finalDocCount)
+                .setSize(docCount)
                 .addSort("value", SortOrder.ASC)
                 .get();
-            assertHitCount(response, finalDocCount);
+            assertHitCount(response, docCount);
 
             int expectedId = 1;
             for (SearchHit hit : response.getHits()) {
@@ -1006,37 +1003,28 @@ public void testRelocationWithDerivedSourceWithUpdates() throws Exception {
         logger.info("--> relocate the shard from node1 to node2");
         client().admin().cluster().prepareReroute().add(new MoveAllocationCommand("test", 0, node1, node2)).get();
 
-        AtomicBoolean stop = new AtomicBoolean(false);
+        final Set<Integer> docsToUpdate = new HashSet<>();
         final Set<Integer> updatedDocs = ConcurrentCollections.newConcurrentSet();
+        int updateCount = randomIntBetween(10, numDocs / 2);
+        for (int i = 0; i < updateCount; i++) {
+            docsToUpdate.add(randomIntBetween(0, numDocs - 1));
+        }
 
-        // Start doc update thread
-        Thread updateThread = new Thread(() -> {
-            while (stop.get() == false) {
-                try {
-                    int docId = randomIntBetween(0, numDocs - 1);
-                    client().prepareUpdate("test", String.valueOf(docId))
-                        .setRetryOnConflict(3)
-                        .setDoc("value", docId * 2)
-                        .execute()
-                        .actionGet();
-                    updatedDocs.add(docId);
-                    Thread.sleep(10);
-                } catch (Exception e) {
-                    if (e instanceof InterruptedException) {
-                        break;
-                    }
-                    logger.warn("Error in update thread", e);
-                }
+        docsToUpdate.stream().forEach(docId -> {
+            try {
+                client().prepareUpdate("test", String.valueOf(docId))
+                    .setRetryOnConflict(3)
+                    .setDoc("value", docId * 2)
+                    .execute()
+                    .actionGet();
+                updatedDocs.add(docId);
+                Thread.sleep(10);
+            } catch (Exception e) {
+                logger.warn("Error while updating doc with id = {}", docId, e);
             }
         });
-        updateThread.start();
-
-        // Wait for some updates to occur
-        Thread.sleep(2000);
-
-        stop.set(true);
-        updateThread.join();
 
+        assertEquals(docsToUpdate.size(), updatedDocs.size());
         refresh("test");
         ensureGreen(TimeValue.timeValueMinutes(2));