Add support for task cancellation to TransportNodesAction (elastic#71895)

Backport of elastic#71695

Author: fcofdez

server/src/internalClusterTest/java/org/elasticsearch/action/admin/cluster/node/tasks/TasksIT.java

@@ -43,6 +43,7 @@
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.tasks.TaskResult;
@@ -496,7 +497,7 @@ public void testTasksCancellation() throws Exception {
                 .setActions(TestTaskPlugin.TestTaskAction.NAME).get();
         assertEquals(1, cancelTasksResponse.getTasks().size());
 
-        future.get();
+        expectThrows(TaskCancelledException.class, future::actionGet);
 
         logger.info("--> checking that test tasks are not running");
         assertEquals(0,

server/src/main/java/org/elasticsearch/action/support/nodes/TransportNodesAction.java

@@ -18,7 +18,9 @@
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportChannel;
 import org.elasticsearch.transport.TransportException;
@@ -203,9 +205,7 @@ class AsyncAction {
         void start() {
             final DiscoveryNode[] nodes = request.concreteNodes();
             if (nodes.length == 0) {
-                // nothing to notify, so respond immediately, but always fork even if finalExecutor == SAME
-                final String executor = finalExecutor.equals(ThreadPool.Names.SAME) ? ThreadPool.Names.GENERIC : finalExecutor;
-                threadPool.executor(executor).execute(() -> newResponse(task, request, responses, listener));
+                finishHim();
                 return;
             }
             final TransportRequestOptions transportRequestOptions = TransportRequestOptions.timeout(request.timeout());
@@ -258,14 +258,27 @@ private void onFailure(int idx, String nodeId, Throwable t) {
         }
 
         private void finishHim() {
-            threadPool.executor(finalExecutor).execute(() -> newResponse(task, request, responses, listener));
+            if (isCancelled(task)) {
+                listener.onFailure(new TaskCancelledException("task cancelled"));
+                return;
+            }
+
+            final String executor = finalExecutor.equals(ThreadPool.Names.SAME) ? ThreadPool.Names.GENERIC : finalExecutor;
+            threadPool.executor(executor).execute(() -> newResponse(task, request, responses, listener));
         }
     }
 
-    class NodeTransportHandler implements TransportRequestHandler<NodeRequest> {
+    private boolean isCancelled(Task task) {
+        return task instanceof CancellableTask && ((CancellableTask) task).isCancelled();
+    }
 
+    class NodeTransportHandler implements TransportRequestHandler<NodeRequest> {
         @Override
         public void messageReceived(NodeRequest request, TransportChannel channel, Task task) throws Exception {
+            if (isCancelled(task)) {
+                throw new TaskCancelledException("task cancelled");
+            }
+
             channel.sendResponse(nodeOperation(request, task));
         }
     }

server/src/test/java/org/elasticsearch/action/admin/cluster/node/tasks/CancellableTasksTests.java

@@ -54,6 +54,7 @@
 import static org.elasticsearch.test.ClusterServiceUtils.setState;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.hamcrest.Matchers.startsWith;
 
@@ -224,7 +225,9 @@ public void testBasicTaskCancellation() throws Exception {
         final AtomicReference<NodesResponse> responseReference = new AtomicReference<>();
         final AtomicReference<Throwable> throwableReference = new AtomicReference<>();
         int runNodesCount = randomIntBetween(1, nodesCount);
-        int blockedNodesCount = randomIntBetween(0, runNodesCount);
+        // Block at least 1 node, otherwise it's quite easy to end up in a race condition where the node tasks
+        // have finished before the cancel request has arrived
+        int blockedNodesCount = randomIntBetween(1, runNodesCount);
         Task mainTask = startCancellableTestNodesAction(waitForActionToStart, runNodesCount, blockedNodesCount,
             new ActionListener<NodesResponse>() {
                 @Override
@@ -261,12 +264,7 @@ public void onFailure(Exception e) {
             assertEquals(runNodesCount, responseReference.get().getNodes().size());
             assertEquals(0, responseReference.get().failureCount());
         } else {
-            // We canceled the request, in this case it should have fail, but we should get partial response
-            assertNull(throwableReference.get());
-            assertEquals(runNodesCount, responseReference.get().failureCount() + responseReference.get().getNodes().size());
-            // and we should have at least as many failures as the number of blocked operations
-            // (we might have cancelled some non-blocked operations before they even started and that's ok)
-            assertThat(responseReference.get().failureCount(), greaterThanOrEqualTo(blockedNodesCount));
+            assertThat(throwableReference.get(), instanceOf(TaskCancelledException.class));
 
             // We should have the information about the cancelled task in the cancel operation response
             assertEquals(1, response.getTasks().size());

server/src/test/java/org/elasticsearch/action/admin/cluster/node/tasks/TaskManagerTestCase.java

@@ -154,7 +154,6 @@ protected NodeResponse newNodeResponse(StreamInput in) throws IOException {
 
         @Override
         protected abstract NodeResponse nodeOperation(NodeRequest request);
-
     }
 
     public static class TestNode implements Releasable {

server/src/test/java/org/elasticsearch/action/support/nodes/TransportNodesActionTests.java

@@ -22,7 +22,9 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.transport.CapturingTransport;
 import org.elasticsearch.threadpool.TestThreadPool;
@@ -41,6 +43,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 import java.util.function.Supplier;
@@ -138,6 +141,37 @@ public void testCustomResolving() throws Exception {
         assertEquals(clusterService.state().nodes().getDataNodes().size(), capturedRequests.size());
     }
 
+    public void testTaskCancellationThrowsException() {
+        TransportNodesAction<TestNodesRequest, TestNodesResponse, TestNodeRequest, TestNodeResponse> action = getTestTransportNodesAction();
+        List<String> nodeIds = new ArrayList<>();
+        for (DiscoveryNode node : clusterService.state().nodes()) {
+            nodeIds.add(node.getId());
+        }
+
+        TestNodesRequest request = new TestNodesRequest(nodeIds.toArray(new String[0]));
+        PlainActionFuture<TestNodesResponse> listener = new PlainActionFuture<>();
+        Task cancellableTask = new CancellableTask(randomLong(), "transport", "action", "", null, emptyMap()) {
+            @Override
+            public boolean isCancelled() {
+                return true;
+            }
+        };
+        action.doExecute(cancellableTask, request, listener);
+        Map<String, List<CapturingTransport.CapturedRequest>> capturedRequests = transport.getCapturedRequestsByTargetNodeAndClear();
+        for (List<CapturingTransport.CapturedRequest> requests : capturedRequests.values()) {
+            for (CapturingTransport.CapturedRequest capturedRequest : requests) {
+                if (randomBoolean()) {
+                    transport.handleResponse(capturedRequest.requestId, new TestNodeResponse(capturedRequest.node));
+                } else {
+                    transport.handleRemoteError(capturedRequest.requestId, new TaskCancelledException("simulated"));
+                }
+            }
+        }
+
+        assertTrue(listener.isDone());
+        expectThrows(ExecutionException.class, TaskCancelledException.class, listener::get);
+    }
+
     private <T> List<T> mockList(Supplier<T> supplier, int size) {
         List<T> failures = new ArrayList<>(size);
         for (int i = 0; i < size; ++i) {
@@ -321,8 +355,13 @@ private static class TestNodeRequest extends BaseNodeRequest {
 
     private static class TestNodeResponse extends BaseNodeResponse {
         TestNodeResponse() {
-            super(mock(DiscoveryNode.class));
+            this(mock(DiscoveryNode.class));
         }
+
+        TestNodeResponse(DiscoveryNode node) {
+            super(node);
+        }
+
         protected TestNodeResponse(StreamInput in) throws IOException {
             super(in);
         }