Async search: Add ID and "is running" http headers (#112431)

Add the async execution ID and "is running" flag in the response as HTTP headers.
This allows users to know the request status without having to parse the response body.
It was also implemented in the `/_async_search/status/<id>` endpoint for consistency.

Continuation of #111840, which implemented this same thing for ESQL.
Fixes #109576

Author: ivancea

docs/changelog/112431.yaml

@@ -0,0 +1,6 @@
+pr: 112431
+summary: "Async search: Add ID and \"is running\" http headers"
+area: Search
+type: feature
+issues:
+ - 109576

x-pack/plugin/async-search/qa/rest/build.gradle

@@ -1,6 +1,7 @@
 import org.elasticsearch.gradle.internal.info.BuildParams
 
 apply plugin: 'elasticsearch.base-internal-es-plugin'
+apply plugin: 'elasticsearch.internal-java-rest-test'
 apply plugin: 'elasticsearch.legacy-yaml-rest-test'
 apply plugin: 'elasticsearch.legacy-yaml-rest-compat-test'
 
@@ -10,6 +11,10 @@ esplugin {
   classname 'org.elasticsearch.query.DeprecatedQueryPlugin'
 }
 
+dependencies {
+  clusterPlugins project(xpackModule('async-search'))
+}
+
 restResources {
   restApi {
     include '_common', 'indices', 'index', 'async_search'

x-pack/plugin/async-search/qa/rest/src/javaRestTest/java/org/elasticsearch/qa/AsyncSearchHeadersIT.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.qa;
+
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.test.cluster.ElasticsearchCluster;
+import org.elasticsearch.test.rest.ESRestTestCase;
+import org.junit.Before;
+import org.junit.ClassRule;
+
+import java.io.IOException;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class AsyncSearchHeadersIT extends ESRestTestCase {
+    @ClassRule
+    public static ElasticsearchCluster cluster = ElasticsearchCluster.local().plugin("x-pack-async-search").build();
+
+    @Override
+    protected String getTestRestCluster() {
+        return cluster.getHttpAddresses();
+    }
+
+    @Before
+    public void createIndex() throws IOException {
+        client().performRequest(new Request("PUT", "/test_index"));
+    }
+
+    public void testAsyncHeaders() throws IOException {
+        Response submitResponse = client().performRequest(new Request("POST", "/test_index/_async_search?keep_on_completion=true"));
+        var asyncExecutionId = assertAsyncHeaders(submitResponse);
+
+        Response statusResponse = client().performRequest(new Request("GET", "/_async_search/status/" + asyncExecutionId));
+        assertAsyncHeaders(statusResponse);
+
+        Response resultResponse = client().performRequest(new Request("GET", "/_async_search/" + asyncExecutionId));
+        assertAsyncHeaders(resultResponse);
+    }
+
+    private String assertAsyncHeaders(Response response) throws IOException {
+        var json = entityAsMap(response);
+
+        var asyncExecutionId = (String) json.get("id");
+        var isRunning = (boolean) json.get("is_running");
+
+        if (asyncExecutionId != null) {
+            assertThat(response.getHeader("X-ElasticSearch-Async-Id"), equalTo(asyncExecutionId));
+        }
+        assertThat(response.getHeader("X-ElasticSearch-Async-Is-Running"), equalTo(isRunning ? "?1" : "?0"));
+
+        return asyncExecutionId;
+    }
+}

x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/TransportGetAsyncSearchAction.java

@@ -16,11 +16,13 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.core.XPackPlugin;
+import org.elasticsearch.xpack.core.async.AsyncExecutionId;
 import org.elasticsearch.xpack.core.async.AsyncResultsService;
 import org.elasticsearch.xpack.core.async.AsyncTaskIndexService;
 import org.elasticsearch.xpack.core.async.GetAsyncResultRequest;
@@ -32,6 +34,7 @@
 public class TransportGetAsyncSearchAction extends HandledTransportAction<GetAsyncResultRequest, AsyncSearchResponse> {
     private final AsyncResultsService<AsyncSearchTask, AsyncSearchResponse> resultsService;
     private final TransportService transportService;
+    private final ThreadContext threadContext;
 
     @Inject
     public TransportGetAsyncSearchAction(
@@ -45,6 +48,7 @@ public TransportGetAsyncSearchAction(
     ) {
         super(GetAsyncSearchAction.NAME, transportService, actionFilters, GetAsyncResultRequest::new, EsExecutors.DIRECT_EXECUTOR_SERVICE);
         this.transportService = transportService;
+        this.threadContext = threadPool.getThreadContext();
         this.resultsService = createResultsService(transportService, clusterService, registry, client, threadPool, bigArrays);
     }
 
@@ -78,15 +82,23 @@ static AsyncResultsService<AsyncSearchTask, AsyncSearchResponse> createResultsSe
 
     @Override
     protected void doExecute(Task task, GetAsyncResultRequest request, ActionListener<AsyncSearchResponse> listener) {
+        ActionListener<AsyncSearchResponse> listenerWithHeaders = listener.map(response -> {
+            threadContext.addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_IS_RUNNING_HEADER, response.isRunning() ? "?1" : "?0");
+            if (response.getId() != null) {
+                threadContext.addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_ID_HEADER, response.getId());
+            }
+            return response;
+        });
+
         DiscoveryNode node = resultsService.getNode(request.getId());
         if (node == null || resultsService.isLocalNode(node)) {
-            resultsService.retrieveResult(request, listener);
+            resultsService.retrieveResult(request, listenerWithHeaders);
         } else {
             transportService.sendRequest(
                 node,
                 GetAsyncSearchAction.NAME,
                 request,
-                new ActionListenerResponseHandler<>(listener, AsyncSearchResponse::new, EsExecutors.DIRECT_EXECUTOR_SERVICE)
+                new ActionListenerResponseHandler<>(listenerWithHeaders, AsyncSearchResponse::new, EsExecutors.DIRECT_EXECUTOR_SERVICE)
             );
         }
     }

x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/TransportGetAsyncStatusAction.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.concurrent.EsExecutors;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.tasks.Task;
@@ -40,6 +41,7 @@
 public class TransportGetAsyncStatusAction extends HandledTransportAction<GetAsyncStatusRequest, AsyncStatusResponse> {
     private final TransportService transportService;
     private final ClusterService clusterService;
+    private final ThreadContext threadContext;
     private final AsyncTaskIndexService<AsyncSearchResponse> store;
 
     @Inject
@@ -55,6 +57,7 @@ public TransportGetAsyncStatusAction(
         super(GetAsyncStatusAction.NAME, transportService, actionFilters, GetAsyncStatusRequest::new, EsExecutors.DIRECT_EXECUTOR_SERVICE);
         this.transportService = transportService;
         this.clusterService = clusterService;
+        this.threadContext = threadPool.getThreadContext();
         this.store = new AsyncTaskIndexService<>(
             XPackPlugin.ASYNC_RESULTS_INDEX,
             clusterService,
@@ -73,6 +76,12 @@ protected void doExecute(Task task, GetAsyncStatusRequest request, ActionListene
         DiscoveryNode node = clusterService.state().nodes().get(searchId.getTaskId().getNodeId());
         DiscoveryNode localNode = clusterService.state().getNodes().getLocalNode();
 
+        ActionListener<AsyncStatusResponse> listenerWithHeaders = listener.map(response -> {
+            threadContext.addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_IS_RUNNING_HEADER, response.isRunning() ? "?1" : "?0");
+            threadContext.addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_ID_HEADER, response.getId());
+            return response;
+        });
+
         if (node == null || Objects.equals(node, localNode)) {
             if (request.getKeepAlive() != null && request.getKeepAlive().getMillis() > 0) {
                 long expirationTime = System.currentTimeMillis() + request.getKeepAlive().getMillis();
@@ -87,17 +96,17 @@ protected void doExecute(Task task, GetAsyncStatusRequest request, ActionListene
                         AsyncSearchTask.class,
                         AsyncSearchTask::getStatusResponse,
                         AsyncStatusResponse::getStatusFromStoredSearch,
-                        listener
+                        listenerWithHeaders
                     );
                 }, exc -> {
                     RestStatus status = ExceptionsHelper.status(ExceptionsHelper.unwrapCause(exc));
                     if (status != RestStatus.NOT_FOUND) {
                         logger.error(() -> format("failed to update expiration time for async-search [%s]", searchId.getEncoded()), exc);
-                        listener.onFailure(exc);
+                        listenerWithHeaders.onFailure(exc);
                     } else {
                         // the async search document or its index is not found.
                         // That can happen if an invalid/deleted search id is provided.
-                        listener.onFailure(new ResourceNotFoundException(searchId.getEncoded()));
+                        listenerWithHeaders.onFailure(new ResourceNotFoundException(searchId.getEncoded()));
                     }
                 }));
             } else {
@@ -107,15 +116,15 @@ protected void doExecute(Task task, GetAsyncStatusRequest request, ActionListene
                     AsyncSearchTask.class,
                     AsyncSearchTask::getStatusResponse,
                     AsyncStatusResponse::getStatusFromStoredSearch,
-                    listener
+                    listenerWithHeaders
                 );
             }
         } else {
             transportService.sendRequest(
                 node,
                 GetAsyncStatusAction.NAME,
                 request,
-                new ActionListenerResponseHandler<>(listener, AsyncStatusResponse::new, EsExecutors.DIRECT_EXECUTOR_SERVICE)
+                new ActionListenerResponseHandler<>(listenerWithHeaders, AsyncStatusResponse::new, EsExecutors.DIRECT_EXECUTOR_SERVICE)
             );
         }
     }

x-pack/plugin/async-search/src/main/java/org/elasticsearch/xpack/search/TransportSubmitAsyncSearchAction.java

@@ -92,6 +92,14 @@ protected void doExecute(Task submitTask, SubmitAsyncSearchRequest request, Acti
                 searchRequest
             );
             searchAction.execute(searchTask, searchRequest, searchTask.getSearchProgressActionListener());
+
+            ActionListener<AsyncSearchResponse> submitListenerWithHeaders = submitListener.map(response -> {
+                threadContext.addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_IS_RUNNING_HEADER, response.isRunning() ? "?1" : "?0");
+                if (response.getId() != null) {
+                    threadContext.addResponseHeader(AsyncExecutionId.ASYNC_EXECUTION_ID_HEADER, response.getId());
+                }
+                return response;
+            });
             searchTask.addCompletionListener(new ActionListener<>() {
                 @Override
                 public void onResponse(AsyncSearchResponse searchResponse) {
@@ -119,14 +127,14 @@ public void onResponse(DocWriteResponse r) {
                                                         finalResponse -> onFinalResponse(searchTask, finalResponse, () -> {})
                                                     );
                                                 } finally {
-                                                    submitListener.onResponse(searchResponse);
+                                                    submitListenerWithHeaders.onResponse(searchResponse);
                                                 }
                                             } else {
                                                 searchResponse.mustIncRef();
                                                 onFinalResponse(
                                                     searchTask,
                                                     searchResponse,
-                                                    () -> ActionListener.respondAndRelease(submitListener, searchResponse)
+                                                    () -> ActionListener.respondAndRelease(submitListenerWithHeaders, searchResponse)
                                                 );
                                             }
                                         }
@@ -138,7 +146,7 @@ public void onFailure(Exception exc) {
                                                 exc,
                                                 searchResponse.isRunning(),
                                                 "fatal failure: unable to store initial response",
-                                                submitListener
+                                                submitListenerWithHeaders
                                             );
                                         }
                                     }, searchResponse::decRef)
@@ -147,14 +155,20 @@ public void onFailure(Exception exc) {
                                 initialResp.decRef();
                             }
                         } catch (Exception exc) {
-                            onFatalFailure(searchTask, exc, searchResponse.isRunning(), "fatal failure: generic error", submitListener);
+                            onFatalFailure(
+                                searchTask,
+                                exc,
+                                searchResponse.isRunning(),
+                                "fatal failure: generic error",
+                                submitListenerWithHeaders
+                            );
                         }
                     } else {
                         try (searchTask) {
                             // the task completed within the timeout so the response is sent back to the user
                             // with a null id since nothing was stored on the cluster.
                             taskManager.unregister(searchTask);
-                            ActionListener.respondAndRelease(submitListener, searchResponse.clone(null));
+                            ActionListener.respondAndRelease(submitListenerWithHeaders, searchResponse.clone(null));
                         }
                     }
                 }
@@ -163,7 +177,7 @@ public void onFailure(Exception exc) {
                 public void onFailure(Exception exc) {
                     // this will only ever be called if there is an issue scheduling the thread that executes
                     // the completion listener once the wait for completion timeout expires.
-                    onFatalFailure(searchTask, exc, true, "fatal failure: addCompletionListener", submitListener);
+                    onFatalFailure(searchTask, exc, true, "fatal failure: addCompletionListener", submitListenerWithHeaders);
                 }
             }, request.getWaitForCompletionTimeout());
         }