Reindex task state initialized before reindex (#46043)

Currently the process to execute a reindex process is tightly coupled to
step of initializing the task state. This creates problems when this
process is asynchronous. It is possible that the task state has not been
initialized which prevents follow-up actions such as rethrottle. This
commit separates the task initialization so that it can be executed as a
first step in the persistent reindex process.

Author: Tim-Brooks

modules/reindex/src/main/java/org/elasticsearch/index/reindex/BulkByScrollParallelizationHelper.java

@@ -57,51 +57,102 @@ private BulkByScrollParallelizationHelper() {}
      *
      * If slices are set as {@code "auto"}, this class will resolve that to a specific number based on characteristics of the source
      * indices. A request with {@code "auto"} slices may end up being sliced or unsliced.
+     *
+     * This method is equivalent to calling {@link #initTaskState} followed by {@link #executeSlicedAction}
      */
     static <Request extends AbstractBulkByScrollRequest<Request>> void startSlicedAction(
-            Request request,
-            BulkByScrollTask task,
-            ActionType<BulkByScrollResponse> action,
-            ActionListener<BulkByScrollResponse> listener,
-            Client client,
-            DiscoveryNode node,
-            Runnable workerAction) {
+        Request request,
+        BulkByScrollTask task,
+        ActionType<BulkByScrollResponse> action,
+        ActionListener<BulkByScrollResponse> listener,
+        Client client,
+        DiscoveryNode node,
+        Runnable workerAction) {
+        initTaskState(task, request, client, new ActionListener<Void>() {
+            @Override
+            public void onResponse(Void aVoid) {
+                executeSlicedAction(task, request, action, listener, client, node, workerAction);
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                listener.onFailure(e);
+            }
+        });
+    }
+
+    /**
+     * Takes an action and a {@link BulkByScrollTask} and runs it with regard to whether this task is a
+     * leader or worker.
+     *
+     * If this task is a worker, the worker action in the given {@link Runnable} will be started on the local
+     * node. If the task is a leader (i.e. the number of slices is more than 1), then a subrequest will be
+     * created for each slice and sent.
+     *
+     * This method can only be called after the task state is initialized {@link #initTaskState}.
+     */
+    static <Request extends AbstractBulkByScrollRequest<Request>> void executeSlicedAction(
+        BulkByScrollTask task,
+        Request request,
+        ActionType<BulkByScrollResponse> action,
+        ActionListener<BulkByScrollResponse> listener,
+        Client client,
+        DiscoveryNode node,
+        Runnable workerAction) {
+        if (task.isLeader()) {
+            sendSubRequests(client, action, node.getId(), task, request, listener);
+        } else if (task.isWorker()) {
+            workerAction.run();
+        } else {
+            throw new AssertionError("Task should have been initialized at this point.");
+        }
+    }
 
-        if (request.getSlices() == AbstractBulkByScrollRequest.AUTO_SLICES) {
+    /**
+     * Takes a {@link BulkByScrollTask} and ensures that its initial task state (leader or worker) is set.
+     *
+     * If slices are set as {@code "auto"}, this method will resolve that to a specific number based on
+     * characteristics of the source indices. A request with {@code "auto"} slices may end up being sliced or
+     * unsliced. This method does not execute the action. In order to execute the action see
+     * {@link #executeSlicedAction}
+     */
+    static <Request extends AbstractBulkByScrollRequest<Request>> void initTaskState(
+        BulkByScrollTask task,
+        Request request,
+        Client client,
+        ActionListener<Void> listener) {
+        int configuredSlices = request.getSlices();
+        if (configuredSlices == AbstractBulkByScrollRequest.AUTO_SLICES) {
             ClusterSearchShardsRequest shardsRequest = new ClusterSearchShardsRequest();
             shardsRequest.indices(request.getSearchRequest().indices());
-            client.admin().cluster().searchShards(shardsRequest, ActionListener.wrap(
-                response -> {
-                    int actualNumSlices = countSlicesBasedOnShards(response);
-                    sliceConditionally(request, task, action, listener, client, node, workerAction, actualNumSlices);
-                },
-                listener::onFailure
-            ));
+            client.admin().cluster().searchShards(shardsRequest, new ActionListener<ClusterSearchShardsResponse>() {
+                @Override
+                public void onResponse(ClusterSearchShardsResponse response) {
+                    setWorkerCount(request, task, countSlicesBasedOnShards(response));
+                    listener.onResponse(null);
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    listener.onFailure(e);
+                }
+            });
         } else {
-            sliceConditionally(request, task, action, listener, client, node, workerAction, request.getSlices());
+            setWorkerCount(request, task, configuredSlices);
+            listener.onResponse(null);
         }
     }
 
-    private static <Request extends AbstractBulkByScrollRequest<Request>> void sliceConditionally(
-            Request request,
-            BulkByScrollTask task,
-            ActionType<BulkByScrollResponse> action,
-            ActionListener<BulkByScrollResponse> listener,
-            Client client,
-            DiscoveryNode node,
-            Runnable workerAction,
-            int slices) {
-
+    private static <Request extends AbstractBulkByScrollRequest<Request>> void setWorkerCount(
+        Request request,
+        BulkByScrollTask task,
+        int slices) {
         if (slices > 1) {
             task.setWorkerCount(slices);
-            sendSubRequests(client, action, node.getId(), task, request, listener);
         } else {
             SliceBuilder sliceBuilder = request.getSearchRequest().source().slice();
-            Integer sliceId = sliceBuilder == null
-                ? null
-                : sliceBuilder.getId();
+            Integer sliceId = sliceBuilder == null ? null : sliceBuilder.getId();
             task.setWorker(request.getRequestsPerSecond(), sliceId);
-            workerAction.run();
         }
     }
 

modules/reindex/src/main/java/org/elasticsearch/index/reindex/Reindexer.java

@@ -87,16 +87,19 @@ public class Reindexer {
         this.reindexSslConfig = reindexSslConfig;
     }
 
+    public void initTask(BulkByScrollTask task, ReindexRequest request, ActionListener<Void> listener) {
+        BulkByScrollParallelizationHelper.initTaskState(task, request, client, listener);
+    }
+
     public void execute(BulkByScrollTask task, ReindexRequest request, ActionListener<BulkByScrollResponse> listener) {
-        BulkByScrollParallelizationHelper.startSlicedAction(request, task, ReindexAction.INSTANCE, listener, client,
+        BulkByScrollParallelizationHelper.executeSlicedAction(task, request, ReindexAction.INSTANCE, listener, client,
             clusterService.localNode(),
             () -> {
                 ParentTaskAssigningClient assigningClient = new ParentTaskAssigningClient(client, clusterService.localNode(), task);
                 AsyncIndexBySearchAction searchAction = new AsyncIndexBySearchAction(task, logger, assigningClient, threadPool,
                     scriptService, reindexSslConfig, request, listener);
                 searchAction.start();
-            }
-        );
+            });
 
     }
 

modules/reindex/src/main/java/org/elasticsearch/index/reindex/TransportReindexAction.java

@@ -54,13 +54,22 @@ public TransportReindexAction(Settings settings, ThreadPool threadPool, ActionFi
         super(ReindexAction.NAME, transportService, actionFilters, ReindexRequest::new);
         this.reindexValidator = new ReindexValidator(settings, clusterService, indexNameExpressionResolver, autoCreateIndex);
         this.reindexer = new Reindexer(clusterService, client, threadPool, scriptService, sslConfig);
-
     }
 
     @Override
     protected void doExecute(Task task, ReindexRequest request, ActionListener<BulkByScrollResponse> listener) {
         reindexValidator.initialValidation(request);
         BulkByScrollTask bulkByScrollTask = (BulkByScrollTask) task;
-        reindexer.execute(bulkByScrollTask, request, listener);
+        reindexer.initTask(bulkByScrollTask, request, new ActionListener<Void>() {
+            @Override
+            public void onResponse(Void v) {
+                reindexer.execute(bulkByScrollTask, request, listener);
+            }
+
+            @Override
+            public void onFailure(Exception e) {
+                listener.onFailure(e);
+            }
+        });
     }
 }