Improve 404 on missing scroll id

This relates to #6040, the fix is twofold, first, not handling missing context specifically in the search code, but behave the same as we do in non scroll search, where if all the shards failed, raise an exception. The second is to apply this logic in both scroll cases.

src/main/java/org/elasticsearch/action/search/type/TransportSearchScrollQueryAndFetchAction.java

@@ -181,7 +181,11 @@ public class TransportSearchScrollQueryAndFetchAction extends AbstractComponent
             addShardFailure(shardIndex, new ShardSearchFailure(t));
             successfulOps.decrementAndGet();
             if (counter.decrementAndGet() == 0) {
-                finishHim();
+                if (successfulOps.get() == 0) {
+                    listener.onFailure(new SearchPhaseExecutionException("query_fetch", "all shards failed", buildShardFailures()));
+                } else {
+                    finishHim();
+                }
             }
         }
 

src/main/java/org/elasticsearch/action/search/type/TransportSearchScrollQueryThenFetchAction.java

@@ -174,12 +174,7 @@ public class TransportSearchScrollQueryThenFetchAction extends AbstractComponent
 
                 @Override
                 public void onFailure(Throwable t) {
-                    Throwable cause = ExceptionsHelper.unwrapCause(t);
-                    if (cause instanceof SearchContextMissingException) {
-                        listener.onFailure(t);
-                    } else {
-                        onQueryPhaseFailure(shardIndex, counter, searchId, t);
-                    }
+                    onQueryPhaseFailure(shardIndex, counter, searchId, t);
                 }
             });
         }
@@ -191,10 +186,14 @@ public class TransportSearchScrollQueryThenFetchAction extends AbstractComponent
             addShardFailure(shardIndex, new ShardSearchFailure(t));
             successfulOps.decrementAndGet();
             if (counter.decrementAndGet() == 0) {
-                try {
-                    executeFetchPhase();
-                } catch (Throwable e) {
-                    listener.onFailure(new SearchPhaseExecutionException("query", "Fetch failed", e, null));
+                if (successfulOps.get() == 0) {
+                    listener.onFailure(new SearchPhaseExecutionException("query", "all shards failed", buildShardFailures()));
+                } else {
+                    try {
+                        executeFetchPhase();
+                    } catch (Throwable e) {
+                        listener.onFailure(new SearchPhaseExecutionException("query", "Fetch failed", e, null));
+                    }
                 }
             }
         }

src/test/java/org/elasticsearch/search/scroll/SearchScrollTests.java

@@ -291,8 +291,8 @@ public class SearchScrollTests extends ElasticsearchIntegrationTest {
                 .execute().actionGet();
         assertThat(clearResponse.isSucceeded(), equalTo(true));
 
-        assertThrows(client().prepareSearchScroll(searchResponse1.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), SearchContextMissingException.class);
-        assertThrows(client().prepareSearchScroll(searchResponse2.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), SearchContextMissingException.class);
+        assertThrows(client().prepareSearchScroll(searchResponse1.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), RestStatus.NOT_FOUND);
+        assertThrows(client().prepareSearchScroll(searchResponse2.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), RestStatus.NOT_FOUND);
     }
 
     @Test
@@ -397,8 +397,8 @@ public class SearchScrollTests extends ElasticsearchIntegrationTest {
                 .execute().actionGet();
         assertThat(clearResponse.isSucceeded(), equalTo(true));
 
-        assertThrows(cluster().transportClient().prepareSearchScroll(searchResponse1.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), SearchContextMissingException.class);
-        assertThrows(cluster().transportClient().prepareSearchScroll(searchResponse2.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), SearchContextMissingException.class);
+        assertThrows(cluster().transportClient().prepareSearchScroll(searchResponse1.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), RestStatus.NOT_FOUND);
+        assertThrows(cluster().transportClient().prepareSearchScroll(searchResponse2.getScrollId()).setScroll(TimeValue.timeValueMinutes(2)), RestStatus.NOT_FOUND);
     }
 
     @Test
@@ -436,17 +436,12 @@ public class SearchScrollTests extends ElasticsearchIntegrationTest {
         client().prepareIndex("index", "type", "1").setSource("field", "value").execute().get();
         refresh();
 
-        try {
-            SearchResponse searchResponse = client().prepareSearch("index").setSize(1).setScroll("1m").get();
-            assertThat(searchResponse.getScrollId(), is(notNullValue()));
+        SearchResponse searchResponse = client().prepareSearch("index").setSize(1).setScroll("1m").get();
+        assertThat(searchResponse.getScrollId(), is(notNullValue()));
 
-            ClearScrollResponse clearScrollResponse = client().prepareClearScroll().addScrollId(searchResponse.getScrollId()).get();
-            assertThat(clearScrollResponse.isSucceeded(), is(true));
+        ClearScrollResponse clearScrollResponse = client().prepareClearScroll().addScrollId(searchResponse.getScrollId()).get();
+        assertThat(clearScrollResponse.isSucceeded(), is(true));
 
-            cluster().transportClient().prepareSearchScroll(searchResponse.getScrollId()).get();
-            fail("Expected exception to happen due to non-existing scroll id");
-        } catch (ElasticsearchException e) {
-            assertThat(e.status(), is(RestStatus.NOT_FOUND));
-        }
+        assertThrows(cluster().transportClient().prepareSearchScroll(searchResponse.getScrollId()), RestStatus.NOT_FOUND);
     }
 }

src/test/java/org/elasticsearch/test/hamcrest/ElasticsearchAssertions.java

@@ -23,6 +23,7 @@ import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ElasticsearchIllegalStateException;
+import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionFuture;
 import org.elasticsearch.action.ActionRequest;
@@ -413,6 +414,35 @@ public class ElasticsearchAssertions {
         }
     }
 
+    public static <E extends Throwable> void assertThrows(ActionRequestBuilder<?, ?, ?> builder, RestStatus status) {
+        assertThrows(builder.execute(), status);
+    }
+
+    public static <E extends Throwable> void assertThrows(ActionRequestBuilder<?, ?, ?> builder, RestStatus status, String extraInfo) {
+        assertThrows(builder.execute(), status, extraInfo);
+    }
+
+    public static <E extends Throwable> void assertThrows(ActionFuture future, RestStatus status) {
+        assertThrows(future, status, null);
+    }
+
+    public static void assertThrows(ActionFuture future, RestStatus status, String extraInfo) {
+        boolean fail = false;
+        extraInfo = extraInfo == null || extraInfo.isEmpty() ? "" : extraInfo + ": ";
+        extraInfo += "expected a " + status + " status exception to be thrown";
+
+        try {
+            future.actionGet();
+            fail = true;
+        } catch (Throwable e) {
+            assertThat(extraInfo, ExceptionsHelper.status(e), equalTo(status));
+        }
+        // has to be outside catch clause to get a proper message
+        if (fail) {
+            throw new AssertionError(extraInfo);
+        }
+    }
+
     private static BytesReference serialize(Version version, Streamable streamable) throws IOException {
         BytesStreamOutput output = new BytesStreamOutput();
         output.setVersion(version);