Fix async stop sometimes not properly collecting result

Author: smalyshev

muted-tests.yml

@@ -282,9 +282,6 @@ tests:
 - class: org.elasticsearch.xpack.ml.integration.ClassificationIT
   method: testWithDatastreams
   issue: https://github.com/elastic/elasticsearch/issues/121236
-- class: org.elasticsearch.xpack.remotecluster.RemoteClusterSecurityEsqlIT
-  method: testCrossClusterAsyncQueryStop
-  issue: https://github.com/elastic/elasticsearch/issues/121249
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT
   method: test {p0=transform/*}
   issue: https://github.com/elastic/elasticsearch/issues/120816

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

@@ -7,6 +7,8 @@
 
 package org.elasticsearch.xpack.esql.plugin;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionListenerResponseHandler;
@@ -19,7 +21,6 @@
 import org.elasticsearch.compute.EsqlRefCountingListener;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
-import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
@@ -36,7 +37,6 @@
 import org.elasticsearch.xpack.esql.action.EsqlQueryTask;
 
 import java.io.IOException;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static org.elasticsearch.xpack.core.ClientHelper.ASYNC_SEARCH_ORIGIN;
@@ -55,6 +55,8 @@ public class TransportEsqlAsyncStopAction extends HandledTransportAction<AsyncSt
     private final TransportService transportService;
     private final AsyncSearchSecurity security;
 
+    private static final Logger logger = LogManager.getLogger(TransportEsqlAsyncStopAction.class);
+
     @Inject
     public TransportEsqlAsyncStopAction(
         TransportService transportService,
@@ -107,24 +109,25 @@ private String sessionID(AsyncExecutionId asyncId) {
     private void stopQueryAndReturnResult(Task task, AsyncExecutionId asyncId, ActionListener<EsqlQueryResponse> listener) {
         String asyncIdStr = asyncId.getEncoded();
         TransportEsqlQueryAction.EsqlQueryListener asyncListener = queryAction.getAsyncListener(asyncIdStr);
+        EsqlQueryTask asyncTask = getEsqlQueryTask(asyncId);
         if (asyncListener == null) {
+            logger.debug("Async stop for task {}, no listener - collecting the result", asyncIdStr);
             // This should mean one of the two things: either bad request ID, or the query has already finished
             // In both cases, let regular async get deal with it.
             var getAsyncResultRequest = new GetAsyncResultRequest(asyncIdStr);
-            // TODO: this should not be happening, but if the listener is not registered and the query is not finished,
-            // we give it some time to finish
-            getAsyncResultRequest.setWaitForCompletionTimeout(new TimeValue(1, TimeUnit.SECONDS));
-            getResultsAction.execute(task, getAsyncResultRequest, listener);
+            Runnable getResults = () -> getResultsAction.execute(task, getAsyncResultRequest, listener);
+            // If the listener is not present but the task is still alive, this may mean it's not finished writing the response
+            // We will wait for the task to be done and then collect the results.
+            if (asyncTask == null || asyncTask.addCompletionListener(() -> ActionListener.running(getResults)) == false) {
+                getResults.run();
+            }
             return;
         }
-        try {
-            EsqlQueryTask asyncTask = AsyncTaskIndexService.getTask(taskManager, asyncId, EsqlQueryTask.class);
-            if (false == security.currentUserHasAccessToTask(asyncTask)) {
-                throw new ResourceNotFoundException(asyncId + " not found");
-            }
-        } catch (IOException e) {
-            throw new ResourceNotFoundException(asyncId + " not found", e);
+
+        if (asyncTask == null) {
+            throw new ResourceNotFoundException(asyncId + " not found");
         }
+        logger.debug("Async stop for task {} - stopping", asyncIdStr);
         // Here we will wait for both the response to become available and for the finish operation to complete
         var responseHolder = new AtomicReference<EsqlQueryResponse>();
         try (var refs = new EsqlRefCountingListener(listener.map(unused -> responseHolder.get()))) {
@@ -136,4 +139,16 @@ private void stopQueryAndReturnResult(Task task, AsyncExecutionId asyncId, Actio
             exchangeService.finishSessionEarly(sessionID(asyncId), refs.acquire());
         }
     }
+
+    private EsqlQueryTask getEsqlQueryTask(AsyncExecutionId asyncId) {
+        try {
+            EsqlQueryTask asyncTask = AsyncTaskIndexService.getTask(taskManager, asyncId, EsqlQueryTask.class);
+            if (asyncTask == null || false == security.currentUserHasAccessToTask(asyncTask)) {
+                return null;
+            }
+            return asyncTask;
+        } catch (IOException e) {
+            return null;
+        }
+    }
 }