Add support for task cancellation to TransportMasterNodeAction (#72157)

Author: fcofdez

server/src/main/java/org/elasticsearch/action/support/master/TransportMasterNodeAction.java

@@ -31,13 +31,15 @@
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.discovery.MasterNotDiscoveredException;
 import org.elasticsearch.node.NodeClosedException;
+import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.ConnectTransportException;
 import org.elasticsearch.transport.RemoteTransportException;
 import org.elasticsearch.transport.TransportException;
 import org.elasticsearch.transport.TransportService;
 
+import java.util.concurrent.CancellationException;
 import java.util.function.Predicate;
 
 /**
@@ -82,6 +84,15 @@ protected TransportMasterNodeAction(String actionName, boolean canTripCircuitBre
     protected abstract void masterOperation(Task task, Request request, ClusterState state,
                                             ActionListener<Response> listener) throws Exception;
 
+    private void executeMasterOperation(Task task, Request request, ClusterState state,
+                                        ActionListener<Response> listener) throws Exception {
+        if (task instanceof CancellableTask && ((CancellableTask) task).isCancelled()) {
+            throw new CancellationException("Task was cancelled");
+        }
+
+        masterOperation(task, request, state, listener);
+    }
+
     protected boolean localExecute(Request request) {
         return false;
     }
@@ -114,6 +125,10 @@ class AsyncSingleAction {
         }
 
         protected void doStart(ClusterState clusterState) {
+            if (isTaskCancelled()) {
+                listener.onFailure(new CancellationException("Task was cancelled"));
+                return;
+            }
             try {
                 final DiscoveryNodes nodes = clusterState.nodes();
                 if (nodes.isLocalNodeElectedMaster() || localExecute(request)) {
@@ -148,7 +163,7 @@ protected void doStart(ClusterState clusterState) {
                             }
                         });
                         threadPool.executor(executor)
-                            .execute(ActionRunnable.wrap(delegate, l -> masterOperation(task, request, clusterState, l)));
+                            .execute(ActionRunnable.wrap(delegate, l -> executeMasterOperation(task, request, clusterState, l)));
                     }
                 } else {
                     if (nodes.getMasterNode() == null) {
@@ -218,7 +233,11 @@ public void onTimeout(TimeValue timeout) {
                             actionName, timeout), failure);
                         listener.onFailure(new MasterNotDiscoveredException(failure));
                     }
-                }, statePredicate);
+                }, clusterState -> isTaskCancelled() || statePredicate.test(clusterState));
+        }
+
+        private boolean isTaskCancelled() {
+            return task instanceof CancellableTask && ((CancellableTask) task).isCancelled();
         }
     }
 }

server/src/test/java/org/elasticsearch/action/support/master/TransportMasterNodeActionTests.java

@@ -31,12 +31,17 @@
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.common.util.concurrent.EsThreadPoolExecutor;
 import org.elasticsearch.discovery.MasterNotDiscoveredException;
 import org.elasticsearch.indices.TestIndexNameExpressionResolver;
 import org.elasticsearch.node.NodeClosedException;
 import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.transport.CapturingTransport;
 import org.elasticsearch.threadpool.TestThreadPool;
@@ -51,8 +56,12 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.HashSet;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 
@@ -126,6 +135,11 @@ public static class Request extends MasterNodeRequest<Request> {
         public ActionRequestValidationException validate() {
             return null;
         }
+
+        @Override
+        public Task createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
+            return new CancellableTask(id, type, action, "", parentTaskId, headers);
+        }
     }
 
     class Response extends ActionResponse {
@@ -160,12 +174,18 @@ public void writeTo(StreamOutput out) throws IOException {
     class Action extends TransportMasterNodeAction<Request, Response> {
         Action(String actionName, TransportService transportService, ClusterService clusterService,
                ThreadPool threadPool) {
+            this(actionName, transportService, clusterService, threadPool, ThreadPool.Names.SAME);
+        }
+
+        Action(String actionName, TransportService transportService, ClusterService clusterService,
+               ThreadPool threadPool, String executor) {
             super(actionName, transportService, clusterService, threadPool,
                 new ActionFilters(new HashSet<>()), Request::new,
                 TestIndexNameExpressionResolver.newInstance(), Response::new,
-                ThreadPool.Names.SAME);
+                executor);
         }
 
+
         @Override
         protected void doExecute(Task task, final Request request, ActionListener<Response> listener) {
             // remove unneeded threading by wrapping listener with SAME to prevent super.doExecute from wrapping it with LISTENER
@@ -460,4 +480,110 @@ protected void masterOperation(Task task, Request request, ClusterState state,
         assertTrue(listener.isDone());
         assertThat(listener.get(), equalTo(response));
     }
+
+    public void testTaskCancellation() {
+        ClusterBlock block = new ClusterBlock(1,
+            "",
+            true,
+            true,
+            false,
+            randomFrom(RestStatus.values()),
+            ClusterBlockLevel.ALL
+        );
+        ClusterState stateWithBlock = ClusterState.builder(ClusterStateCreationUtils.state(localNode, localNode, allNodes))
+            .blocks(ClusterBlocks.builder().addGlobalBlock(block)).build();
+
+        // Update the cluster state with a block so the request waits until it's unblocked
+        setState(clusterService, stateWithBlock);
+
+        TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet());
+
+        Request request = new Request();
+        final CancellableTask task = (CancellableTask) taskManager.register("type", "internal:testAction", request);
+
+        boolean cancelBeforeStart = randomBoolean();
+        if (cancelBeforeStart) {
+            taskManager.cancel(task, "", () -> {});
+            assertThat(task.isCancelled(), equalTo(true));
+        }
+
+        PlainActionFuture<Response> listener = new PlainActionFuture<>();
+        ActionTestUtils.execute(new Action("internal:testAction", transportService, clusterService, threadPool) {
+            @Override
+            protected ClusterBlockException checkBlock(Request request, ClusterState state) {
+                Set<ClusterBlock> blocks = state.blocks().global();
+                return blocks.isEmpty() ? null : new ClusterBlockException(blocks);
+            }
+        }, task, request, listener);
+
+        final int genericThreads = threadPool.info(ThreadPool.Names.GENERIC).getMax();
+        final EsThreadPoolExecutor executor = (EsThreadPoolExecutor) threadPool.executor(ThreadPool.Names.GENERIC);
+        final CyclicBarrier barrier = new CyclicBarrier(genericThreads + 1);
+        final CountDownLatch latch = new CountDownLatch(1);
+
+        if (cancelBeforeStart == false) {
+            assertThat(listener.isDone(), equalTo(false));
+
+            taskManager.cancel(task, "", () -> {});
+            assertThat(task.isCancelled(), equalTo(true));
+
+            // Now that the task is cancelled, let the request to be executed
+            final ClusterState.Builder newStateBuilder = ClusterState.builder(stateWithBlock);
+
+            // Either unblock the cluster state or just do an unrelated cluster state change that will check
+            // if the task has been cancelled
+            if (randomBoolean()) {
+                newStateBuilder.blocks(ClusterBlocks.EMPTY_CLUSTER_BLOCK);
+            } else {
+                newStateBuilder.incrementVersion();
+            }
+            setState(clusterService, newStateBuilder.build());
+        }
+        expectThrows(CancellationException.class, listener::actionGet);
+    }
+
+    public void testTaskCancellationOnceActionItIsDispatchedToMaster() throws Exception {
+        TaskManager taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet());
+
+        Request request = new Request();
+        final CancellableTask task = (CancellableTask) taskManager.register("type", "internal:testAction", request);
+
+        // Block all the threads of the executor in which the master operation will be dispatched to
+        // ensure that the master operation won't be executed until the threads are released
+        final String executorName = ThreadPool.Names.GENERIC;
+        final Runnable releaseBlockedThreads = blockAllThreads(executorName);
+
+        PlainActionFuture<Response> listener = new PlainActionFuture<>();
+        ActionTestUtils.execute(new Action("internal:testAction", transportService, clusterService, threadPool, executorName),
+            task,
+            request,
+            listener
+        );
+
+        taskManager.cancel(task, "", () -> {});
+        assertThat(task.isCancelled(), equalTo(true));
+
+        releaseBlockedThreads.run();
+
+        expectThrows(CancellationException.class, listener::actionGet);
+    }
+
+    private Runnable blockAllThreads(String executorName) throws Exception {
+        final int numberOfThreads = threadPool.info(executorName).getMax();
+        final EsThreadPoolExecutor executor = (EsThreadPoolExecutor) threadPool.executor(executorName);
+        final CyclicBarrier barrier = new CyclicBarrier(numberOfThreads + 1);
+        final CountDownLatch latch = new CountDownLatch(1);
+        for (int i = 0; i < numberOfThreads; i++) {
+            executor.submit(() -> {
+                try {
+                    barrier.await();
+                    latch.await();
+                } catch (Exception e) {
+                    throw new AssertionError(e);
+                }
+            });
+        }
+        barrier.await();
+        return latch::countDown;
+    }
 }