Change

Author: Tim-Brooks

server/src/main/java/org/elasticsearch/action/support/replication/ReplicationSplitHelper.java

@@ -62,24 +62,22 @@ public ReplicationSplitHelper(
     }
 
     public static <Request extends ReplicationRequest<Request>> boolean needsSplitCoordination(
-        final TransportReplicationAction.ConcreteShardRequest<Request> primaryRequest,
+        final Request primaryRequest,
         final IndexMetadata indexMetadata
     ) {
-        SplitShardCountSummary requestSplitSummary = primaryRequest.getRequest().reshardSplitShardCountSummary();
+        SplitShardCountSummary requestSplitSummary = primaryRequest.reshardSplitShardCountSummary();
         // TODO: We currently only set the request split summary transport shard bulk. Only evaluate this at the moment or else every
         // request would say it needs a split.
         return requestSplitSummary.isUnset() == false
-            && requestSplitSummary.equals(
-                SplitShardCountSummary.forIndexing(indexMetadata, primaryRequest.getRequest().shardId().getId())
-            ) == false;
+            && requestSplitSummary.equals(SplitShardCountSummary.forIndexing(indexMetadata, primaryRequest.shardId().getId())) == false;
     }
 
     public SplitCoordinator newSplitRequest(
         TransportReplicationAction<Request, ReplicaRequest, Response> action,
         ReplicationTask task,
         ProjectMetadata project,
         TransportReplicationAction<Request, ReplicaRequest, Response>.PrimaryShardReference primaryShardReference,
-        TransportReplicationAction.ConcreteShardRequest<Request> primaryRequest,
+        Request primaryRequest,
         CheckedBiConsumer<
             TransportReplicationAction<Request, ReplicaRequest, Response>.PrimaryShardReference,
             ActionListener<Response>,
@@ -103,7 +101,7 @@ public class SplitCoordinator {
         private final ReplicationTask task;
         private final ProjectMetadata project;
         private final TransportReplicationAction<Request, ReplicaRequest, Response>.PrimaryShardReference primaryShardReference;
-        private final TransportReplicationAction.ConcreteShardRequest<Request> originalRequest;
+        private final Request originalRequest;
         private final CheckedBiConsumer<
             TransportReplicationAction<Request, ReplicaRequest, Response>.PrimaryShardReference,
             ActionListener<Response>,
@@ -115,7 +113,7 @@ public SplitCoordinator(
             ReplicationTask task,
             ProjectMetadata project,
             TransportReplicationAction<Request, ReplicaRequest, Response>.PrimaryShardReference primaryShardReference,
-            TransportReplicationAction.ConcreteShardRequest<Request> originalRequest,
+            Request originalRequest,
             CheckedBiConsumer<
                 TransportReplicationAction<Request, ReplicaRequest, Response>.PrimaryShardReference,
                 ActionListener<Response>,
@@ -132,7 +130,7 @@ public SplitCoordinator(
         }
 
         public void coordinate() throws Exception {
-            Map<ShardId, Request> splitRequests = action.splitRequestOnPrimary(originalRequest.getRequest());
+            Map<ShardId, Request> splitRequests = action.splitRequestOnPrimary(originalRequest);
 
             int numSplitRequests = splitRequests.size();
 
@@ -142,7 +140,8 @@ public void coordinate() throws Exception {
 
             if (numSplitRequests == 1) {
                 // If the request is for source, same behavior as before
-                if (splitRequests.containsKey(originalRequest.getRequest().shardId())) {
+                if (splitRequests.containsKey(originalRequest.shardId())) {
+                    TransportReplicationAction.setPhase(task, "primary");
                     doPrimaryRequest.accept(primaryShardReference, onCompletionListener);
                 } else {
                     // If the request is for target, forward request to target.
@@ -198,11 +197,7 @@ public void onFailure(Exception e) {
                     }
 
                     private void finish() {
-                        Tuple<Response, Exception> finalResponse = action.combineSplitResponses(
-                            originalRequest.getRequest(),
-                            splitRequests,
-                            results
-                        );
+                        Tuple<Response, Exception> finalResponse = action.combineSplitResponses(originalRequest, splitRequests, results);
                         TransportReplicationAction.setPhase(task, "finished");
                         if (finalResponse.v1() != null) {
                             onCompletionListener.onResponse(finalResponse.v1());
@@ -211,7 +206,7 @@ private void finish() {
                         }
                     }
                 };
-                if (splitRequest.getKey().equals(originalRequest.getRequest().shardId())) {
+                if (splitRequest.getKey().equals(originalRequest.shardId())) {
                     doPrimaryRequest.accept(primaryShardReference, listener);
                 } else {
                     delegateToTarget(splitRequest.getKey(), splitRequest.getValue(), clusterService::state, project, listener);

server/src/main/java/org/elasticsearch/action/support/replication/TransportReplicationAction.java

@@ -492,6 +492,10 @@ protected void doRun() throws Exception {
         }
 
         void runWithPrimaryShardReference(final PrimaryShardReference primaryShardReference) {
+            ActionListener<Response> setFinishedListener = ActionListener.runBefore(
+                onCompletionListener,
+                () -> setPhase(replicationTask, "finished")
+            );
             try {
                 final ClusterState clusterState = clusterService.state();
                 final Index index = primaryShardReference.routingEntry().index();
@@ -521,51 +525,37 @@ void runWithPrimaryShardReference(final PrimaryShardReference primaryShardRefere
                         new ConcreteShardRequest<>(primaryRequest.getRequest(), allocationID, primaryRequest.getPrimaryTerm()),
                         transportOptions,
                         new ActionListenerResponseHandler<>(
-                            onCompletionListener,
+                            setFinishedListener,
                             TransportReplicationAction.this::newResponseInstance,
                             TransportResponseHandler.TRANSPORT_WORKER
-                        ) {
-
-                            @Override
-                            public void handleResponse(Response response) {
-                                setPhase(replicationTask, "finished");
-                                super.handleResponse(response);
-                            }
-
-                            @Override
-                            public void handleException(TransportException exp) {
-                                setPhase(replicationTask, "finished");
-                                super.handleException(exp);
-                            }
-                        }
+                        )
                     );
-                } else if (ReplicationSplitHelper.needsSplitCoordination(primaryRequest, indexMetadata)) {
+                } else if (ReplicationSplitHelper.needsSplitCoordination(primaryRequest.getRequest(), indexMetadata)) {
                     ReplicationSplitHelper<Request, ReplicaRequest, Response>.SplitCoordinator splitCoordinator = splitHelper
                         .newSplitRequest(
                             TransportReplicationAction.this,
                             replicationTask,
                             project,
                             primaryShardReference,
-                            primaryRequest,
+                            primaryRequest.getRequest(),
                             this::executePrimaryRequest,
-                            onCompletionListener
+                            setFinishedListener
                         );
                     splitCoordinator.coordinate();
                 } else {
-                    executePrimaryRequest(primaryShardReference, onCompletionListener);
+                    setPhase(replicationTask, "primary");
+                    executePrimaryRequest(primaryShardReference, setFinishedListener);
                 }
             } catch (Exception e) {
                 Releasables.closeWhileHandlingException(primaryShardReference);
-                onFailure(e);
+                setFinishedListener.onFailure(e);
             }
         }
 
         private void executePrimaryRequest(
             final TransportReplicationAction<Request, ReplicaRequest, Response>.PrimaryShardReference primaryShardReference,
             final ActionListener<Response> listener
         ) throws Exception {
-            setPhase(replicationTask, "primary");
-
             final ActionListener<Response> responseListener = ActionListener.wrap(response -> {
                 adaptResponse(response, primaryShardReference.indexShard);
 
@@ -589,9 +579,11 @@ private void executePrimaryRequest(
 
                 assert primaryShardReference.indexShard.isPrimaryMode();
                 primaryShardReference.close(); // release shard operation lock before responding to caller
-                setPhase(replicationTask, "finished");
                 listener.onResponse(response);
-            }, e -> handleException(primaryShardReference, e));
+            }, e -> {
+                Releasables.closeWhileHandlingException(primaryShardReference);
+                listener.onFailure(e);
+            });
 
             new ReplicationOperation<>(
                 primaryRequest.getRequest(),
@@ -607,11 +599,6 @@ private void executePrimaryRequest(
             ).execute();
         }
 
-        private void handleException(PrimaryShardReference primaryShardReference, Exception e) {
-            Releasables.closeWhileHandlingException(primaryShardReference); // release shard operation lock before responding to caller
-            onFailure(e);
-        }
-
         @Override
         public void onFailure(Exception e) {
             setPhase(replicationTask, "finished");