Aggs: Fix Filters agg cancellation flaky test

Author: ivancea

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/FiltersCancellationIT.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
+import org.junit.Before;
 
 import java.util.Collection;
 import java.util.List;
@@ -39,6 +40,7 @@
 import static org.elasticsearch.search.aggregations.AggregationBuilders.filters;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
 import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.not;
 
@@ -55,11 +57,12 @@ public class FiltersCancellationIT extends ESIntegTestCase {
 
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
-        return CollectionUtils.appendToCopy(super.nodePlugins(), pausableFieldPluginClass());
+        return CollectionUtils.appendToCopy(super.nodePlugins(), PauseScriptPlugin.class);
     }
 
-    protected Class<? extends Plugin> pausableFieldPluginClass() {
-        return PauseScriptPlugin.class;
+    @Override
+    public Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
+        return Settings.builder().put(super.nodeSettings(nodeOrdinal, otherSettings)).put("thread_pool.search.size", 4).build();
     }
 
     @Override
@@ -99,6 +102,11 @@ public void setupSuiteScopeCluster() throws Exception {
         client().admin().indices().prepareForceMerge(INDEX).setMaxNumSegments(1).get();
     }
 
+    @Before
+    public void reset() {
+        SCRIPT_SEMAPHORE.drainPermits();
+    }
+
     public void testFiltersCountCancellation() throws Exception {
         ensureProperCancellation(
             client().prepareSearch(INDEX)
@@ -129,14 +137,14 @@ public void testFiltersSubAggsCancellation() throws Exception {
 
     private void ensureProperCancellation(SearchRequestBuilder searchRequestBuilder) throws Exception {
         var searchRequestFuture = searchRequestBuilder.setTimeout(TimeValue.timeValueSeconds(1)).execute();
-        assertFalse(searchRequestFuture.isCancelled());
-        assertFalse(searchRequestFuture.isDone());
+        assertThat(searchRequestFuture.isCancelled(), equalTo(false));
+        assertThat(searchRequestFuture.isDone(), equalTo(false));
 
         // Check that there are search tasks running
         assertThat(getSearchTasks(), not(empty()));
 
         // Wait for the script field to get blocked
-        assertBusy(() -> { assertThat(SCRIPT_SEMAPHORE.getQueueLength(), greaterThan(0)); });
+        assertBusy(() -> assertThat(SCRIPT_SEMAPHORE.getQueueLength(), greaterThan(0)));
 
         // Cancel the tasks
         // Warning: Adding a waitForCompletion(true)/execute() here sometimes causes tasks to not get canceled and threads to get stuck
@@ -146,8 +154,8 @@ private void ensureProperCancellation(SearchRequestBuilder searchRequestBuilder)
 
         // Ensure the search request finished and that there are no more search tasks
         assertBusy(() -> {
-            assertTrue(searchRequestFuture.isDone());
-            assertThat(getSearchTasks(), empty());
+            assertThat("Search request didn't finish", searchRequestFuture.isDone(), equalTo(true));
+            assertThat("There are dangling search tasks", getSearchTasks(), empty());
         });
     }