Ensure search contexts are released after tests

These assertions are introduced in elastic#71354
dnhatn · Apr 7, 2021 · ab2b698 · ab2b698
1 parent 6a081e1
commit ab2b698
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 0 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java
@@ -1152,6 +1152,13 @@ public int getActiveContexts() {
         return this.activeReaders.size();
     }
 
+    /**
+     * Returns the number of scroll contexts opened on the node
+     */
+    public int getOpenScrollContexts() {
+        return openScrollContexts.get();
+    }
+
     public ResponseCollectorService getResponseCollectorService() {
         return this.responseCollectorService;
     }

diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESSingleNodeTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESSingleNodeTestCase.java
@@ -41,6 +41,7 @@
 import org.elasticsearch.node.NodeValidationException;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.script.MockScriptService;
+import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.transport.TransportSettings;
@@ -117,6 +118,9 @@ public void setUp() throws Exception {
     public void tearDown() throws Exception {
         logger.trace("[{}#{}]: cleaning up after test", getTestClass().getSimpleName(), getTestName());
         ensureNoInitializingShards();
+        SearchService searchService = getInstanceFromNode(SearchService.class);
+        assertThat(searchService.getActiveContexts(), equalTo(0));
+        assertThat(searchService.getOpenScrollContexts(), equalTo(0));
         super.tearDown();
         assertAcked(
             client().admin().indices().prepareDelete("*")

diff --git a/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java b/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java
@@ -2315,6 +2315,7 @@ public void ensureEstimatedStats() {
     public synchronized void assertAfterTest() throws Exception {
         super.assertAfterTest();
         assertRequestsFinished();
+        assertSearchContextsReleased();
         assertNoInFlightDocsInEngine();
         for (NodeAndClient nodeAndClient : nodes.values()) {
             NodeEnvironment env = nodeAndClient.node().getNodeEnvironment();
@@ -2356,4 +2357,18 @@ public void assertRequestsFinished() {
             }
         }
     }
+
+    private void assertSearchContextsReleased() {
+        for (NodeAndClient nodeAndClient : nodes.values()) {
+            SearchService searchService = getInstance(SearchService.class, nodeAndClient.name);
+            try {
+                assertBusy(() -> {
+                    assertThat(searchService.getActiveContexts(), equalTo(0));
+                    assertThat(searchService.getOpenScrollContexts(), equalTo(0));
+                });
+            } catch (Exception e) {
+                throw new AssertionError("Failed to verify search contexts", e);
+            }
+        }
+    }
 }