Safely ensure search contexts are release in ESSingleNodeTestCase (#111153)

The transport action for freeing contexts is forking to GENERIC nowadays and executed in a fire-and-forget manner in production code. When it was running on the direct executor, no busy wait was needed for a single node since the freeing was fully synchronous inside the searches. With forking, the waiting on free tasks does not help because the task is created inside the forked execution for a local node execution, so we need to busy wait like we do for multi-node clusters here too. -> dry up the code and use the same method for both.

Author: original-brownbear

test/framework/src/main/java/org/elasticsearch/test/ESSingleNodeTestCase.java

@@ -133,9 +133,7 @@ public void tearDown() throws Exception {
  ensureNoInitializingShards();
  ensureAllFreeContextActionsAreConsumed();
 
- SearchService searchService = getInstanceFromNode(SearchService.class);
- assertThat(searchService.getActiveContexts(), equalTo(0));
- assertThat(searchService.getOpenScrollContexts(), equalTo(0));
+ ensureAllContextsReleased(getInstanceFromNode(SearchService.class));
  super.tearDown();
  var deleteDataStreamsRequest = new DeleteDataStreamAction.Request("*");
  deleteDataStreamsRequest.indicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED_HIDDEN);

test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java

@@ -115,6 +115,7 @@
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.ScriptType;
 import org.elasticsearch.search.MockSearchService;
+import org.elasticsearch.search.SearchService;
 import org.elasticsearch.test.junit.listeners.LoggingListener;
 import org.elasticsearch.test.junit.listeners.ReproduceInfoPrinter;
 import org.elasticsearch.threadpool.ExecutorBuilder;
@@ -208,6 +209,7 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.startsWith;
+import static org.junit.Assert.assertThat;
 
 /**
  * Base testcase for randomized unit testing with Elasticsearch
@@ -2515,4 +2517,15 @@ public static void runInParallel(int numberOfTasks, IntConsumer taskFactory) thr
  throw new AssertionError(e);
  }
  }
+
+ public static void ensureAllContextsReleased(SearchService searchService) {
+ try {
+ assertBusy(() -> {
+ assertThat(searchService.getActiveContexts(), equalTo(0));
+ assertThat(searchService.getOpenScrollContexts(), equalTo(0));
+ });
+ } catch (Exception e) {
+ throw new AssertionError("Failed to verify search contexts", e);
+ }
+ }
 }

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

@@ -2545,15 +2545,7 @@ 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);
- }
+ ESTestCase.ensureAllContextsReleased(getInstance(SearchService.class, nodeAndClient.name));
  }
  }