From 8b5a10c15873922076520a35589973788a181486 Mon Sep 17 00:00:00 2001
From: Rishikesh Pasham <62345295+Rishikesh1159@users.noreply.github.com>
Date: Thu, 21 Jul 2022 13:59:41 +0000
Subject: [PATCH] Fix possible flaky test for
 testBeforeIndexShardClosed_CancelsOngoingReplications() (#3963)

* Fixing flaky test failure happening for testShardAlreadyReplicating()

Signed-off-by: Rishikesh1159 <rishireddy1159@gmail.com>
---
 .../indices/replication/SegmentReplicationTargetService.java | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java
index c44b27911bb7a..0c3350f224b11 100644
--- a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java
+++ b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java
@@ -172,6 +172,11 @@ public void run() {
 
     private void start(final long replicationId) {
         try (ReplicationRef<SegmentReplicationTarget> replicationRef = onGoingReplications.get(replicationId)) {
+            // This check is for handling edge cases where the reference is removed before the ReplicationRunner is started by the
+            // threadpool.
+            if (replicationRef == null) {
+                return;
+            }
             replicationRef.get().startReplication(new ActionListener<>() {
                 @Override
                 public void onResponse(Void o) {