From e96cc08e93add69097a6ddd9d1f5b38a3c59f42a Mon Sep 17 00:00:00 2001
From: Suraj Singh <surajrider@gmail.com>
Date: Wed, 15 Mar 2023 15:10:54 -0700
Subject: [PATCH 1/2] [Segment Replication] Override segment replication
 handler for new request from same replica

Signed-off-by: Suraj Singh <surajrider@gmail.com>
---
 .../replication/OngoingSegmentReplications.java   | 15 +++++++--------
 .../OngoingSegmentReplicationsTests.java          |  3 ++-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java b/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java
index 60078c082a7e3..99d32b20b17ef 100644
--- a/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java
+++ b/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java
@@ -147,16 +147,15 @@ void startSegmentCopy(GetSegmentFilesRequest request, ActionListener<GetSegmentF
      */
     CopyState prepareForReplication(CheckpointInfoRequest request, FileChunkWriter fileChunkWriter) throws IOException {
         final CopyState copyState = getCachedCopyState(request.getCheckpoint());
-        if (allocationIdToHandlers.putIfAbsent(
+        // Clear (if any) existing handler
+        if (allocationIdToHandlers.containsKey(request.getTargetAllocationId())) {
+            logger.warn("Override handler for allocation id {}", request.getTargetAllocationId());
+            cancel(request.getTargetAllocationId(), "cancel due to retry");
+        }
+        allocationIdToHandlers.putIfAbsent(
             request.getTargetAllocationId(),
             createTargetHandler(request.getTargetNode(), copyState, request.getTargetAllocationId(), fileChunkWriter)
-        ) != null) {
-            throw new OpenSearchException(
-                "Shard copy {} on node {} already replicating",
-                request.getCheckpoint().getShardId(),
-                request.getTargetNode()
-            );
-        }
+        );
         return copyState;
     }
 
diff --git a/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java b/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java
index 838dabef1cf75..45cc607943e3f 100644
--- a/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java
+++ b/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java
@@ -277,7 +277,8 @@ public void testShardAlreadyReplicatingToNode() throws IOException {
             listener.onResponse(null);
         };
         replications.prepareForReplication(request, segmentSegmentFileChunkWriter);
-        assertThrows(OpenSearchException.class, () -> { replications.prepareForReplication(request, segmentSegmentFileChunkWriter); });
+        CopyState copyState = replications.prepareForReplication(request, segmentSegmentFileChunkWriter);
+        assertEquals(1, copyState.refCount());
     }
 
     public void testStartReplicationWithNoFilesToFetch() throws IOException {

From dcc837a0acd0298c32f923d42fd44f6f36b311de Mon Sep 17 00:00:00 2001
From: Suraj Singh <surajrider@gmail.com>
Date: Wed, 15 Mar 2023 16:41:00 -0700
Subject: [PATCH 2/2] Spotless and use map.compute

Signed-off-by: Suraj Singh <surajrider@gmail.com>
---
 .../replication/OngoingSegmentReplications.java  | 16 +++++++---------
 .../OngoingSegmentReplicationsTests.java         |  1 -
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java b/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java
index 99d32b20b17ef..bf9e4607b695a 100644
--- a/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java
+++ b/server/src/main/java/org/opensearch/indices/replication/OngoingSegmentReplications.java
@@ -147,15 +147,13 @@ void startSegmentCopy(GetSegmentFilesRequest request, ActionListener<GetSegmentF
      */
     CopyState prepareForReplication(CheckpointInfoRequest request, FileChunkWriter fileChunkWriter) throws IOException {
         final CopyState copyState = getCachedCopyState(request.getCheckpoint());
-        // Clear (if any) existing handler
-        if (allocationIdToHandlers.containsKey(request.getTargetAllocationId())) {
-            logger.warn("Override handler for allocation id {}", request.getTargetAllocationId());
-            cancel(request.getTargetAllocationId(), "cancel due to retry");
-        }
-        allocationIdToHandlers.putIfAbsent(
-            request.getTargetAllocationId(),
-            createTargetHandler(request.getTargetNode(), copyState, request.getTargetAllocationId(), fileChunkWriter)
-        );
+        allocationIdToHandlers.compute(request.getTargetAllocationId(), (allocationId, segrepHandler) -> {
+            if (segrepHandler != null) {
+                logger.warn("Override handler for allocation id {}", request.getTargetAllocationId());
+                cancelHandlers(handler -> handler.getAllocationId().equals(request.getTargetAllocationId()), "cancel due to retry");
+            }
+            return createTargetHandler(request.getTargetNode(), copyState, request.getTargetAllocationId(), fileChunkWriter);
+        });
         return copyState;
     }
 
diff --git a/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java b/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java
index 45cc607943e3f..8cb7a5f6d8929 100644
--- a/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java
+++ b/server/src/test/java/org/opensearch/indices/replication/OngoingSegmentReplicationsTests.java
@@ -9,7 +9,6 @@
 package org.opensearch.indices.replication;
 
 import org.junit.Assert;
-import org.opensearch.OpenSearchException;
 import org.opensearch.action.ActionListener;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.cluster.node.DiscoveryNode;