Ensure unregister child node if failed to register task

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

@@ -146,14 +146,19 @@ Task registerAndExecute(String type, TransportAction<Request, Response> action,
         } else {
             unregisterChildNode = () -> {};
         }
-        Task task = register(type, action.actionName, request);
+        final Task task;
+        try {
+            task = register(type, action.actionName, request);
+        } catch (TaskCancelledException e) {
+            unregisterChildNode.close();
+            throw e;
+        }
         // NOTE: ActionListener cannot infer Response, see https://bugs.openjdk.java.net/browse/JDK-8203195
         action.execute(task, request, new ActionListener<Response>() {
             @Override
             public void onResponse(Response response) {
                 try {
-                    unregisterChildNode.close();
-                    unregister(task);
+                    Releasables.close(unregisterChildNode, () -> unregister(task));
                 } finally {
                     onResponse.accept(task, response);
                 }

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

@@ -52,7 +52,6 @@
 import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.test.ESIntegTestCase;
-import org.elasticsearch.test.junit.annotations.TestIssueLogging;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportException;
 import org.elasticsearch.transport.TransportResponseHandler;
@@ -80,9 +79,6 @@
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.instanceOf;
 
-@TestIssueLogging(
-    value = "org.elasticsearch.action.admin.cluster.node.tasks.cancel:TRACE,org.elasticsearch.tasks:TRACE",
-    issueUrl = "https://github.com/elastic/elasticsearch/issues/55875")
 public class CancellableTasksIT extends ESIntegTestCase {
 
     static int idGenerator = 0;
@@ -241,7 +237,6 @@ public void testCancelTaskMultipleTimes() throws Exception {
         ensureAllBansRemoved();
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/55875")
     public void testDoNotWaitForCompletion() throws Exception {
         Set<DiscoveryNode> nodes = StreamSupport.stream(clusterService().state().nodes().spliterator(), false).collect(Collectors.toSet());
         TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 3));

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

@@ -53,6 +53,7 @@
 import java.util.Random;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Phaser;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicIntegerArray;
 import java.util.concurrent.atomic.AtomicReference;
@@ -61,6 +62,7 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
+import static org.hamcrest.Matchers.startsWith;
 
 public class CancellableTasksTests extends TaskManagerTestCase {
 
@@ -347,6 +349,26 @@ public void onFailure(Exception e) {
         });
     }
 
+    public void testRegisterAndExecuteChildTaskWhileParentTaskIsBeingCanceled() throws Exception {
+        setupTestNodes(Settings.EMPTY);
+        connectNodes(testNodes);
+        final TaskManager taskManager = testNodes[0].transportService.getTaskManager();
+        CancellableNodesRequest parentRequest = new CancellableNodesRequest("parent");
+        final Task parentTask = taskManager.register("test", "test", parentRequest);
+        final TaskId parentTaskId = parentTask.taskInfo(testNodes[0].getNodeId(), false).getTaskId();
+        taskManager.setBan(new TaskId(testNodes[0].getNodeId(), parentTask.getId()), "test");
+        CancellableNodesRequest childRequest = new CancellableNodesRequest("child");
+        childRequest.setParentTask(parentTaskId);
+        CancellableTestNodesAction testAction = new CancellableTestNodesAction("internal:testAction", threadPool, testNodes[1]
+            .clusterService, testNodes[0].transportService, false, new CountDownLatch(1));
+        TaskCancelledException cancelledException = expectThrows(TaskCancelledException.class,
+            () -> taskManager.registerAndExecute("test", testAction, childRequest, (task, response) -> {}, (task, e) -> {}));
+        assertThat(cancelledException.getMessage(), startsWith("Task cancelled before it started:"));
+        CountDownLatch latch = new CountDownLatch(1);
+        taskManager.startBanOnChildrenNodes(parentTaskId.getId(), latch::countDown);
+        assertTrue("onChildTasksCompleted() is not invoked", latch.await(1, TimeUnit.SECONDS));
+    }
+
     public void testTaskCancellationOnCoordinatingNodeLeavingTheCluster() throws Exception {
         setupTestNodes(Settings.EMPTY);
         connectNodes(testNodes);