Use cancel reason in ESQL (elastic#103824)

ComputeService in ESQL should use the cancel reason when cancelling its exchanges. Closes elastic#103746
tvernum · Jan 2, 2024 · 9a7c624 · 9a7c624
1 parent 05c79cc
commit 9a7c624
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/...rnalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClustersCancellationIT.java b/...rnalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClustersCancellationIT.java
@@ -157,7 +157,6 @@ private void createRemoteIndex(int numDocs) throws Exception {
         bulk.get();
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/103746")
     public void testCancel() throws Exception {
         createRemoteIndex(between(10, 100));
         EsqlQueryRequest request = new EsqlQueryRequest();

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java
@@ -575,7 +575,9 @@ public void messageReceived(DataNodeRequest request, TransportChannel channel, T
             final var parentTask = (CancellableTask) task;
             final var sessionId = request.sessionId();
             final var exchangeSink = exchangeService.getSinkHandler(sessionId);
-            parentTask.addListener(() -> exchangeService.finishSinkHandler(sessionId, new TaskCancelledException("task cancelled")));
+            parentTask.addListener(
+                () -> exchangeService.finishSinkHandler(sessionId, new TaskCancelledException(parentTask.getReasonCancelled()))
+            );
             final ActionListener<ComputeResponse> listener = new ChannelActionListener<>(channel);
             final EsqlConfiguration configuration = request.configuration();
             acquireSearchContexts(
@@ -651,7 +653,9 @@ void runComputeOnRemoteCluster(
         ActionListener<ComputeResponse> listener
     ) {
         final var exchangeSink = exchangeService.getSinkHandler(globalSessionId);
-        parentTask.addListener(() -> exchangeService.finishSinkHandler(globalSessionId, new TaskCancelledException("request cancelled")));
+        parentTask.addListener(
+            () -> exchangeService.finishSinkHandler(globalSessionId, new TaskCancelledException(parentTask.getReasonCancelled()))
+        );
         ThreadPool threadPool = transportService.getThreadPool();
         final var responseHeadersCollector = new ResponseHeadersCollector(threadPool.getThreadContext());
         listener = ActionListener.runBefore(listener, responseHeadersCollector::finish);