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.

1 files changed, 100 insertions, 0 deletions

diff --git a/core/src/test/java/org/elasticsearch/search/SearchServiceTests.java b/core/src/test/java/org/elasticsearch/search/SearchServiceTests.java
index a222ff5272..2ac00b541a 100644
--- a/core/src/test/java/org/elasticsearch/search/SearchServiceTests.java
+++ b/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;
+        }
+    }
 }