Close SearchContext if query rewrite failed

If a query failed to be rewritten and throws an exception, the SearchContext is not properly closed, skewing the ref count on the underlying Store.

core/src/main/java/org/elasticsearch/index/query/WrapperQueryBuilder.java

@@ -126,7 +126,7 @@ public class WrapperQueryBuilder extends AbstractQueryBuilder<WrapperQueryBuilde
         }
         String fieldName = parser.currentName();
         if (! parseContext.getParseFieldMatcher().match(fieldName, QUERY_FIELD)) {
-            throw new ParsingException(parser.getTokenLocation(), "[wrapper] query malformed, expected `query` but was" + fieldName);
+            throw new ParsingException(parser.getTokenLocation(), "[wrapper] query malformed, expected `query` but was " + fieldName);
         }
         parser.nextToken();
 

core/src/main/java/org/elasticsearch/rest/action/search/RestMultiSearchAction.java

@@ -157,7 +157,7 @@ public class RestMultiSearchAction extends BaseRestHandler {
                         Object value = entry.getValue();
                         if ("index".equals(entry.getKey()) || "indices".equals(entry.getKey())) {
                             if (!allowExplicitIndex) {
-                                throw new IllegalArgumentException("explicit index in multi percolate is not allowed");
+                                throw new IllegalArgumentException("explicit index in multi search is not allowed");
                             }
                             searchRequest.indices(nodeStringArrayValue(value));
                         } else if ("type".equals(entry.getKey()) || "types".equals(entry.getKey())) {

core/src/main/java/org/elasticsearch/search/SearchService.java

@@ -546,14 +546,14 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> imp
                 indexShard, scriptService, bigArrays, threadPool.estimatedTimeInMillisCounter(), parseFieldMatcher,
                 defaultSearchTimeout, fetchPhase);
         SearchContext.setCurrent(context);
-        request.rewrite(context.getQueryShardContext());
-        // reset that we have used nowInMillis from the context since it may
-        // have been rewritten so its no longer in the query and the request can
-        // be cached. If it is still present in the request (e.g. in a range
-        // aggregation) it will still be caught when the aggregation is
-        // evaluated.
-        context.resetNowInMillisUsed();
         try {
+            request.rewrite(context.getQueryShardContext());
+            // reset that we have used nowInMillis from the context since it may
+            // have been rewritten so its no longer in the query and the request can
+            // be cached. If it is still present in the request (e.g. in a range
+            // aggregation) it will still be caught when the aggregation is
+            // evaluated.
+            context.resetNowInMillisUsed();
             if (request.scroll() != null) {
                 context.scrollContext(new ScrollContext());
                 context.scrollContext().scroll = request.scroll();

core/src/test/java/org/elasticsearch/search/SearchServiceTests.java

@@ -19,9 +19,25 @@
 package org.elasticsearch.search;
 
 
+import org.apache.lucene.search.Query;
+import org.elasticsearch.action.search.SearchPhaseExecutionException;
 import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.index.IndexService;
+import org.elasticsearch.index.query.AbstractQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryRewriteContext;
+import org.elasticsearch.index.query.QueryShardContext;
+import org.elasticsearch.index.shard.IndexShard;
+import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 
+import java.io.IOException;
+import java.util.Collection;
 import java.util.concurrent.ExecutionException;
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
@@ -35,6 +51,11 @@ public class SearchServiceTests extends ESSingleNodeTestCase {
         return true;
     }
 
+    @Override
+    protected Collection<Class<? extends Plugin>> getPlugins() {
+        return pluginList(FailOnRewriteQueryPlugin.class);
+    }
+
     public void testClearOnClose() throws ExecutionException, InterruptedException {
         createIndex("index");
         client().prepareIndex("index", "type", "1").setSource("field", "value").setRefresh(true).get();
@@ -70,4 +91,83 @@ public class SearchServiceTests extends ESSingleNodeTestCase {
         assertAcked(client().admin().indices().prepareDelete("index"));
         assertEquals(0, service.getActiveContexts());
     }
+
+    public void testCloseSearchContextOnRewriteException() {
+        createIndex("index");
+        client().prepareIndex("index", "type", "1").setSource("field", "value").setRefresh(true).get();
+
+        SearchService service = getInstanceFromNode(SearchService.class);
+        IndicesService indicesService = getInstanceFromNode(IndicesService.class);
+        IndexService indexService = indicesService.indexServiceSafe(resolveIndex("index"));
+        IndexShard indexShard = indexService.getShard(0);
+
+        final int activeContexts = service.getActiveContexts();
+        final int activeRefs = indexShard.store().refCount();
+        expectThrows(SearchPhaseExecutionException.class, () ->
+                client().prepareSearch("index").setQuery(new FailOnRewriteQueryBuilder()).get());
+        assertEquals(activeContexts, service.getActiveContexts());
+        assertEquals(activeRefs, indexShard.store().refCount());
+    }
+
+    public static class FailOnRewriteQueryPlugin extends Plugin {
+
+        @Override
+        public String name() {
+            return FailOnRewriteQueryPlugin.class.getSimpleName();
+        }
+
+        @Override
+        public String description() {
+            return "This plugin registers a query that always fails at rewrite phase";
+        }
+
+        public void onModule(SearchModule module) {
+            module.registerQuery(FailOnRewriteQueryBuilder::new, parseContext -> {
+                throw new UnsupportedOperationException("No query parser for this plugin");
+            }, new ParseField("fail_on_rewrite_query"));
+        }
+    }
+
+    public static class FailOnRewriteQueryBuilder extends AbstractQueryBuilder<FailOnRewriteQueryBuilder> {
+
+        public FailOnRewriteQueryBuilder(StreamInput in) throws IOException {
+            super(in);
+        }
+
+        public FailOnRewriteQueryBuilder() {
+        }
+
+        @Override
+        protected QueryBuilder doRewrite(QueryRewriteContext queryShardContext) throws IOException {
+            throw new IllegalStateException("Fail on rewrite phase");
+        }
+
+        @Override
+        protected void doWriteTo(StreamOutput out) throws IOException {
+        }
+
+        @Override
+        protected void doXContent(XContentBuilder builder, Params params) throws IOException {
+        }
+
+        @Override
+        protected Query doToQuery(QueryShardContext context) throws IOException {
+            return null;
+        }
+
+        @Override
+        protected boolean doEquals(FailOnRewriteQueryBuilder other) {
+            return false;
+        }
+
+        @Override
+        protected int doHashCode() {
+            return 0;
+        }
+
+        @Override
+        public String getWriteableName() {
+            return null;
+        }
+    }
 }