From 3f2a94969c20ca757f81362ad24ff40d83ac7f56 Mon Sep 17 00:00:00 2001
From: Armin Braun <me@obrown.io>
Date: Mon, 22 Jul 2024 13:13:46 +0200
Subject: [PATCH] Safely ensure search contexts are release in
 ESSingleNodeTestCase

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.
---
 .../elasticsearch/test/ESSingleNodeTestCase.java    |  4 +---
 .../java/org/elasticsearch/test/ESTestCase.java     | 13 +++++++++++++
 .../org/elasticsearch/test/InternalTestCluster.java | 10 +---------
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESSingleNodeTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESSingleNodeTestCase.java
index 7fdc5765a90e8..a538c39704a73 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/ESSingleNodeTestCase.java
+++ b/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);
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
index dc8a1dc29d233..dfa29fd26367a 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
+++ b/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);
+        }
+    }
 }
diff --git a/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java b/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java
index af37fb6feefbd..0b69245177c7a 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java
+++ b/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));
         }
     }