Merge pull request #9595 from jpountz/test/context_closed

Tests: Assert that we do not leak SearchContexts.

Close #9595

Author: jpountz

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

@@ -186,7 +186,7 @@ public Node(Settings preparedSettings, boolean loadConfigSettings) throws Elasti
             }
             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 @@
  */
 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 void afterIndexDeleted(Index index, @IndexSettings Settings indexSettings
         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 @@ protected void doStop() throws ElasticsearchException {
         for (final SearchContext context : activeContexts.values()) {
             freeContext(context.id());
         }
-        activeContexts.clear();
     }
 
     @Override
@@ -525,7 +533,7 @@ final SearchContext createAndPutContext(ShardSearchRequest request) throws Elast
         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 @@ private void freeAllContextForIndex(Index index) {
 
 
     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 void testScanScrollWithShardExceptions() throws Exception {
             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 final void before() throws Exception {
 
     @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.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 void ensureAllArraysReleased() throws Exception {
         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.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.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 @@ private static Settings getRandomNodeSettings(long seed) {
             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();
+    }
+
+}