Make Transport Shard Bulk Action Async (#39793)

This is a dependency of #39504 

Motivation: 
By refactoring `TransportShardBulkAction#shardOperationOnPrimary` to async, we enable using `DeterministicTaskQueue` based tests to run indexing operations. This was previously impossible since we were blocking on the `write` thread until the `update` thread finished the mapping update.
With this change, the mapping update will trigger a new task in the `write` queue instead. 
This change significantly enhances the amount of coverage we get from `SnapshotResiliencyTests` (and other potential future tests) when it comes to tracking down concurrency issues with distributed state machines.

The logical change is effectively all in `TransportShardBulkAction`, the rest of the changes is then simply mechanically moving the caller code and tests to being async and passing the `ActionListener` down.

Since the move to async would've added more parameters to the `private static` steps in this logic, I decided to inline and dry up (between delete and update) the logic as much as I could instead of passing the listener + wait-consumer down through all of them.

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,10 +85,10 @@ public ClusterBlockLevel indexBlockLevel() {
     }
 
     @Override
-    protected WritePrimaryResult<ResyncReplicationRequest, ResyncReplicationResponse> shardOperationOnPrimary(
-        ResyncReplicationRequest request, IndexShard primary) {
-        final ResyncReplicationRequest replicaRequest = performOnPrimary(request);
-        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) {

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.action.ActionListener;
@@ -189,8 +190,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
@@ -416,7 +417,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;
@@ -915,11 +916,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.

server/src/main/java/org/elasticsearch/cluster/action/index/MappingUpdatedAction.java

@@ -19,7 +19,9 @@
 
 package org.elasticsearch.cluster.action.index;
 
-import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequestBuilder;
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.action.support.master.MasterNodeRequest;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.client.IndicesAdminClient;
@@ -29,6 +31,7 @@
 import org.elasticsearch.common.settings.Setting.Property;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.common.util.concurrent.FutureUtils;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.mapper.MapperService;
@@ -57,34 +60,36 @@ private void setDynamicMappingUpdateTimeout(TimeValue dynamicMappingUpdateTimeou
         this.dynamicMappingUpdateTimeout = dynamicMappingUpdateTimeout;
     }
 
-
     public void setClient(Client client) {
         this.client = client.admin().indices();
     }
 
-    private PutMappingRequestBuilder updateMappingRequest(Index index, String type, Mapping mappingUpdate, final TimeValue timeout) {
-        if (type.equals(MapperService.DEFAULT_MAPPING)) {
-            throw new IllegalArgumentException("_default_ mapping should not be updated");
-        }
-        return client.preparePutMapping().setConcreteIndex(index).setType(type).setSource(mappingUpdate.toString(), XContentType.JSON)
-                .setMasterNodeTimeout(timeout).setTimeout(TimeValue.ZERO);
-    }
-
-    /**
-     * Same as {@link #updateMappingOnMaster(Index, String, Mapping, TimeValue)}
-     * using the default timeout.
-     */
-    public void updateMappingOnMaster(Index index, String type, Mapping mappingUpdate) {
-        updateMappingOnMaster(index, type, mappingUpdate, dynamicMappingUpdateTimeout);
-    }
-
     /**
      * Update mappings on the master node, waiting for the change to be committed,
      * but not for the mapping update to be applied on all nodes. The timeout specified by
      * {@code timeout} is the master node timeout ({@link MasterNodeRequest#masterNodeTimeout()}),
      * potentially waiting for a master node to be available.
      */
-    public void updateMappingOnMaster(Index index, String type, Mapping mappingUpdate, TimeValue masterNodeTimeout) {
-        updateMappingRequest(index, type, mappingUpdate, masterNodeTimeout).get();
+    public void updateMappingOnMaster(Index index, String type, Mapping mappingUpdate, ActionListener<Void> listener) {
+        if (type.equals(MapperService.DEFAULT_MAPPING)) {
+            throw new IllegalArgumentException("_default_ mapping should not be updated");
+        }
+        client.preparePutMapping().setConcreteIndex(index).setType(type).setSource(mappingUpdate.toString(), XContentType.JSON)
+            .setMasterNodeTimeout(dynamicMappingUpdateTimeout).setTimeout(TimeValue.ZERO)
+            .execute(new ActionListener<AcknowledgedResponse>() {
+                @Override
+                public void onResponse(AcknowledgedResponse acknowledgedResponse) {
+                    listener.onResponse(null);
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    listener.onFailure(unwrapException(e));
+                }
+            });
+    }
+
+    private static Exception unwrapException(Exception cause) {
+        return cause instanceof ElasticsearchException ? FutureUtils.unwrapEsException((ElasticsearchException) cause) : cause;
     }
 }

server/src/main/java/org/elasticsearch/common/util/concurrent/FutureUtils.java

@@ -88,17 +88,19 @@ public static  <T> T get(Future<T> future, long timeout, TimeUnit unit) {
     public static RuntimeException rethrowExecutionException(ExecutionException e) {
         if (e.getCause() instanceof ElasticsearchException) {
             ElasticsearchException esEx = (ElasticsearchException) e.getCause();
-            Throwable root = esEx.unwrapCause();
-            if (root instanceof ElasticsearchException) {
-                return (ElasticsearchException) root;
-            } else if (root instanceof RuntimeException) {
-                return (RuntimeException) root;
-            }
-            return new UncategorizedExecutionException("Failed execution", root);
+            return unwrapEsException(esEx);
         } else if (e.getCause() instanceof RuntimeException) {
             return (RuntimeException) e.getCause();
         } else {
             return new UncategorizedExecutionException("Failed execution", e);
         }
     }
+
+    public static RuntimeException unwrapEsException(ElasticsearchException esEx) {
+        Throwable root = esEx.unwrapCause();
+        if (root instanceof ElasticsearchException || root instanceof RuntimeException) {
+            return (RuntimeException) root;
+        }
+        return new UncategorizedExecutionException("Failed execution", root);
+    }
 }

server/src/main/java/org/elasticsearch/index/seqno/GlobalCheckpointSyncAction.java

@@ -102,10 +102,12 @@ protected ReplicationResponse newResponseInstance() {
     }
 
     @Override
-    protected PrimaryResult<Request, ReplicationResponse> shardOperationOnPrimary(
-            final Request request, final IndexShard indexShard) throws Exception {
-        maybeSyncTranslog(indexShard);
-        return new PrimaryResult<>(request, new ReplicationResponse());
+    protected void shardOperationOnPrimary(Request request, IndexShard indexShard,
+                                           ActionListener<PrimaryResult<Request, ReplicationResponse>> listener) {
+        ActionListener.completeWith(listener, () -> {
+            maybeSyncTranslog(indexShard);
+            return new PrimaryResult<>(request, new ReplicationResponse());
+        });
     }
 
     @Override

server/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseBackgroundSyncAction.java

@@ -127,14 +127,16 @@ public void backgroundSync(
     }
 
     @Override
-    protected PrimaryResult<Request, ReplicationResponse> shardOperationOnPrimary(
+    protected void shardOperationOnPrimary(
             final Request request,
-            final IndexShard primary) throws WriteStateException {
-        assert request.waitForActiveShards().equals(ActiveShardCount.NONE) : request.waitForActiveShards();
-        Objects.requireNonNull(request);
-        Objects.requireNonNull(primary);
-        primary.persistRetentionLeases();
-        return new PrimaryResult<>(request, new ReplicationResponse());
+            final IndexShard primary, ActionListener<PrimaryResult<Request, ReplicationResponse>> listener) {
+        ActionListener.completeWith(listener, () -> {
+            assert request.waitForActiveShards().equals(ActiveShardCount.NONE) : request.waitForActiveShards();
+            Objects.requireNonNull(request);
+            Objects.requireNonNull(primary);
+            primary.persistRetentionLeases();
+            return new PrimaryResult<>(request, new ReplicationResponse());
+        });
     }
 
     @Override