Fix sporadic failures in AsyncSearchAsyncTests

server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java

@@ -375,6 +375,11 @@ private void onShardFailure(final int shardIndex, @Nullable ShardRouting shard,
         // we do make sure to clean it on a successful response from a shard
         SearchShardTarget shardTarget = shardIt.newSearchShardTarget(nodeId);
         onShardFailure(shardIndex, shardTarget, e);
+        final ShardRouting nextShard = shardIt.nextOrNull();
+        final boolean lastShard = nextShard == null;
+        if (lastShard) {
+            onShardGroupFailure(shardIndex, shardTarget, e);
+        }
 
         if (totalOps.incrementAndGet() == expectedTotalOps) {
             if (logger.isDebugEnabled()) {
@@ -385,11 +390,8 @@ private void onShardFailure(final int shardIndex, @Nullable ShardRouting shard,
                     logger.trace(new ParameterizedMessage("{}: Failed to execute [{}]", shard, request), e);
                 }
             }
-            onShardGroupFailure(shardIndex, shardTarget, e);
             onPhaseDone();
         } else {
-            final ShardRouting nextShard = shardIt.nextOrNull();
-            final boolean lastShard = nextShard == null;
             // trace log this exception
             logger.trace(() -> new ParameterizedMessage(
                 "{}: Failed to execute [{}] lastShard [{}]",
@@ -405,7 +407,6 @@ private void onShardFailure(final int shardIndex, @Nullable ShardRouting shard,
                             shard != null ? shard.shortSummary() : shardIt.shardId(), request, lastShard), e);
                     }
                 }
-                onShardGroupFailure(shardIndex, shardTarget, e);
             }
         }
     }

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

@@ -83,7 +83,6 @@ synchronized void updatePartialResponse(int successfulShards, SearchResponseSect
             throw new IllegalStateException("received partial response out of order: "
                 + newSections.getNumReducePhases() + " < " + sections.getNumReducePhases());
         }
-        failIfFrozen();
         ++ version;
         this.successfulShards = successfulShards;
         this.sections = newSections;

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

@@ -83,7 +83,10 @@ public void onResponse(AsyncSearchResponse searchResponse) {
                                 onFatalFailure(searchTask, cause, false, submitListener);
                             } else {
                                 final String docId = searchTask.getSearchId().getDocId();
-                                store.storeInitialResponse(docId, searchTask.getOriginHeaders(), searchResponse,
+                                // creates the fallback response if the node crashes/restarts in the middle of the request
+                                // TODO: store intermediate results ?
+                                AsyncSearchResponse initialResp = searchResponse.clone(searchResponse.getId());
+                                store.storeInitialResponse(docId, searchTask.getOriginHeaders(), initialResp,
                                     new ActionListener<>() {
                                         @Override
                                         public void onResponse(IndexResponse r) {

x-pack/plugin/async-search/src/test/java/org/elasticsearch/xpack/search/AsyncSearchActionTests.java

@@ -6,7 +6,6 @@
 
 package org.elasticsearch.xpack.search;
 
-import org.apache.lucene.util.LuceneTestCase;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.common.settings.Settings;
@@ -37,7 +36,6 @@
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 // TODO: add tests for keepAlive and expiration
-@LuceneTestCase.AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/53360")
 public class AsyncSearchActionTests extends AsyncSearchIntegTestCase {
     private String indexName;
     private int numShards;
@@ -51,7 +49,7 @@ public class AsyncSearchActionTests extends AsyncSearchIntegTestCase {
     public void indexDocuments() throws InterruptedException {
         indexName = "test-async";
         numShards = randomIntBetween(internalCluster().numDataNodes(), internalCluster().numDataNodes()*10);
-        int numDocs = randomIntBetween(numShards, numShards*10);
+        int numDocs = randomIntBetween(numShards, numShards*3);
         createIndex(indexName, Settings.builder().put("index.number_of_shards", numShards).build());
         numKeywords = randomIntBetween(1, 100);
         keywordFreqs = new HashMap<>();
@@ -216,6 +214,7 @@ public void testCleanupOnFailure() throws Exception {
         }
         ensureTaskCompletion(initial.getId());
         AsyncSearchResponse response = getAsyncSearch(initial.getId());
+        assertFalse(response.isRunning());
         assertNotNull(response.getFailure());
         assertTrue(response.isPartial());
         assertThat(response.getSearchResponse().getTotalShards(), equalTo(numShards));

x-pack/plugin/async-search/src/test/java/org/elasticsearch/xpack/search/AsyncSearchIntegTestCase.java

@@ -361,7 +361,7 @@ public int hashCode() {
 
         @Override
         protected boolean doEquals(BlockQueryBuilder other) {
-            return true;
+            return false;
         }
 
         @Override

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/search/action/AsyncSearchResponse.java

@@ -100,7 +100,7 @@ public void writeTo(StreamOutput out) throws IOException {
     }
 
     public AsyncSearchResponse clone(String id) {
-        return new AsyncSearchResponse(id, version, searchResponse, error, isPartial, isRunning, startTimeMillis, expirationTimeMillis);
+        return new AsyncSearchResponse(id, version, searchResponse, error, isPartial, false, startTimeMillis, expirationTimeMillis);
     }
 
     /**