From 9aee4e94b5ed123a321420c5c7c82671e743df51 Mon Sep 17 00:00:00 2001
From: Rishab Nahata <rnnahata@amazon.com>
Date: Thu, 8 Aug 2024 15:19:22 +0530
Subject: [PATCH] Add test

Signed-off-by: Rishab Nahata <rnnahata@amazon.com>
---
 .../gateway/GatewayAllocatorTests.java        | 21 +++++++++++++++++++
 .../TestShardBatchGatewayAllocator.java       | 21 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java b/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java
index 1596a0b566b28..c7eae77d6deba 100644
--- a/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java
+++ b/server/src/test/java/org/opensearch/gateway/GatewayAllocatorTests.java
@@ -32,6 +32,7 @@
 import org.opensearch.cluster.routing.allocation.decider.AllocationDeciders;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.BatchRunnableExecutor;
 import org.opensearch.common.util.set.Sets;
 import org.opensearch.core.index.shard.ShardId;
@@ -45,6 +46,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 import static org.opensearch.gateway.ShardsBatchGatewayAllocator.PRIMARY_BATCH_ALLOCATOR_TIMEOUT_SETTING;
@@ -423,6 +426,24 @@ public void testReplicaAllocatorTimeout() {
         assertEquals(-1, REPLICA_BATCH_ALLOCATOR_TIMEOUT_SETTING.get(build).getMillis());
     }
 
+    public void testCollectTimedOutShards() throws InterruptedException {
+        createIndexAndUpdateClusterState(2, 5, 2);
+        CountDownLatch latch = new CountDownLatch(10);
+        testShardsBatchGatewayAllocator = new TestShardBatchGatewayAllocator(latch);
+        testShardsBatchGatewayAllocator.setPrimaryBatchAllocatorTimeout(TimeValue.ZERO);
+        testShardsBatchGatewayAllocator.setReplicaBatchAllocatorTimeout(TimeValue.ZERO);
+        BatchRunnableExecutor executor = testShardsBatchGatewayAllocator.allocateAllUnassignedShards(testAllocation, true);
+        executor.run();
+        assertTrue(latch.await(1, TimeUnit.MINUTES));
+        latch = new CountDownLatch(10);
+        testShardsBatchGatewayAllocator = new TestShardBatchGatewayAllocator(latch);
+        testShardsBatchGatewayAllocator.setPrimaryBatchAllocatorTimeout(TimeValue.ZERO);
+        testShardsBatchGatewayAllocator.setReplicaBatchAllocatorTimeout(TimeValue.ZERO);
+        executor = testShardsBatchGatewayAllocator.allocateAllUnassignedShards(testAllocation, false);
+        executor.run();
+        assertTrue(latch.await(1, TimeUnit.MINUTES));
+    }
+
     private void createIndexAndUpdateClusterState(int count, int numberOfShards, int numberOfReplicas) {
         if (count == 0) return;
         Metadata.Builder metadata = Metadata.builder();
diff --git a/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java b/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java
index 0eb4bb6935bac..156b1d7c620e6 100644
--- a/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java
+++ b/test/framework/src/main/java/org/opensearch/test/gateway/TestShardBatchGatewayAllocator.java
@@ -29,13 +29,20 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 
 public class TestShardBatchGatewayAllocator extends ShardsBatchGatewayAllocator {
 
+    CountDownLatch latch;
+
     public TestShardBatchGatewayAllocator() {
 
     }
 
+    public TestShardBatchGatewayAllocator(CountDownLatch latch) {
+        this.latch = latch;
+    }
+
     public TestShardBatchGatewayAllocator(long maxBatchSize) {
         super(maxBatchSize);
     }
@@ -83,6 +90,13 @@ protected AsyncShardFetch.FetchResult<TransportNodesListGatewayStartedShardsBatc
             }
             return new AsyncShardFetch.FetchResult<>(foundShards, shardsToIgnoreNodes);
         }
+
+        @Override
+        protected void allocateUnassignedBatchOnTimeout(Set<ShardId> shardIds, RoutingAllocation allocation, boolean primary) {
+            for (int i = 0; i < shardIds.size(); i++) {
+                latch.countDown();
+            }
+        }
     };
 
     ReplicaShardBatchAllocator replicaBatchShardAllocator = new ReplicaShardBatchAllocator() {
@@ -100,6 +114,13 @@ protected AsyncShardFetch.FetchResult<TransportNodesListShardStoreMetadataBatch.
         protected boolean hasInitiatedFetching(ShardRouting shard) {
             return true;
         }
+
+        @Override
+        protected void allocateUnassignedBatchOnTimeout(Set<ShardId> shardIds, RoutingAllocation allocation, boolean primary) {
+            for (int i = 0; i < shardIds.size(); i++) {
+                latch.countDown();
+            }
+        }
     };
 
     @Override