TransportGetTaskAction: Wait for the task asynchronously

modules/reindex/build.gradle

@@ -184,6 +184,7 @@ tasks.named("yamlRestTestV7CompatTest").configure {
     'update_by_query/80_scripting/Can\'t change _id',
     'update_by_query/80_scripting/Set unsupported operation type',
     'update_by_query/80_scripting/Setting bogus context is an error',
-
+    // `tasks.get` in 7.x doesn't inherit deprecation warnings
+    'reindex/30_search/Sorting deprecated wait_for_completion false'
   ].join(',')
 }

modules/reindex/src/yamlRestTest/resources/rest-api-spec/test/reindex/30_search.yml

@@ -152,6 +152,10 @@
   - set: {task: task}
 
   - do:
+      # Inherits the warning from the previous task
+      warnings:
+        - The sort option in reindex is deprecated. Instead consider using query
+          filtering to find the desired subset of data.
       tasks.get:
         wait_for_completion: true
         task_id: $task

server/src/main/java/org/elasticsearch/action/admin/cluster/node/tasks/get/TransportGetTaskAction.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.action.admin.cluster.node.tasks.get;
 
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.ElasticsearchTimeoutException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.action.ActionListener;
@@ -17,15 +18,17 @@
 import org.elasticsearch.action.get.GetResponse;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
+import org.elasticsearch.action.support.ListenableActionFuture;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.client.internal.OriginSettingClient;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.inject.Inject;
-import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.IndexNotFoundException;
+import org.elasticsearch.tasks.RemovedTaskListener;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.tasks.TaskInfo;
@@ -39,8 +42,9 @@
 
 import java.io.IOException;
 
+import static java.util.Objects.requireNonNullElse;
 import static org.elasticsearch.action.admin.cluster.node.tasks.get.GetTaskAction.TASKS_ORIGIN;
-import static org.elasticsearch.action.admin.cluster.node.tasks.list.TransportListTasksAction.waitForCompletionTimeout;
+import static org.elasticsearch.core.TimeValue.timeValueSeconds;
 
 /**
  * ActionType to get a single task. If the task isn't running then it'll try to request the status from request index.
@@ -53,6 +57,9 @@
  * </ul>
  */
 public class TransportGetTaskAction extends HandledTransportAction<GetTaskRequest, GetTaskResponse> {
+
+    private static final TimeValue DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT = timeValueSeconds(30);
+
     private final ThreadPool threadPool;
     private final ClusterService clusterService;
     private final TransportService transportService;
@@ -130,19 +137,42 @@ void getRunningTaskFromNode(Task thisTask, GetTaskRequest request, ActionListene
             getFinishedTaskFromIndex(thisTask, request, listener);
         } else {
             if (request.getWaitForCompletion()) {
-                // Shift to the generic thread pool and let it wait for the task to complete so we don't block any important threads.
-                threadPool.generic().execute(new AbstractRunnable() {
-                    @Override
-                    protected void doRun() {
-                        taskManager.waitForTaskCompletion(runningTask, waitForCompletionTimeout(request.getTimeout()));
-                        waitedForCompletion(thisTask, request, runningTask.taskInfo(clusterService.localNode().getId(), true), listener);
-                    }
-
-                    @Override
-                    public void onFailure(Exception e) {
-                        listener.onFailure(e);
+                final ListenableActionFuture<Void> future = new ListenableActionFuture<>();
+                RemovedTaskListener removedTaskListener = task -> {
+                    if (task.equals(runningTask)) {
+                        future.onResponse(null);
                     }
-                });
+                };
+                taskManager.registerRemovedTaskListener(removedTaskListener);
+                // Check if the task had finished before we registered the listener,
+                // so we wouldn't wait for a task removal event that would never come
+                if (taskManager.getTask(request.getTaskId().getId()) == null) {
+                    future.onResponse(null);
+                }
+                final ActionListener<Void> waitedForCompletionListener = ActionListener.runBefore(
+                    ActionListener.wrap(
+                        v -> waitedForCompletion(
+                            thisTask,
+                            request,
+                            runningTask.taskInfo(clusterService.localNode().getId(), true),
+                            listener
+                        ),
+                        listener::onFailure
+                    ),
+                    () -> taskManager.unregisterRemovedTaskListener(removedTaskListener)
+                );
+                if (future.isDone()) {
+                    // The task has already finished, we can run the completion listener in the same thread
+                    waitedForCompletionListener.onResponse(null);
+                } else {
+                    future.addListener(waitedForCompletionListener);
+                    var cancellable = threadPool.schedule(
+                        () -> future.onFailure(new ElasticsearchTimeoutException("Timed out waiting for completion of task")),
+                        requireNonNullElse(request.getTimeout(), DEFAULT_WAIT_FOR_COMPLETION_TIMEOUT),
+                        ThreadPool.Names.SAME
+                    );
+                    future.addListener(ActionListener.wrap(cancellable::cancel));
+                }
             } else {
                 TaskInfo info = runningTask.taskInfo(clusterService.localNode().getId(), true);
                 listener.onResponse(new GetTaskResponse(new TaskResult(false, info)));

server/src/main/java/org/elasticsearch/tasks/TaskManager.java

@@ -12,8 +12,6 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.Assertions;
-import org.elasticsearch.ElasticsearchException;
-import org.elasticsearch.ElasticsearchTimeoutException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
@@ -547,23 +545,6 @@ public void applyClusterState(ClusterChangedEvent event) {
         lastDiscoveryNodes = event.state().getNodes();
     }
 
-    /**
-     * Blocks the calling thread, waiting for the task to vanish from the TaskManager.
-     */
-    public void waitForTaskCompletion(Task task, long untilInNanos) {
-        while (System.nanoTime() - untilInNanos < 0) {
-            if (getTask(task.getId()) == null) {
-                return;
-            }
-            try {
-                Thread.sleep(WAIT_FOR_COMPLETION_POLL.millis());
-            } catch (InterruptedException e) {
-                throw new ElasticsearchException("Interrupted waiting for completion of [{}]", e, task);
-            }
-        }
-        throw new ElasticsearchTimeoutException("Timed out waiting for completion of [{}]", task);
-    }
-
     private static class CancellableTaskHolder {
         private final CancellableTask task;
         private boolean finished = false;

test/framework/src/main/java/org/elasticsearch/test/tasks/MockTaskManager.java

@@ -78,21 +78,6 @@ public Task unregister(Task task) {
         return removedTask;
     }
 
-    @Override
-    public void waitForTaskCompletion(Task task, long untilInNanos) {
-        for (MockTaskManagerListener listener : listeners) {
-            try {
-                listener.waitForTaskCompletion(task);
-            } catch (Exception e) {
-                logger.warn(
-                    () -> format("failed to notify task manager listener about waitForTaskCompletion the task with id %s", task.getId()),
-                    e
-                );
-            }
-        }
-        super.waitForTaskCompletion(task, untilInNanos);
-    }
-
     @Override
     public void registerRemovedTaskListener(RemovedTaskListener removedTaskListener) {
         for (MockTaskManagerListener listener : listeners) {