Handle exceptions when register compute group

Author: dnhatn

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ClusterComputeHandler.java

@@ -109,7 +109,13 @@ void startComputeOnRemoteCluster(
                     groupTask = rootTask;
                     onGroupFailure = cancelQueryOnFailure;
                 } else {
-                    groupTask = computeService.createGroupTask(rootTask, () -> "compute group: cluster [" + clusterAlias + "]");
+                    try {
+                        groupTask = computeService.createGroupTask(rootTask, () -> "compute group: cluster [" + clusterAlias + "]");
+                    } catch (Exception e) {
+                        assert e instanceof TaskCancelledException : new AssertionError(e);
+                        l.onFailure(e);
+                        return;
+                    }
                     onGroupFailure = computeService.cancelQueryOnFailure(groupTask);
                     l = ActionListener.runAfter(l, () -> transportService.getTaskManager().unregister(groupTask));
                 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java

@@ -33,6 +33,7 @@
 import org.elasticsearch.search.lookup.SourceProvider;
 import org.elasticsearch.tasks.CancellableTask;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.tasks.TaskManager;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -431,7 +432,7 @@ Runnable cancelQueryOnFailure(CancellableTask task) {
         });
     }
 
-    CancellableTask createGroupTask(Task parentTask, Supplier<String> description) {
+    CancellableTask createGroupTask(Task parentTask, Supplier<String> description) throws TaskCancelledException {
         final TaskManager taskManager = transportService.getTaskManager();
         return (CancellableTask) taskManager.register(
             "transport",

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java

@@ -130,10 +130,16 @@ protected void sendRequest(
                         final Runnable onGroupFailure;
                         final CancellableTask groupTask;
                         if (allowPartialResults) {
-                            groupTask = computeService.createGroupTask(
-                                parentTask,
-                                () -> "compute group: data-node [" + node.getName() + "], " + shardIds + " [" + shardIds + "]"
-                            );
+                            try {
+                                groupTask = computeService.createGroupTask(
+                                    parentTask,
+                                    () -> "compute group: data-node [" + node.getName() + "], " + shardIds + " [" + shardIds + "]"
+                                );
+                            } catch (Exception e) {
+                                assert e instanceof TaskCancelledException : new AssertionError(e);
+                                l.onFailure(e);
+                                return;
+                            }
                             onGroupFailure = computeService.cancelQueryOnFailure(groupTask);
                             l = ActionListener.runAfter(l, () -> transportService.getTaskManager().unregister(groupTask));
                         } else {