Tests: Assert that we do not leak SearchContexts.

Even if there is a background thread that periodically closes search contexts
that seem unused (every minute by default), it is important to close search
contexts as soon as possible in order to not keep unnecessary open files or
to prevent segments from being deleted.

This check would help ensure that refactorings of the SearchContext management
like #9296 are correct.

src/main/java/org/elasticsearch/node/Node.java

@@ -186,7 +186,7 @@ public class Node implements Releasable {
             }
             modules.add(new RiversModule(settings));
             modules.add(new IndicesModule(settings));
-            modules.add(new SearchModule());
+            modules.add(new SearchModule(settings));
             modules.add(new ActionModule(false));
             modules.add(new MonitorModule(settings));
             modules.add(new GatewayModule());

src/main/java/org/elasticsearch/search/DefaultSearchServiceModule.java

@@ -0,0 +1,31 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search;
+
+import org.elasticsearch.common.inject.AbstractModule;
+
+public class DefaultSearchServiceModule extends AbstractModule {
+
+    @Override
+    protected void configure() {
+        bind(SearchService.class).asEagerSingleton();
+    }
+
+}

src/main/java/org/elasticsearch/search/SearchModule.java

@@ -20,9 +20,11 @@
 package org.elasticsearch.search;
 
 import com.google.common.collect.ImmutableList;
+
 import org.elasticsearch.common.inject.AbstractModule;
 import org.elasticsearch.common.inject.Module;
 import org.elasticsearch.common.inject.SpawnModules;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.query.functionscore.FunctionScoreModule;
 import org.elasticsearch.index.search.morelikethis.MoreLikeThisFetchService;
 import org.elasticsearch.search.action.SearchServiceTransportAction;
@@ -47,16 +49,27 @@ import org.elasticsearch.search.suggest.SuggestModule;
  */
 public class SearchModule extends AbstractModule implements SpawnModules {
 
+    private final Settings settings;
+
+    public SearchModule(Settings settings) {
+        this.settings = settings;
+    }
+
     @Override
     public Iterable<? extends Module> spawnModules() {
-        return ImmutableList.of(new TransportSearchModule(), new HighlightModule(), new SuggestModule(), new FunctionScoreModule(), new AggregationModule());
+        return ImmutableList.of(
+                new SearchServiceModule(settings),
+                new TransportSearchModule(),
+                new HighlightModule(),
+                new SuggestModule(),
+                new FunctionScoreModule(),
+                new AggregationModule());
     }
 
     @Override
     protected void configure() {
         bind(DfsPhase.class).asEagerSingleton();
         bind(QueryPhase.class).asEagerSingleton();
-        bind(SearchService.class).asEagerSingleton();
         bind(SearchPhaseController.class).asEagerSingleton();
 
         bind(FetchPhase.class).asEagerSingleton();

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

@@ -182,6 +182,15 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
         this.indicesWarmer.addListener(new SearchWarmer());
     }
 
+    protected void putContext(SearchContext context) {
+        final SearchContext previous = activeContexts.put(context.id(), context);
+        assert previous == null;
+    }
+
+    protected SearchContext removeContext(long id) {
+        return activeContexts.remove(id);
+    }
+
     @Override
     protected void doStart() throws ElasticsearchException {
     }
@@ -191,7 +200,6 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
         for (final SearchContext context : activeContexts.values()) {
             freeContext(context.id());
         }
-        activeContexts.clear();
     }
 
     @Override
@@ -525,7 +533,7 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
         SearchContext context = createContext(request, null);
         boolean success = false;
         try {
-            activeContexts.put(context.id(), context);
+            putContext(context);
             context.indexShard().searchService().onNewContext(context);
             success = true;
             return context;
@@ -590,7 +598,7 @@ public class SearchService extends AbstractLifecycleComponent<SearchService> {
 
 
     public boolean freeContext(long id) {
-        final SearchContext context = activeContexts.remove(id);
+        final SearchContext context = removeContext(id);
         if (context != null) {
             try {
                 context.indexShard().searchService().onFreeContext(context);

src/main/java/org/elasticsearch/search/SearchServiceModule.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search;
+
+import com.google.common.collect.ImmutableList;
+
+import org.elasticsearch.common.inject.AbstractModule;
+import org.elasticsearch.common.inject.Module;
+import org.elasticsearch.common.inject.SpawnModules;
+import org.elasticsearch.common.settings.Settings;
+
+import static org.elasticsearch.common.inject.Modules.createModule;
+
+public class SearchServiceModule extends AbstractModule implements SpawnModules {
+
+    public static final String IMPL = "search.service_impl";
+
+    private final Settings settings;
+
+    public SearchServiceModule(Settings settings) {
+        this.settings = settings;
+    }
+
+    @Override
+    protected void configure() {
+    }
+
+    @Override
+    public Iterable<? extends Module> spawnModules() {
+        return ImmutableList.of(createModule(settings.getAsClass(IMPL, DefaultSearchServiceModule.class), settings));
+    }
+}

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

@@ -106,6 +106,8 @@ public class SearchScrollWithFailingNodesTests extends ElasticsearchIntegrationT
             assertThat(searchResponse.getSuccessfulShards(), equalTo(numberOfSuccessfulShards));
         } while (searchResponse.getHits().hits().length > 0);
         assertThat(numHits, greaterThan(0l));
+
+        clearScroll(searchResponse.getScrollId());
     }
 
 }

src/test/java/org/elasticsearch/test/ElasticsearchIntegrationTest.java

@@ -1838,6 +1838,10 @@ public abstract class ElasticsearchIntegrationTest extends ElasticsearchTestCase
 
     @After
     public final void after() throws Exception {
+        // Deleting indices is going to clear search contexts implicitely so we
+        // need to check that there are no more in-flight search contexts before
+        // we remove indices
+        super.ensureAllSearchContextsReleased();
         if (runTestScopeLifecycle()) {
             afterInternal(false);
         }

src/test/java/org/elasticsearch/test/ElasticsearchTestCase.java

@@ -46,6 +46,7 @@ import org.elasticsearch.env.NodeEnvironment;
 import org.elasticsearch.test.cache.recycler.MockBigArrays;
 import org.elasticsearch.test.cache.recycler.MockPageCacheRecycler;
 import org.elasticsearch.test.junit.listeners.LoggingListener;
+import org.elasticsearch.test.search.MockSearchService;
 import org.elasticsearch.test.store.MockDirectoryHelper;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.junit.*;
@@ -212,6 +213,16 @@ public abstract class ElasticsearchTestCase extends AbstractRandomizedTest {
         MockBigArrays.ensureAllArraysAreReleased();
     }
 
+    @After
+    public void ensureAllSearchContextsReleased() throws Exception {
+        assertBusy(new Runnable() {
+            @Override
+            public void run() {
+                MockSearchService.assertNoInFLightContext();
+            }
+        });
+    }
+
     public static boolean hasUnclosedWrapper() {
         for (MockDirectoryWrapper w : MockDirectoryHelper.wrappers) {
             if (w.isOpen()) {

src/test/java/org/elasticsearch/test/InternalTestCluster.java

@@ -29,6 +29,7 @@ import com.google.common.collect.*;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
+
 import org.apache.lucene.util.AbstractRandomizedTest;
 import org.apache.lucene.util.IOUtils;
 import org.elasticsearch.ElasticsearchException;
@@ -86,10 +87,12 @@ import org.elasticsearch.node.Node;
 import org.elasticsearch.node.service.NodeService;
 import org.elasticsearch.plugins.PluginsService;
 import org.elasticsearch.search.SearchService;
+import org.elasticsearch.search.SearchServiceModule;
 import org.elasticsearch.test.cache.recycler.MockBigArraysModule;
 import org.elasticsearch.test.cache.recycler.MockPageCacheRecyclerModule;
 import org.elasticsearch.test.disruption.ServiceDisruptionScheme;
 import org.elasticsearch.test.engine.MockEngineFactory;
+import org.elasticsearch.test.search.MockSearchServiceModule;
 import org.elasticsearch.test.store.MockFSIndexStoreModule;
 import org.elasticsearch.test.transport.AssertingLocalTransport;
 import org.elasticsearch.test.transport.MockTransportService;
@@ -360,6 +363,7 @@ public final class InternalTestCluster extends TestCluster {
             builder.put(PageCacheRecyclerModule.CACHE_IMPL, MockPageCacheRecyclerModule.class.getName());
             builder.put(BigArraysModule.IMPL, MockBigArraysModule.class.getName());
             builder.put(TransportModule.TRANSPORT_SERVICE_TYPE_KEY, MockTransportService.class.getName());
+            builder.put(SearchServiceModule.IMPL, MockSearchServiceModule.class.getName());
         }
         if (isLocalTransportConfigured()) {
             builder.put(TransportModule.TRANSPORT_TYPE_KEY, AssertingLocalTransport.class.getName());

src/test/java/org/elasticsearch/test/search/MockSearchService.java

@@ -0,0 +1,76 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.test.search;
+
+import org.elasticsearch.cache.recycler.PageCacheRecycler;
+import org.elasticsearch.cluster.ClusterService;
+import org.elasticsearch.common.inject.Inject;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.indices.IndicesService;
+import org.elasticsearch.indices.IndicesWarmer;
+import org.elasticsearch.indices.cache.query.IndicesQueryCache;
+import org.elasticsearch.script.ScriptService;
+import org.elasticsearch.search.SearchService;
+import org.elasticsearch.search.dfs.DfsPhase;
+import org.elasticsearch.search.fetch.FetchPhase;
+import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.query.QueryPhase;
+import org.elasticsearch.threadpool.ThreadPool;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class MockSearchService extends SearchService {
+
+    private static final Map<SearchContext, Throwable> ACTIVE_SEARCH_CONTEXTS = new ConcurrentHashMap<>();
+
+    /** Throw an {@link AssertionError} if there are still in-flight contexts. */
+    public static void assertNoInFLightContext() {
+        final Map<SearchContext, Throwable> copy = new HashMap<>(ACTIVE_SEARCH_CONTEXTS);
+        if (copy.isEmpty() == false) {
+            throw new AssertionError("There are still " + copy.size() + " in-flight contexts", copy.values().iterator().next());
+        }
+    }
+
+    @Inject
+    public MockSearchService(Settings settings, ClusterService clusterService, IndicesService indicesService, IndicesWarmer indicesWarmer,
+            ThreadPool threadPool, ScriptService scriptService, PageCacheRecycler pageCacheRecycler, BigArrays bigArrays,
+            DfsPhase dfsPhase, QueryPhase queryPhase, FetchPhase fetchPhase, IndicesQueryCache indicesQueryCache) {
+        super(settings, clusterService, indicesService, indicesWarmer, threadPool, scriptService, pageCacheRecycler, bigArrays, dfsPhase,
+                queryPhase, fetchPhase, indicesQueryCache);
+    }
+ 
+    @Override
+    protected void putContext(SearchContext context) {
+        super.putContext(context);
+        ACTIVE_SEARCH_CONTEXTS.put(context, new RuntimeException());
+    }
+
+    @Override
+    protected SearchContext removeContext(long id) {
+        final SearchContext removed = super.removeContext(id);
+        if (removed != null) {
+            ACTIVE_SEARCH_CONTEXTS.remove(removed);
+        }
+        return removed;
+    }
+}

src/test/java/org/elasticsearch/test/search/MockSearchServiceModule.java

@@ -0,0 +1,32 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.test.search;
+
+import org.elasticsearch.common.inject.AbstractModule;
+import org.elasticsearch.search.SearchService;
+
+public class MockSearchServiceModule extends AbstractModule {
+
+    @Override
+    protected void configure() {
+        bind(SearchService.class).to(MockSearchService.class).asEagerSingleton();
+    }
+
+}