Improve stop action

Author: smalyshev

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

@@ -9,7 +9,6 @@
 
 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;
 import org.elasticsearch.action.support.ActionFilters;
@@ -18,7 +17,6 @@
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
-import org.elasticsearch.compute.EsqlRefCountingListener;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
 import org.elasticsearch.injection.guice.Inject;
@@ -33,11 +31,11 @@
 import org.elasticsearch.xpack.core.async.GetAsyncResultRequest;
 import org.elasticsearch.xpack.core.security.SecurityContext;
 import org.elasticsearch.xpack.esql.action.EsqlAsyncStopAction;
+import org.elasticsearch.xpack.esql.action.EsqlExecutionInfo;
 import org.elasticsearch.xpack.esql.action.EsqlQueryResponse;
 import org.elasticsearch.xpack.esql.action.EsqlQueryTask;
 
 import java.io.IOException;
-import java.util.concurrent.atomic.AtomicReference;
 
 import static org.elasticsearch.xpack.core.ClientHelper.ASYNC_SEARCH_ORIGIN;
 
@@ -108,36 +106,26 @@ 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);
+        GetAsyncResultRequest getAsyncResultRequest = new GetAsyncResultRequest(asyncIdStr);
+        if (asyncTask == null) {
             // 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);
-            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();
-            }
+            logger.debug("Async stop for task {}, no task present - passing to GetAsyncResultRequest", asyncIdStr);
+            getResultsAction.execute(task, getAsyncResultRequest, listener);
             return;
         }
-
-        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()))) {
-            asyncListener.addListener(refs.acquire().map(r -> {
-                responseHolder.set(r);
-                return null;
-            }));
-            asyncListener.markAsPartial();
-            exchangeService.finishSessionEarly(sessionID(asyncId), refs.acquire());
+        final EsqlExecutionInfo esqlExecutionInfo = asyncTask.executionInfo();
+        if (esqlExecutionInfo != null) {
+            esqlExecutionInfo.markAsPartial();
         }
+        Runnable getResults = () -> getResultsAction.execute(task, getAsyncResultRequest, listener);
+        exchangeService.finishSessionEarly(sessionID(asyncId), ActionListener.running(() -> {
+            if (asyncTask.addCompletionListener(() -> ActionListener.running(getResults)) == false) {
+                getResults.run();
+            }
+        }));
     }
 
     private EsqlQueryTask getEsqlQueryTask(AsyncExecutionId asyncId) {

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

@@ -13,14 +13,12 @@
 import org.elasticsearch.action.admin.cluster.stats.CCSUsageTelemetry;
 import org.elasticsearch.action.support.ActionFilters;
 import org.elasticsearch.action.support.HandledTransportAction;
-import org.elasticsearch.action.support.SubscribableListener;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.util.BigArrays;
-import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
@@ -82,8 +80,6 @@ public class TransportEsqlQueryAction extends HandledTransportAction<EsqlQueryRe
     private final RemoteClusterService remoteClusterService;
     private final UsageService usageService;
     private final TransportActionServices services;
-    // Listeners for active async queries, key being the async task execution ID
-    private final Map<String, EsqlQueryListener> asyncListeners = ConcurrentCollections.newConcurrentMap();
 
     @Inject
     @SuppressWarnings("this-escape")
@@ -190,41 +186,11 @@ private void doExecuteForked(Task task, EsqlQueryRequest request, ActionListener
         }
     }
 
-    // Subscribable listener that can keep track of the EsqlExecutionInfo
-    // Used to mark an async query as partial if it is stopped
-    public static class EsqlQueryListener extends SubscribableListener<EsqlQueryResponse> {
-        private EsqlExecutionInfo executionInfo;
-
-        public EsqlQueryListener(EsqlExecutionInfo executionInfo) {
-            this.executionInfo = executionInfo;
-        }
-
-        public EsqlExecutionInfo getExecutionInfo() {
-            return executionInfo;
-        }
-
-        public void markAsPartial() {
-            if (executionInfo != null) {
-                executionInfo.markAsPartial();
-            }
-        }
-    }
-
     @Override
     public void execute(EsqlQueryRequest request, EsqlQueryTask task, ActionListener<EsqlQueryResponse> listener) {
         // set EsqlExecutionInfo on async-search task so that it is accessible to GET _query/async while the query is still running
         task.setExecutionInfo(createEsqlExecutionInfo(request));
-        // Since the request is async here, we need to wrap the listener in a SubscribableListener so that we can collect the results from
-        // other endpoints, such as _query/async/stop
-        EsqlQueryListener subListener = new EsqlQueryListener(task.executionInfo());
-        String asyncExecutionId = task.getExecutionId().getEncoded();
-        subListener.addListener(ActionListener.runAfter(listener, () -> asyncListeners.remove(asyncExecutionId)));
-        asyncListeners.put(asyncExecutionId, subListener);
-        ActionListener.run(subListener, l -> innerExecute(task, request, l));
-    }
-
-    public EsqlQueryListener getAsyncListener(String executionId) {
-        return asyncListeners.get(executionId);
+        ActionListener.run(listener, l -> innerExecute(task, request, l));
     }
 
     private void innerExecute(Task task, EsqlQueryRequest request, ActionListener<EsqlQueryResponse> listener) {