Addressed comments

Signed-off-by: Ankit Kala <ankikala@amazon.com>
opensearch-project · Jun 14, 2023 · 2d564bc · 2d564bc
1 parent 4a350a8
commit 2d564bc
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 13 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/index/SegmentReplicationPressureIT.java b/server/src/internalClusterTest/java/org/opensearch/index/SegmentReplicationPressureIT.java
@@ -30,7 +30,6 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
-import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -326,15 +325,6 @@ public void testBulkWritesRejected() throws Exception {
         verifyStoreContent();
     }
 
-    private void assertReplicaCheckpointUpdated(IndexShard primaryShard) throws Exception {
-        assertBusy(() -> {
-            Set<SegmentReplicationShardStats> groupStats = primaryShard.getReplicationStats();
-            for (SegmentReplicationShardStats shardStat : groupStats) {
-                assertEquals(0, shardStat.getCheckpointsBehindCount());
-            }
-        }, 30, TimeUnit.SECONDS);
-    }
-
     private BulkResponse executeBulkRequest(List<String> nodes, int docsPerBatch) {
         final BulkRequest bulkRequest = new BulkRequest();
         for (int j = 0; j < docsPerBatch; ++j) {

diff --git a/...internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationBaseIT.java b/...internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationBaseIT.java
@@ -23,6 +23,7 @@
 import org.opensearch.index.IndexModule;
 import org.opensearch.index.IndexService;
 import org.opensearch.index.SegmentReplicationPerGroupStats;
+import org.opensearch.index.SegmentReplicationShardStats;
 import org.opensearch.index.shard.IndexShard;
 import org.opensearch.index.store.Store;
 import org.opensearch.index.store.StoreFileMetadata;
@@ -38,6 +39,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
@@ -222,4 +224,13 @@ protected Releasable blockReplication(List<String> nodes, CountDownLatch latch)
         };
     }
 
+    protected void assertReplicaCheckpointUpdated(IndexShard primaryShard) throws Exception {
+        assertBusy(() -> {
+            Set<SegmentReplicationShardStats> groupStats = primaryShard.getReplicationStats();
+            assertEquals(primaryShard.indexSettings().getNumberOfReplicas(), groupStats.size());
+            for (SegmentReplicationShardStats shardStat : groupStats) {
+                assertEquals(0, shardStat.getCheckpointsBehindCount());
+            }
+        }, 30, TimeUnit.SECONDS);
+    }
 }
diff --git a/...src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationIT.java b/...src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationIT.java
@@ -78,7 +78,6 @@
 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 java.util.Arrays.asList;
@@ -877,7 +876,7 @@ public void testPressureServiceStats() throws Exception {
                 assertEquals(1, shardStatsSet.size());
                 final SegmentReplicationShardStats stats = shardStatsSet.stream().findFirst().get();
                 assertEquals(0, stats.getCheckpointsBehindCount());
-            }, 30, TimeUnit.SECONDS);
+            });
         }
     }
 

diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java
@@ -330,7 +330,6 @@ public void onFailure(Exception e) {
                     request.getCheckpoint(),
                     e
                 );
-                transportClient.cancel();
             }
         };