Reindex rethrottle persistent task

This adds support for rethrottling resilient/persistent reindex
through updating the .reindex index and notifying the task. This
ensures that the new throttle value sticks on failovers while
also ensuring that the task wakes up immediately if it had a very
low throttle value.

Related to elastic#42612

Author: henningandersen

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

@@ -68,7 +68,7 @@ static <Request extends AbstractBulkByScrollRequest<Request>> void startSlicedAc
         Client client,
         DiscoveryNode node,
         Runnable workerAction) {
-        initTaskState(task, request, client, new ActionListener<>() {
+        initTaskState(task, request, request.getRequestsPerSecond(), client, new ActionListener<>() {
             @Override
             public void onResponse(Void aVoid) {
                 executeSlicedAction(task, request, action, listener, client, node, workerAction);
@@ -119,6 +119,7 @@ static <Request extends AbstractBulkByScrollRequest<Request>> void executeSliced
     static <Request extends AbstractBulkByScrollRequest<Request>> void initTaskState(
         BulkByScrollTask task,
         Request request,
+        float requestsPerSecond,
         Client client,
         ActionListener<Void> listener) {
         int configuredSlices = request.getSlices();
@@ -128,7 +129,7 @@ static <Request extends AbstractBulkByScrollRequest<Request>> void initTaskState
             client.admin().cluster().searchShards(shardsRequest, new ActionListener<>() {
                 @Override
                 public void onResponse(ClusterSearchShardsResponse response) {
-                    setWorkerCount(request, task, countSlicesBasedOnShards(response));
+                    setWorkerCount(request, requestsPerSecond, task, countSlicesBasedOnShards(response));
                     listener.onResponse(null);
                 }
 
@@ -138,21 +139,22 @@ public void onFailure(Exception e) {
                 }
             });
         } else {
-            setWorkerCount(request, task, configuredSlices);
+            setWorkerCount(request, requestsPerSecond, task, configuredSlices);
             listener.onResponse(null);
         }
     }
 
     private static <Request extends AbstractBulkByScrollRequest<Request>> void setWorkerCount(
         Request request,
+        float requestsPerSecond,
         BulkByScrollTask task,
         int slices) {
         if (slices > 1) {
             task.setWorkerCount(slices);
         } else {
             SliceBuilder sliceBuilder = request.getSearchRequest().source().slice();
             Integer sliceId = sliceBuilder == null ? null : sliceBuilder.getId();
-            task.setWorker(request.getRequestsPerSecond(), sliceId);
+            task.setWorker(requestsPerSecond, sliceId);
         }
     }
 

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

@@ -72,7 +72,8 @@ public class ReindexPlugin extends Plugin implements ActionPlugin, PersistentTas
                 new ActionHandler<>(UpdateByQueryAction.INSTANCE, TransportUpdateByQueryAction.class),
                 new ActionHandler<>(DeleteByQueryAction.INSTANCE, TransportDeleteByQueryAction.class),
                 new ActionHandler<>(RethrottleAction.INSTANCE, TransportRethrottleAction.class),
-                new ActionHandler<>(StartReindexTaskAction.INSTANCE, TransportStartReindexTaskAction.class)
+                new ActionHandler<>(StartReindexTaskAction.INSTANCE, TransportStartReindexTaskAction.class),
+                new ActionHandler<>(RethrottlePersistentReindexAction.INSTANCE, TransportRethrottlePersistentReindexAction.class)
             );
     }
 

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

@@ -36,6 +36,7 @@
 import org.elasticsearch.script.ScriptService;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.threadpool.ThreadPool;
 
@@ -138,11 +139,25 @@ BulkByScrollTask getChildTask() {
         return childTask;
     }
 
+    private void rethrottle(float requestsPerSecond) {
+        TransportRethrottleAction.rethrottle(logger, client.getLocalNodeId(), client, childTask, requestsPerSecond,
+            new ActionListener<>() {
+                @Override
+                public void onResponse(TaskInfo taskInfo) {
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    assert false : e;
+                    logger.error("Unable to rethrottle [{}]", getPersistentTaskId());
+                }
+            });
+    }
+
     private void execute(ReindexTaskParams reindexTaskParams) {
         long allocationId = getAllocationId();
-
         ReindexTaskStateUpdater taskUpdater = new ReindexTaskStateUpdater(reindexIndexClient, client.threadPool(), getPersistentTaskId(),
-            allocationId, new ActionListener<>() {
+            allocationId, taskId, new ActionListener<>() {
             @Override
             public void onResponse(ReindexTaskStateDoc stateDoc) {
                 reindexDone(stateDoc, reindexTaskParams.shouldStoreResult());
@@ -153,14 +168,14 @@ public void onFailure(Exception e) {
                 logger.info("Reindex task failed", e);
                 updateClusterStateToFailed(reindexTaskParams.shouldStoreResult(), ReindexPersistentTaskState.Status.DONE, e);
             }
-        }, this::handleCheckpointAssignmentConflict);
+        }, this::handleCheckpointAssignmentConflict, this::rethrottle);
 
         taskUpdater.assign(new ActionListener<>() {
             @Override
             public void onResponse(ReindexTaskStateDoc stateDoc) {
                 ReindexRequest reindexRequest = stateDoc.getReindexRequest();
                 description = reindexRequest.getDescription();
-                reindexer.initTask(childTask, reindexRequest, new ActionListener<>() {
+                reindexer.initTask(childTask, reindexRequest, stateDoc.getRequestsPerSecond(), new ActionListener<>() {
                     @Override
                     public void onResponse(Void aVoid) {
                         // TODO: need to store status in state so we can continue from it.

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

@@ -27,63 +27,81 @@
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.tasks.TaskId;
 
 import java.io.IOException;
 
 public class ReindexTaskStateDoc implements ToXContentObject {
 
     public static final ConstructingObjectParser<ReindexTaskStateDoc, Void> PARSER =
         new ConstructingObjectParser<>("reindex/index_state", a -> new ReindexTaskStateDoc((ReindexRequest) a[0], (Long) a[1],
-            (BulkByScrollResponse) a[2], (ElasticsearchException) a[3], (Integer) a[4], (ScrollableHitSource.Checkpoint) a[5]));
+            toTaskId((String) a[2]), (BulkByScrollResponse) a[3], (ElasticsearchException) a[4], (Integer) a[5],
+            (ScrollableHitSource.Checkpoint) a[6],
+            (float) a[7]));
 
     private static final String REINDEX_REQUEST = "request";
     private static final String ALLOCATION = "allocation";
+    private static final String EPHEMERAL_TASK_ID = "ephemeral_task_id";
     private static final String REINDEX_RESPONSE = "response";
     private static final String REINDEX_EXCEPTION = "exception";
     private static final String FAILURE_REST_STATUS = "failure_rest_status";
     private static final String REINDEX_CHECKPOINT = "checkpoint";
+    private static final String REQUESTS_PER_SECOND = "requests_per_second";
 
     static {
         PARSER.declareObject(ConstructingObjectParser.constructorArg(), (p, c) -> ReindexRequest.fromXContentWithParams(p),
             new ParseField(REINDEX_REQUEST));
         PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), new ParseField(ALLOCATION));
+        PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), new ParseField(EPHEMERAL_TASK_ID));
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), (p, c) -> BulkByScrollResponse.fromXContent(p),
             new ParseField(REINDEX_RESPONSE));
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), (p, c) -> ElasticsearchException.fromXContent(p),
             new ParseField(REINDEX_EXCEPTION));
         PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), new ParseField(FAILURE_REST_STATUS));
         PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), (p, c) -> ScrollableHitSource.Checkpoint.fromXContent(p),
             new ParseField(REINDEX_CHECKPOINT));
+        PARSER.declareFloat(ConstructingObjectParser.constructorArg(), new ParseField(REQUESTS_PER_SECOND));
     }
 
     private final ReindexRequest reindexRequest;
     private final Long allocationId;
+    private final TaskId ephemeralTaskId;
     private final BulkByScrollResponse reindexResponse;
     private final ElasticsearchException exception;
     private final RestStatus failureStatusCode;
     private final ScrollableHitSource.Checkpoint checkpoint;
+    private final float requestsPerSecond;
 
     public ReindexTaskStateDoc(ReindexRequest reindexRequest) {
-        this(reindexRequest, null, null, null, (RestStatus) null, null);
+        this(reindexRequest, null, null, null, null, (RestStatus) null, null, reindexRequest.getRequestsPerSecond());
     }
 
     public ReindexTaskStateDoc(ReindexRequest reindexRequest, @Nullable Long allocationId,
-                               @Nullable BulkByScrollResponse reindexResponse, @Nullable ElasticsearchException exception,
-                               @Nullable Integer failureStatusCode, ScrollableHitSource.Checkpoint checkpoint) {
-        this(reindexRequest, allocationId, reindexResponse, exception,
-            failureStatusCode == null ? null : RestStatus.fromCode(failureStatusCode), checkpoint);
+                               @Nullable TaskId ephemeralTaskId, @Nullable BulkByScrollResponse reindexResponse,
+                               @Nullable ElasticsearchException exception,
+                               @Nullable Integer failureStatusCode, ScrollableHitSource.Checkpoint checkpoint, float requestsPerSecond) {
+        this(reindexRequest, allocationId, ephemeralTaskId, reindexResponse, exception,
+            failureStatusCode == null ? null : RestStatus.fromCode(failureStatusCode), checkpoint, requestsPerSecond);
     }
 
     public ReindexTaskStateDoc(ReindexRequest reindexRequest, @Nullable Long allocationId,
-                               @Nullable BulkByScrollResponse reindexResponse, @Nullable ElasticsearchException exception,
-                               @Nullable ScrollableHitSource.Checkpoint checkpoint) {
-        this(reindexRequest, allocationId, reindexResponse, exception, exception != null ? exception.status() : null, checkpoint);
+                               @Nullable TaskId ephemeralTaskId, @Nullable BulkByScrollResponse reindexResponse,
+                               @Nullable ElasticsearchException exception,
+                               @Nullable ScrollableHitSource.Checkpoint checkpoint, float requestsPerSecond) {
+        this(reindexRequest, allocationId, ephemeralTaskId, reindexResponse, exception, exception != null ? exception.status() : null,
+            checkpoint, requestsPerSecond);
     }
 
     private ReindexTaskStateDoc(ReindexRequest reindexRequest, @Nullable Long allocationId,
-                                @Nullable BulkByScrollResponse reindexResponse, @Nullable ElasticsearchException exception,
-                                @Nullable RestStatus failureStatusCode, @Nullable ScrollableHitSource.Checkpoint checkpoint) {
+                                @Nullable TaskId ephemeralTaskId, @Nullable BulkByScrollResponse reindexResponse,
+                                @Nullable ElasticsearchException exception,
+                                @Nullable RestStatus failureStatusCode, @Nullable ScrollableHitSource.Checkpoint checkpoint,
+                                float requestsPerSecond) {
+        assert (allocationId == null) == (ephemeralTaskId == null);
         this.allocationId = allocationId;
+        this.ephemeralTaskId = ephemeralTaskId;
+        assert Float.isNaN(requestsPerSecond) == false && requestsPerSecond >= 0;
+        this.requestsPerSecond = requestsPerSecond;
         assert (reindexResponse == null) || (exception == null) : "Either response or exception must be null";
         this.reindexRequest = reindexRequest;
         this.reindexResponse = reindexResponse;
@@ -100,6 +118,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (allocationId != null) {
             builder.field(ALLOCATION, allocationId);
         }
+        if (ephemeralTaskId != null) {
+            builder.field(EPHEMERAL_TASK_ID, ephemeralTaskId.toString());
+        }
         if (reindexResponse != null) {
             builder.field(REINDEX_RESPONSE);
             builder.startObject();
@@ -117,13 +138,18 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             builder.field(REINDEX_CHECKPOINT);
             checkpoint.toXContent(builder, params);
         }
+        builder.field(REQUESTS_PER_SECOND, requestsPerSecond);
         return builder.endObject();
     }
 
     public static ReindexTaskStateDoc fromXContent(XContentParser parser) {
         return PARSER.apply(parser, null);
     }
 
+    private static TaskId toTaskId(String s) {
+        return s != null ? new TaskId(s) : null;
+    }
+
     public ReindexRequest getReindexRequest() {
         return reindexRequest;
     }
@@ -148,17 +174,33 @@ public Long getAllocationId() {
         return allocationId;
     }
 
+    public TaskId getEphemeralTaskId() {
+        return ephemeralTaskId;
+    }
+
+    public float getRequestsPerSecond() {
+        return requestsPerSecond;
+    }
+
     public ReindexTaskStateDoc withCheckpoint(ScrollableHitSource.Checkpoint checkpoint, BulkByScrollTask.Status status) {
         // todo: also store and resume from status.
-        return new ReindexTaskStateDoc(reindexRequest, allocationId, reindexResponse, exception, failureStatusCode, checkpoint);
+        return new ReindexTaskStateDoc(reindexRequest, allocationId, ephemeralTaskId, reindexResponse, exception, failureStatusCode,
+            checkpoint, requestsPerSecond);
     }
 
-    public ReindexTaskStateDoc withNewAllocation(long newAllocationId) {
-        return new ReindexTaskStateDoc(reindexRequest, newAllocationId, reindexResponse, exception, failureStatusCode, checkpoint);
+    public ReindexTaskStateDoc withNewAllocation(long newAllocationId, TaskId ephemeralTaskId) {
+        return new ReindexTaskStateDoc(reindexRequest, newAllocationId, ephemeralTaskId, reindexResponse, exception, failureStatusCode,
+            checkpoint, requestsPerSecond);
     }
 
     public ReindexTaskStateDoc withFinishedState(@Nullable BulkByScrollResponse reindexResponse,
                                                  @Nullable ElasticsearchException exception) {
-        return new ReindexTaskStateDoc(reindexRequest, allocationId, reindexResponse, exception, checkpoint);
+        return new ReindexTaskStateDoc(reindexRequest, allocationId, ephemeralTaskId, reindexResponse, exception, checkpoint,
+            requestsPerSecond);
+    }
+
+    public ReindexTaskStateDoc withRequestsPerSecond(float requestsPerSecond) {
+        return new ReindexTaskStateDoc(reindexRequest, allocationId, ephemeralTaskId, reindexResponse, exception, checkpoint,
+            requestsPerSecond);
     }
 }