Make Transport Shard Bulk Action Async (#39793)

server/src/main/java/org/elasticsearch/action/admin/indices/close/TransportVerifyShardBeforeCloseAction.java

@@ -85,14 +85,16 @@ protected void acquireReplicaOperationPermit(final IndexShard replica,
     }
 
     @Override
-    protected PrimaryResult<ShardRequest, ReplicationResponse> shardOperationOnPrimary(final ShardRequest shardRequest,
-                                                                                       final IndexShard primary) throws Exception {
-        executeShardOperation(shardRequest, primary);
-        return new PrimaryResult<>(shardRequest, new ReplicationResponse());
+    protected void shardOperationOnPrimary(final ShardRequest shardRequest, final IndexShard primary,
+            ActionListener<PrimaryResult<ShardRequest, ReplicationResponse>> listener) {
+        ActionListener.completeWith(listener, () -> {
+            executeShardOperation(shardRequest, primary);
+            return new PrimaryResult<>(shardRequest, new ReplicationResponse());
+        });
     }
 
     @Override
-    protected ReplicaResult shardOperationOnReplica(final ShardRequest shardRequest, final IndexShard replica) throws Exception {
+    protected ReplicaResult shardOperationOnReplica(final ShardRequest shardRequest, final IndexShard replica) {
         executeShardOperation(shardRequest, replica);
         return new ReplicaResult();
     }

server/src/main/java/org/elasticsearch/action/admin/indices/flush/TransportShardFlushAction.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.action.admin.indices.flush;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.replication.ReplicationResponse;
 import org.elasticsearch.action.support.replication.TransportReplicationAction;
@@ -51,11 +52,13 @@ protected ReplicationResponse newResponseInstance() {
     }
 
     @Override
-    protected PrimaryResult<ShardFlushRequest, ReplicationResponse> shardOperationOnPrimary(ShardFlushRequest shardRequest,
-            IndexShard primary) {
-        primary.flush(shardRequest.getRequest());
-        logger.trace("{} flush request executed on primary", primary.shardId());
-        return new PrimaryResult<>(shardRequest, new ReplicationResponse());
+    protected void shardOperationOnPrimary(ShardFlushRequest shardRequest, IndexShard primary,
+            ActionListener<PrimaryResult<ShardFlushRequest, ReplicationResponse>> listener) {
+        ActionListener.completeWith(listener, () -> {
+            primary.flush(shardRequest.getRequest());
+            logger.trace("{} flush request executed on primary", primary.shardId());
+            return new PrimaryResult<>(shardRequest, new ReplicationResponse());
+        });
     }
 
     @Override

server/src/main/java/org/elasticsearch/action/admin/indices/refresh/TransportShardRefreshAction.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.action.admin.indices.refresh;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.replication.BasicReplicationRequest;
 import org.elasticsearch.action.support.replication.ReplicationResponse;
@@ -53,11 +54,13 @@ protected ReplicationResponse newResponseInstance() {
     }
 
     @Override
-    protected PrimaryResult<BasicReplicationRequest, ReplicationResponse> shardOperationOnPrimary(
-            BasicReplicationRequest shardRequest, IndexShard primary) {
-        primary.refresh("api");
-        logger.trace("{} refresh request executed on primary", primary.shardId());
-        return new PrimaryResult<>(shardRequest, new ReplicationResponse());
+    protected void shardOperationOnPrimary(BasicReplicationRequest shardRequest, IndexShard primary,
+            ActionListener<PrimaryResult<BasicReplicationRequest, ReplicationResponse>> listener) {
+        ActionListener.completeWith(listener, () -> {
+            primary.refresh("api");
+            logger.trace("{} refresh request executed on primary", primary.shardId());
+            return new PrimaryResult<>(shardRequest, new ReplicationResponse());
+        });
     }
 
     @Override

server/src/main/java/org/elasticsearch/action/bulk/MappingUpdatePerformer.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.action.bulk;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.index.mapper.Mapping;
 import org.elasticsearch.index.shard.ShardId;
 
@@ -27,6 +28,6 @@ public interface MappingUpdatePerformer {
     /**
      * Update the mappings on the master.
      */
-    void updateMappings(Mapping update, ShardId shardId, String type);
+    void updateMappings(Mapping update, ShardId shardId, String type, ActionListener<Void> listener);
 
 }

server/src/main/java/org/elasticsearch/action/resync/TransportResyncReplicationAction.java

@@ -85,13 +85,13 @@ public ClusterBlockLevel indexBlockLevel() {
     }
 
     @Override
-    protected WritePrimaryResult<ResyncReplicationRequest, ResyncReplicationResponse> shardOperationOnPrimary(
-        ResyncReplicationRequest request, IndexShard primary) throws Exception {
-        final ResyncReplicationRequest replicaRequest = performOnPrimary(request, primary);
-        return new WritePrimaryResult<>(replicaRequest, new ResyncReplicationResponse(), null, null, primary, logger);
+    protected void shardOperationOnPrimary(ResyncReplicationRequest request, IndexShard primary,
+            ActionListener<PrimaryResult<ResyncReplicationRequest, ResyncReplicationResponse>> listener) {
+        ActionListener.completeWith(listener,
+            () -> new WritePrimaryResult<>(performOnPrimary(request), new ResyncReplicationResponse(), null, null, primary, logger));
     }
 
-    public static ResyncReplicationRequest performOnPrimary(ResyncReplicationRequest request, IndexShard primary) {
+    public static ResyncReplicationRequest performOnPrimary(ResyncReplicationRequest request) {
         return request;
     }
 

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

@@ -101,14 +101,17 @@ public void execute() throws Exception {
 
         totalShards.incrementAndGet();
         pendingActions.incrementAndGet(); // increase by 1 until we finish all primary coordination
-        primaryResult = primary.perform(request);
-        primary.updateLocalCheckpointForShard(primaryRouting.allocationId().getId(), primary.localCheckpoint());
+        primary.perform(request, ActionListener.wrap(this::handlePrimaryResult, resultListener::onFailure));
+    }
+
+    private void handlePrimaryResult(final PrimaryResultT primaryResult) {
+        this.primaryResult = primaryResult;
+        primary.updateLocalCheckpointForShard(primary.routingEntry().allocationId().getId(), primary.localCheckpoint());
         final ReplicaRequest replicaRequest = primaryResult.replicaRequest();
         if (replicaRequest != null) {
             if (logger.isTraceEnabled()) {
-                logger.trace("[{}] op [{}] completed on primary for request [{}]", primaryId, opType, request);
+                logger.trace("[{}] op [{}] completed on primary for request [{}]", primary.routingEntry().shardId(), opType, request);
             }
-
             // we have to get the replication group after successfully indexing into the primary in order to honour recovery semantics.
             // we have to make sure that every operation indexed into the primary after recovery start will also be replicated
             // to the recovery target. If we used an old replication group, we may miss a recovery that has started since then.
@@ -118,14 +121,14 @@ public void execute() throws Exception {
             // This would entail that some shards could learn about a global checkpoint that would be higher than its local checkpoint.
             final long globalCheckpoint = primary.globalCheckpoint();
             // we have to capture the max_seq_no_of_updates after this request was completed on the primary to make sure the value of
-            // max_seq_no_of_updates on replica when this request is executed is at least the value on the primary when it was executed on.
+            // max_seq_no_of_updates on replica when this request is executed is at least the value on the primary when it was executed
+            // on.
             final long maxSeqNoOfUpdatesOrDeletes = primary.maxSeqNoOfUpdatesOrDeletes();
             assert maxSeqNoOfUpdatesOrDeletes != SequenceNumbers.UNASSIGNED_SEQ_NO : "seqno_of_updates still uninitialized";
             final ReplicationGroup replicationGroup = primary.getReplicationGroup();
             markUnavailableShardsAsStale(replicaRequest, replicationGroup);
             performOnReplicas(replicaRequest, globalCheckpoint, maxSeqNoOfUpdatesOrDeletes, replicationGroup);
         }
-
         successfulShards.incrementAndGet();  // mark primary as successful
         decPendingAndFinishIfNeeded();
     }
@@ -310,9 +313,9 @@ public interface Primary<
          * also complete after. Deal with it.
          *
          * @param request the request to perform
-         * @return the request to send to the replicas
+         * @param listener result listener
          */
-        PrimaryResultT perform(RequestT request) throws Exception;
+        void perform(RequestT request, ActionListener<PrimaryResultT> listener);
 
         /**
          * Notifies the primary of a local checkpoint for the given allocation.

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

@@ -21,6 +21,7 @@
 
 import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.apache.lucene.store.AlreadyClosedException;
+import org.elasticsearch.Assertions;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.Version;
@@ -190,8 +191,8 @@ protected void resolveRequest(final IndexMetaData indexMetaData, final Request r
      * @param shardRequest the request to the primary shard
      * @param primary      the primary shard to perform the operation on
      */
-    protected abstract PrimaryResult<ReplicaRequest, Response> shardOperationOnPrimary(
-            Request shardRequest, IndexShard primary) throws Exception;
+    protected abstract void shardOperationOnPrimary(Request shardRequest, IndexShard primary,
+        ActionListener<PrimaryResult<ReplicaRequest, Response>> listener);
 
     /**
      * Synchronously execute the specified replica operation. This is done under a permit from
@@ -357,58 +358,54 @@ public void handleException(TransportException exp) {
                         });
                 } else {
                     setPhase(replicationTask, "primary");
-                    final ActionListener<Response> listener = createResponseListener(primaryShardReference);
                     createReplicatedOperation(primaryRequest.getRequest(),
-                            ActionListener.wrap(result -> result.respond(listener), listener::onFailure),
-                            primaryShardReference)
-                            .execute();
+                        ActionListener.wrap(result -> result.respond(
+                            new ActionListener<Response>() {
+                                @Override
+                                public void onResponse(Response response) {
+                                    if (syncGlobalCheckpointAfterOperation) {
+                                        final IndexShard shard = primaryShardReference.indexShard;
+                                        try {
+                                            shard.maybeSyncGlobalCheckpoint("post-operation");
+                                        } catch (final Exception e) {
+                                            // only log non-closed exceptions
+                                            if (ExceptionsHelper.unwrap(
+                                                e, AlreadyClosedException.class, IndexShardClosedException.class) == null) {
+                                                // intentionally swallow, a missed global checkpoint sync should not fail this operation
+                                                logger.info(
+                                                    new ParameterizedMessage(
+                                                        "{} failed to execute post-operation global checkpoint sync", shard.shardId()), e);
+                                            }
+                                        }
+                                    }
+                                    primaryShardReference.close(); // release shard operation lock before responding to caller
+                                    setPhase(replicationTask, "finished");
+                                    onCompletionListener.onResponse(response);
+                                }
+
+                                @Override
+                                public void onFailure(Exception e) {
+                                    handleException(primaryShardReference, e);
+                                }
+                            }), e -> handleException(primaryShardReference, e)
+                        ), primaryShardReference).execute();
                 }
             } catch (Exception e) {
-                Releasables.closeWhileHandlingException(primaryShardReference); // release shard operation lock before responding to caller
-                onFailure(e);
+                handleException(primaryShardReference, e);
             }
         }
 
+        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");
             onCompletionListener.onFailure(e);
         }
 
-        private ActionListener<Response> createResponseListener(final PrimaryShardReference primaryShardReference) {
-            return new ActionListener<Response>() {
-                @Override
-                public void onResponse(Response response) {
-                    if (syncGlobalCheckpointAfterOperation) {
-                        final IndexShard shard = primaryShardReference.indexShard;
-                        try {
-                            shard.maybeSyncGlobalCheckpoint("post-operation");
-                        } catch (final Exception e) {
-                            // only log non-closed exceptions
-                            if (ExceptionsHelper.unwrap(e, AlreadyClosedException.class, IndexShardClosedException.class) == null) {
-                                logger.info(
-                                        new ParameterizedMessage(
-                                                "{} failed to execute post-operation global checkpoint sync",
-                                                shard.shardId()),
-                                        e);
-                                // intentionally swallow, a missed global checkpoint sync should not fail this operation
-                            }
-                        }
-                    }
-                    primaryShardReference.close(); // release shard operation lock before responding to caller
-                    setPhase(replicationTask, "finished");
-                    onCompletionListener.onResponse(response);
-                }
-
-                @Override
-                public void onFailure(Exception e) {
-                    primaryShardReference.close(); // release shard operation lock before responding to caller
-                    setPhase(replicationTask, "finished");
-                    onCompletionListener.onFailure(e);
-                }
-            };
-        }
-
         protected ReplicationOperation<Request, ReplicaRequest, PrimaryResult<ReplicaRequest, Response>> createReplicatedOperation(
             Request request, ActionListener<PrimaryResult<ReplicaRequest, Response>> listener,
             PrimaryShardReference primaryShardReference) {
@@ -417,7 +414,7 @@ protected ReplicationOperation<Request, ReplicaRequest, PrimaryResult<ReplicaReq
         }
     }
 
-    protected static class PrimaryResult<ReplicaRequest extends ReplicationRequest<ReplicaRequest>,
+    public static class PrimaryResult<ReplicaRequest extends ReplicationRequest<ReplicaRequest>,
             Response extends ReplicationResponse>
             implements ReplicationOperation.PrimaryResult<ReplicaRequest> {
         final ReplicaRequest replicaRequest;
@@ -916,11 +913,15 @@ public void failShard(String reason, Exception e) {
         }
 
         @Override
-        public PrimaryResult<ReplicaRequest, Response> perform(Request request) throws Exception {
-            PrimaryResult<ReplicaRequest, Response> result = shardOperationOnPrimary(request, indexShard);
-            assert result.replicaRequest() == null || result.finalFailure == null : "a replica request [" + result.replicaRequest()
-                + "] with a primary failure [" + result.finalFailure + "]";
-            return result;
+        public void perform(Request request, ActionListener<PrimaryResult<ReplicaRequest, Response>> listener) {
+            if (Assertions.ENABLED) {
+                listener = ActionListener.map(listener, result -> {
+                    assert result.replicaRequest() == null || result.finalFailure == null : "a replica request [" + result.replicaRequest()
+                        + "] with a primary failure [" + result.finalFailure + "]";
+                    return result;
+                });
+            }
+            shardOperationOnPrimary(request, indexShard, listener);
         }
 
         @Override

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

@@ -103,12 +103,12 @@ protected ReplicationOperation.Replicas newReplicasProxy(long primaryTerm) {
     /**
      * Called on the primary with a reference to the primary {@linkplain IndexShard} to modify.
      *
-     * @return the result of the operation on primary, including current translog location and operation response and failure
-     * async refresh is performed on the <code>primary</code> shard according to the <code>Request</code> refresh policy
+     * @param listener listener for the result of the operation on primary, including current translog location and operation response
+     * and failure async refresh is performed on the <code>primary</code> shard according to the <code>Request</code> refresh policy
      */
     @Override
-    protected abstract WritePrimaryResult<ReplicaRequest, Response> shardOperationOnPrimary(
-            Request request, IndexShard primary) throws Exception;
+    protected abstract void shardOperationOnPrimary(
+            Request request, IndexShard primary, ActionListener<PrimaryResult<ReplicaRequest, Response>> listener);
 
     /**
      * Called once per replica with a reference to the replica {@linkplain IndexShard} to modify.