Dry up search error trace ITs (#122138)

This logic will need a bit of adjustment for bulk query execution.
Lets dry it up before so we don't have to copy and paste the fix which
will be a couple lines.

qa/smoke-test-http/src/javaRestTest/java/org/elasticsearch/http/SearchErrorTraceIT.java

@@ -11,43 +11,26 @@ package org.elasticsearch.http;
 
 import org.apache.http.entity.ContentType;
 import org.apache.http.nio.entity.NByteArrayEntity;
-import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.search.MultiSearchRequest;
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.client.Request;
+import org.elasticsearch.search.ErrorTraceHelper;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
-import org.elasticsearch.transport.TransportMessageListener;
-import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xcontent.XContentType;
 import org.junit.Before;
 
 import java.io.IOException;
 import java.nio.charset.Charset;
-import java.util.Optional;
+import java.util.function.BooleanSupplier;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.elasticsearch.index.query.QueryBuilders.simpleQueryStringQuery;
 
 public class SearchErrorTraceIT extends HttpSmokeTestCase {
-    private AtomicBoolean hasStackTrace;
+    private BooleanSupplier hasStackTrace;
 
     @Before
-    private void setupMessageListener() {
+    public void setupMessageListener() {
-        internalCluster().getDataNodeInstances(TransportService.class).forEach(ts -> {
+        hasStackTrace = ErrorTraceHelper.setupErrorTraceListener(internalCluster());
-            ts.addMessageListener(new TransportMessageListener() {
-                @Override
-                public void onResponseSent(long requestId, String action, Exception error) {
-                    TransportMessageListener.super.onResponseSent(requestId, action, error);
-                    if (action.startsWith("indices:data/read/search")) {
-                        Optional<Throwable> throwable = ExceptionsHelper.unwrapCausesAndSuppressed(
-                            error,
-                            t -> t.getStackTrace().length > 0
-                        );
-                        hasStackTrace.set(throwable.isPresent());
-                    }
-                }
-            });
-        });
     }
 
     private void setupIndexWithDocs() {
@@ -61,7 +44,6 @@ public class SearchErrorTraceIT extends HttpSmokeTestCase {
     }
 
     public void testSearchFailingQueryErrorTraceDefault() throws IOException {
-        hasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_search");
@@ -76,11 +58,10 @@ public class SearchErrorTraceIT extends HttpSmokeTestCase {
             }
             """);
         getRestClient().performRequest(searchRequest);
-        assertFalse(hasStackTrace.get());
+        assertFalse(hasStackTrace.getAsBoolean());
     }
 
     public void testSearchFailingQueryErrorTraceTrue() throws IOException {
-        hasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_search");
@@ -96,11 +77,10 @@ public class SearchErrorTraceIT extends HttpSmokeTestCase {
             """);
         searchRequest.addParameter("error_trace", "true");
         getRestClient().performRequest(searchRequest);
-        assertTrue(hasStackTrace.get());
+        assertTrue(hasStackTrace.getAsBoolean());
     }
 
     public void testSearchFailingQueryErrorTraceFalse() throws IOException {
-        hasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_search");
@@ -116,11 +96,10 @@ public class SearchErrorTraceIT extends HttpSmokeTestCase {
             """);
         searchRequest.addParameter("error_trace", "false");
         getRestClient().performRequest(searchRequest);
-        assertFalse(hasStackTrace.get());
+        assertFalse(hasStackTrace.getAsBoolean());
     }
 
     public void testMultiSearchFailingQueryErrorTraceDefault() throws IOException {
-        hasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         XContentType contentType = XContentType.JSON;
@@ -133,11 +112,10 @@ public class SearchErrorTraceIT extends HttpSmokeTestCase {
             new NByteArrayEntity(requestBody, ContentType.create(contentType.mediaTypeWithoutParameters(), (Charset) null))
         );
         getRestClient().performRequest(searchRequest);
-        assertFalse(hasStackTrace.get());
+        assertFalse(hasStackTrace.getAsBoolean());
     }
 
     public void testMultiSearchFailingQueryErrorTraceTrue() throws IOException {
-        hasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         XContentType contentType = XContentType.JSON;
@@ -151,11 +129,10 @@ public class SearchErrorTraceIT extends HttpSmokeTestCase {
         );
         searchRequest.addParameter("error_trace", "true");
         getRestClient().performRequest(searchRequest);
-        assertTrue(hasStackTrace.get());
+        assertTrue(hasStackTrace.getAsBoolean());
     }
 
     public void testMultiSearchFailingQueryErrorTraceFalse() throws IOException {
-        hasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         XContentType contentType = XContentType.JSON;
@@ -170,6 +147,6 @@ public class SearchErrorTraceIT extends HttpSmokeTestCase {
         searchRequest.addParameter("error_trace", "false");
         getRestClient().performRequest(searchRequest);
 
-        assertFalse(hasStackTrace.get());
+        assertFalse(hasStackTrace.getAsBoolean());
     }
 }

test/framework/src/main/java/org/elasticsearch/search/ErrorTraceHelper.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search;
+
+import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.test.InternalTestCluster;
+import org.elasticsearch.transport.TransportMessageListener;
+import org.elasticsearch.transport.TransportService;
+
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.BooleanSupplier;
+
+/**
+ * Utilities around testing the `error_trace` message header in search.
+ */
+public enum ErrorTraceHelper {
+    ;
+
+    public static BooleanSupplier setupErrorTraceListener(InternalTestCluster internalCluster) {
+        final AtomicBoolean transportMessageHasStackTrace = new AtomicBoolean(false);
+        internalCluster.getDataNodeInstances(TransportService.class).forEach(ts -> ts.addMessageListener(new TransportMessageListener() {
+            @Override
+            public void onResponseSent(long requestId, String action, Exception error) {
+                TransportMessageListener.super.onResponseSent(requestId, action, error);
+                if (action.startsWith("indices:data/read/search")) {
+                    Optional<Throwable> throwable = ExceptionsHelper.unwrapCausesAndSuppressed(error, t -> t.getStackTrace().length > 0);
+                    transportMessageHasStackTrace.set(throwable.isPresent());
+                }
+            }
+        }));
+        return transportMessageHasStackTrace::get;
+    }
+}

x-pack/plugin/async-search/src/internalClusterTest/java/org/elasticsearch/xpack/search/AsyncSearchErrorTraceIT.java

@@ -7,15 +7,13 @@
 
 package org.elasticsearch.xpack.search;
 
-import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.search.ErrorTraceHelper;
 import org.elasticsearch.test.ESIntegTestCase;
-import org.elasticsearch.transport.TransportMessageListener;
-import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xcontent.XContentType;
 import org.junit.Before;
 
@@ -23,8 +21,7 @@ import java.io.IOException;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
+import java.util.function.BooleanSupplier;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
 
@@ -38,25 +35,11 @@ public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
         return List.of(AsyncSearch.class);
     }
 
-    private AtomicBoolean transportMessageHasStackTrace;
+    private BooleanSupplier transportMessageHasStackTrace;
 
     @Before
-    private void setupMessageListener() {
+    public void setupMessageListener() {
-        internalCluster().getDataNodeInstances(TransportService.class).forEach(ts -> {
+        transportMessageHasStackTrace = ErrorTraceHelper.setupErrorTraceListener(internalCluster());
-            ts.addMessageListener(new TransportMessageListener() {
-                @Override
-                public void onResponseSent(long requestId, String action, Exception error) {
-                    TransportMessageListener.super.onResponseSent(requestId, action, error);
-                    if (action.startsWith("indices:data/read/search")) {
-                        Optional<Throwable> throwable = ExceptionsHelper.unwrapCausesAndSuppressed(
-                            error,
-                            t -> t.getStackTrace().length > 0
-                        );
-                        transportMessageHasStackTrace.set(throwable.isPresent());
-                    }
-                }
-            });
-        });
     }
 
     private void setupIndexWithDocs() {
@@ -70,7 +53,6 @@ public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
     }
 
     public void testAsyncSearchFailingQueryErrorTraceDefault() throws IOException, InterruptedException {
-        transportMessageHasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_async_search");
@@ -93,11 +75,10 @@ public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
             responseEntity = performRequestAndGetResponseEntityAfterDelay(request, TimeValue.timeValueSeconds(1L));
         }
         // check that the stack trace was not sent from the data node to the coordinating node
-        assertFalse(transportMessageHasStackTrace.get());
+        assertFalse(transportMessageHasStackTrace.getAsBoolean());
     }
 
     public void testAsyncSearchFailingQueryErrorTraceTrue() throws IOException, InterruptedException {
-        transportMessageHasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_async_search");
@@ -122,11 +103,10 @@ public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
             responseEntity = performRequestAndGetResponseEntityAfterDelay(request, TimeValue.timeValueSeconds(1L));
         }
         // check that the stack trace was sent from the data node to the coordinating node
-        assertTrue(transportMessageHasStackTrace.get());
+        assertTrue(transportMessageHasStackTrace.getAsBoolean());
     }
 
     public void testAsyncSearchFailingQueryErrorTraceFalse() throws IOException, InterruptedException {
-        transportMessageHasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_async_search");
@@ -151,11 +131,10 @@ public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
             responseEntity = performRequestAndGetResponseEntityAfterDelay(request, TimeValue.timeValueSeconds(1L));
         }
         // check that the stack trace was not sent from the data node to the coordinating node
-        assertFalse(transportMessageHasStackTrace.get());
+        assertFalse(transportMessageHasStackTrace.getAsBoolean());
     }
 
     public void testAsyncSearchFailingQueryErrorTraceFalseOnSubmitAndTrueOnGet() throws IOException, InterruptedException {
-        transportMessageHasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_async_search");
@@ -180,11 +159,10 @@ public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
             responseEntity = performRequestAndGetResponseEntityAfterDelay(request, TimeValue.timeValueSeconds(1L));
         }
         // check that the stack trace was not sent from the data node to the coordinating node
-        assertFalse(transportMessageHasStackTrace.get());
+        assertFalse(transportMessageHasStackTrace.getAsBoolean());
     }
 
     public void testAsyncSearchFailingQueryErrorTraceTrueOnSubmitAndFalseOnGet() throws IOException, InterruptedException {
-        transportMessageHasStackTrace = new AtomicBoolean();
         setupIndexWithDocs();
 
         Request searchRequest = new Request("POST", "/_async_search");
@@ -209,7 +187,7 @@ public class AsyncSearchErrorTraceIT extends ESIntegTestCase {
             responseEntity = performRequestAndGetResponseEntityAfterDelay(request, TimeValue.timeValueSeconds(1L));
         }
         // check that the stack trace was sent from the data node to the coordinating node
-        assertTrue(transportMessageHasStackTrace.get());
+        assertTrue(transportMessageHasStackTrace.getAsBoolean());
     }
 
     private Map<String, Object> performRequestAndGetResponseEntityAfterDelay(Request r, TimeValue sleep) throws IOException,