Verify task hierarchy is cancelled

Author: danielmitterdorfer

x-pack/plugin/profiler/src/internalClusterTest/java/org/elasticsearch/xpack/profiler/GetProfilingActionIT.java

@@ -43,8 +43,10 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -185,8 +187,9 @@ void verifyCancellation(String action, Request restRequest) throws Exception {
         Cancellable cancellable = getRestClient().performRequestAsync(restRequest, wrapAsRestResponseListener(future));
 
         awaitForBlock(plugins);
+        Collection<TaskId> profilingTasks = collectProfilingRelatedTasks(action);
         cancellable.cancel();
-        ensureTaskIsCancelled(action, nodeIdToName::get);
+        ensureTasksAreCancelled(profilingTasks, nodeIdToName::get);
 
         disableBlocks(plugins);
         expectThrows(CancellationException.class, future::actionGet);
@@ -202,22 +205,35 @@ private static Map<String, String> readNodesInfo() {
         return nodeIdToName;
     }
 
-    private static void ensureTaskIsCancelled(String transportAction, Function<String, String> nodeIdToName) throws Exception {
-        SetOnce<TaskInfo> searchTask = new SetOnce<>();
+    private static Collection<TaskId> collectProfilingRelatedTasks(String transportAction) {
+        SetOnce<TaskInfo> profilingTask = new SetOnce<>();
+        Map<TaskId, Set<TaskId>> taskToParent = new HashMap<>();
         ListTasksResponse listTasksResponse = client().admin().cluster().prepareListTasks().get();
         for (TaskInfo task : listTasksResponse.getTasks()) {
+            TaskId parentTaskId = task.parentTaskId();
+            if (parentTaskId != null) {
+                if (taskToParent.containsKey(parentTaskId) == false) {
+                    taskToParent.put(parentTaskId, new HashSet<>());
+                }
+                taskToParent.get(parentTaskId).add(task.taskId());
+            }
             if (task.action().equals(transportAction)) {
-                searchTask.set(task);
+                profilingTask.set(task);
             }
         }
-        assertNotNull(searchTask.get());
-        TaskId taskId = searchTask.get().taskId();
-        String nodeName = nodeIdToName.apply(taskId.getNodeId());
+        assertNotNull(profilingTask.get());
+        return taskToParent.get(profilingTask.get().taskId());
+    }
+
+    private static void ensureTasksAreCancelled(Collection<TaskId> taskIds, Function<String, String> nodeIdToName) throws Exception {
         assertBusy(() -> {
-            TaskManager taskManager = internalCluster().getInstance(TransportService.class, nodeName).getTaskManager();
-            Task task = taskManager.getTask(taskId.getId());
-            assertThat(task, instanceOf(CancellableTask.class));
-            assertTrue(((CancellableTask) task).isCancelled());
+            for (TaskId taskId : taskIds) {
+                String nodeName = nodeIdToName.apply(taskId.getNodeId());
+                TaskManager taskManager = internalCluster().getInstance(TransportService.class, nodeName).getTaskManager();
+                Task task = taskManager.getTask(taskId.getId());
+                assertThat(task, instanceOf(CancellableTask.class));
+                assertTrue(((CancellableTask) task).isCancelled());
+            }
         });
     }
 
@@ -229,6 +245,10 @@ private static List<ScriptedBlockPlugin> initBlockFactory() {
         for (ScriptedBlockPlugin plugin : plugins) {
             plugin.reset();
             plugin.enableBlock();
+            // Allow to execute one search and only block starting with the second one. This
+            // is done so we have at least one child action and can check that all active children
+            // are cancelled with the parent action.
+            plugin.setSlack(1);
         }
         return plugins;
     }
@@ -255,6 +275,8 @@ public static class ScriptedBlockPlugin extends MockScriptPlugin {
 
         private final AtomicInteger hits = new AtomicInteger();
 
+        private final AtomicInteger slack = new AtomicInteger(0);
+
         private final AtomicBoolean shouldBlock = new AtomicBoolean(true);
 
         void reset() {
@@ -269,16 +291,22 @@ void enableBlock() {
             shouldBlock.set(true);
         }
 
+        void setSlack(int slack) {
+            this.slack.set(slack);
+        }
+
         @Override
         public Map<String, Function<Map<String, Object>, Object>> pluginScripts() {
             return Collections.singletonMap(SCRIPT_NAME, params -> {
                 LeafStoredFieldsLookup fieldsLookup = (LeafStoredFieldsLookup) params.get("_fields");
                 LogManager.getLogger(GetProfilingActionIT.class).info("Blocking on the document {}", fieldsLookup.get("_id"));
                 hits.incrementAndGet();
-                try {
-                    waitUntil(() -> shouldBlock.get() == false);
-                } catch (Exception e) {
-                    throw new RuntimeException(e);
+                if (slack.decrementAndGet() < 0) {
+                    try {
+                        waitUntil(() -> shouldBlock.get() == false);
+                    } catch (Exception e) {
+                        throw new RuntimeException(e);
+                    }
                 }
                 return true;
             });