Use opaque id in task cancellation assertion (#110680)

Add use of Opaque ID HTTP header in task cancellation assertion. In some
tests, like this #88201 `testCatSegmentsRestCancellation`, we assert
that all tasks related to specific HTTP request are cancelled. But we do
blanket approach in assertion block catching all tasks by action name. I
think narrowing down assertion to specific http request in this case
would be more accurate.

It is still not clear why test mentioned above failing, but after hours
of investigation and injecting random delays, I'm inclining more to
@DaveCTurner's comment about interference from other tests or cluster
activity. I added additional log that will report when we spot task with
different opaque id.

Author: mhl-b

qa/smoke-test-http/src/javaRestTest/java/org/elasticsearch/http/BlockedSearcherRestCancellationTestCase.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.plugins.EnginePlugin;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.tasks.Task;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -76,6 +77,10 @@ void runTest(Request request, String actionPrefix) throws Exception {
         createIndex("test", Settings.builder().put(BLOCK_SEARCHER_SETTING.getKey(), true).build());
         ensureGreen("test");
 
+        assert request.getOptions().containsHeader(Task.X_OPAQUE_ID_HTTP_HEADER) == false;
+        final var opaqueId = getTestClass().getSimpleName() + "-" + getTestName() + "-" + randomUUID();
+        request.setOptions(request.getOptions().toBuilder().addHeader(Task.X_OPAQUE_ID_HTTP_HEADER, opaqueId));
+
         final List<Semaphore> searcherBlocks = new ArrayList<>();
         for (final IndicesService indicesService : internalCluster().getInstances(IndicesService.class)) {
             for (final IndexService indexService : indicesService) {
@@ -96,7 +101,8 @@ void runTest(Request request, String actionPrefix) throws Exception {
             }
 
             final PlainActionFuture<Response> future = new PlainActionFuture<>();
-            logger.info("--> sending request");
+            logger.info("--> sending request, opaque id={}", opaqueId);
+
             final Cancellable cancellable = getRestClient().performRequestAsync(request, wrapAsRestResponseListener(future));
 
             awaitTaskWithPrefix(actionPrefix);
@@ -108,7 +114,7 @@ void runTest(Request request, String actionPrefix) throws Exception {
             cancellable.cancel();
             expectThrows(CancellationException.class, future::actionGet);
 
-            assertAllCancellableTasksAreCancelled(actionPrefix);
+            assertAllCancellableTasksAreCancelled(actionPrefix, opaqueId);
         } finally {
             Releasables.close(releasables);
         }

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

@@ -109,6 +109,11 @@ public final <T> boolean notifyIfCancelled(ActionListener<T> listener) {
         return true;
     }
 
+    @Override
+    public String toString() {
+        return "CancellableTask{" + super.toString() + ", reason='" + reason + '\'' + ", isCancelled=" + isCancelled + '}';
+    }
+
     private TaskCancelledException getTaskCancelledException() {
         assert Thread.holdsLock(this);
         assert isCancelled;

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

@@ -225,6 +225,8 @@ public String toString() {
             + parentTask
             + ", startTime="
             + startTime
+            + ", headers="
+            + headers
             + ", startTimeNanos="
             + startTimeNanos
             + '}';

test/framework/src/main/java/org/elasticsearch/test/TaskAssertions.java

@@ -10,16 +10,19 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskInfo;
-import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.transport.TransportService;
 
+import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
+import static junit.framework.TestCase.assertFalse;
 import static junit.framework.TestCase.assertTrue;
 import static junit.framework.TestCase.fail;
 import static org.elasticsearch.test.ESIntegTestCase.client;
@@ -59,30 +62,28 @@ private static void awaitTaskWithPrefix(String actionPrefix, Iterable<TransportS
         });
     }
 
-    public static void assertAllCancellableTasksAreCancelled(String actionPrefix) throws Exception {
+    public static void assertAllCancellableTasksAreCancelled(String actionPrefix, @Nullable String opaqueId) throws Exception {
         logger.info("--> checking that all tasks with prefix {} are marked as cancelled", actionPrefix);
 
         assertBusy(() -> {
-            boolean foundTask = false;
+            var tasks = new ArrayList<CancellableTask>();
             for (TransportService transportService : internalCluster().getInstances(TransportService.class)) {
-                final TaskManager taskManager = transportService.getTaskManager();
+                var taskManager = transportService.getTaskManager();
                 assertTrue(taskManager.assertCancellableTaskConsistency());
-                for (CancellableTask cancellableTask : taskManager.getCancellableTasks().values()) {
-                    if (cancellableTask.getAction().startsWith(actionPrefix)) {
-                        logger.trace("--> found task with prefix [{}]: [{}]", actionPrefix, cancellableTask);
-                        foundTask = true;
-                        assertTrue(
-                            "task " + cancellableTask.getId() + "/" + cancellableTask.getAction() + " not cancelled",
-                            cancellableTask.isCancelled()
-                        );
-                        logger.trace("--> Task with prefix [{}] is marked as cancelled: [{}]", actionPrefix, cancellableTask);
-                    }
-                }
+                taskManager.getCancellableTasks().values().stream().filter(t -> t.getAction().startsWith(actionPrefix)).forEach(tasks::add);
             }
-            assertTrue("found no cancellable tasks", foundTask);
+            assertFalse("no tasks found for action: " + actionPrefix, tasks.isEmpty());
+            assertTrue(
+                tasks.toString(),
+                tasks.stream().allMatch(t -> t.isCancelled() && Objects.equals(t.getHeader(Task.X_OPAQUE_ID_HTTP_HEADER), opaqueId))
+            );
         }, 30, TimeUnit.SECONDS);
     }
 
+    public static void assertAllCancellableTasksAreCancelled(String actionPrefix) throws Exception {
+        assertAllCancellableTasksAreCancelled(actionPrefix, null);
+    }
+
     public static void assertAllTasksHaveFinished(String actionPrefix) throws Exception {
         logger.info("--> checking that all tasks with prefix {} have finished", actionPrefix);
         assertBusy(() -> {