Skip to content

Commit 79ebb5e

Browse files
nik9000nik9000
committed
ESQL: Check speed up spec tests
Speeds up the spec tests by: * Randomize `async`/`sync` instead of both * Check the took time rarely * Don't check if index exists every time * Reduce transfer when checking breaker * Cache capabilities checks The first one is most of the speed up but the rest save another couple of minutes. Time drops from 13m45s to 6m or so on my laptop.
1 parent d2e75cc commit 79ebb5e

File tree

12 files changed

+70
-133
lines changed

12 files changed

+70
-133
lines changed

‎x-pack/plugin/esql/qa/server/mixed-cluster/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/mixed/MixedClusterEsqlSpecIT.java‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import org.elasticsearch.test.rest.TestFeatureService;
1313
import org.elasticsearch.xpack.esql.CsvSpecReader.CsvTestCase;
1414
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
15-
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
1615
import org.junit.AfterClass;
1716
import org.junit.Before;
1817
import org.junit.ClassRule;
@@ -43,11 +42,6 @@ public void extractOldClusterFeatures() {
4342
}
4443
}
4544

46-
protected static boolean oldClusterHasFeature(String featureId) {
47-
assert oldClusterTestFeatureService != null;
48-
return oldClusterTestFeatureService.clusterHasFeature(featureId);
49-
}
50-
5145
@AfterClass
5246
public static void cleanUp() {
5347
oldClusterTestFeatureService = null;
@@ -59,10 +53,9 @@ public MixedClusterEsqlSpecIT(
5953
String testName,
6054
Integer lineNumber,
6155
CsvTestCase testCase,
62-
String instructions,
63-
Mode mode
56+
String instructions
6457
) {
65-
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
58+
super(fileName, groupName, testName, lineNumber, testCase, instructions);
6659
}
6760

6861
@Override

‎x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java‎

Lines changed: 3 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import org.elasticsearch.xpack.esql.CsvTestsDataLoader;
2828
import org.elasticsearch.xpack.esql.SpecReader;
2929
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
30-
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
3130
import org.junit.AfterClass;
3231
import org.junit.ClassRule;
3332
import org.junit.rules.RuleChain;
@@ -36,7 +35,6 @@
3635
import java.io.ByteArrayInputStream;
3736
import java.io.IOException;
3837
import java.net.URL;
39-
import java.util.ArrayList;
4038
import java.util.Arrays;
4139
import java.util.List;
4240
import java.util.Locale;
@@ -59,7 +57,6 @@
5957
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.JOIN_PLANNING_V1;
6058
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.METADATA_FIELDS_REMOTE_TEST;
6159
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.UNMAPPED_FIELDS;
62-
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode.SYNC;
6360
import static org.mockito.ArgumentMatchers.any;
6461
import static org.mockito.Mockito.doAnswer;
6562
import static org.mockito.Mockito.mock;
@@ -87,19 +84,7 @@ public class MultiClusterSpecIT extends EsqlSpecTestCase {
8784
public static List<Object[]> readScriptSpec() throws Exception {
8885
List<URL> urls = classpathResources("/*.csv-spec");
8986
assertTrue("Not enough specs found " + urls, urls.size() > 0);
90-
List<Object[]> specs = SpecReader.readScriptSpec(urls, specParser());
91-
92-
int len = specs.get(0).length;
93-
List<Object[]> testcases = new ArrayList<>();
94-
for (var spec : specs) {
95-
for (Mode mode : List.of(SYNC)) { // No async, for now
96-
Object[] obj = new Object[len + 1];
97-
System.arraycopy(spec, 0, obj, 0, len);
98-
obj[len] = mode;
99-
testcases.add(obj);
100-
}
101-
}
102-
return testcases;
87+
return SpecReader.readScriptSpec(urls, specParser());
10388
}
10489

10590
public MultiClusterSpecIT(
@@ -108,10 +93,9 @@ public MultiClusterSpecIT(
10893
String testName,
10994
Integer lineNumber,
11095
CsvTestCase testCase,
111-
String instructions,
112-
Mode mode
96+
String instructions
11397
) {
114-
super(fileName, groupName, testName, lineNumber, convertToRemoteIndices(testCase), instructions, mode);
98+
super(fileName, groupName, testName, lineNumber, convertToRemoteIndices(testCase), instructions);
11599
}
116100

117101
// TODO: think how to handle this better

‎x-pack/plugin/esql/qa/server/multi-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/multi_node/EsqlSpecIT.java‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@
88
package org.elasticsearch.xpack.esql.qa.multi_node;
99

1010
import org.elasticsearch.test.cluster.ElasticsearchCluster;
11+
import org.elasticsearch.test.junit.annotations.TestLogging;
1112
import org.elasticsearch.xpack.esql.CsvSpecReader.CsvTestCase;
1213
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
13-
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
1414
import org.junit.ClassRule;
1515

1616
import java.io.IOException;
@@ -24,16 +24,8 @@ protected String getTestRestCluster() {
2424
return cluster.getHttpAddresses();
2525
}
2626

27-
public EsqlSpecIT(
28-
String fileName,
29-
String groupName,
30-
String testName,
31-
Integer lineNumber,
32-
CsvTestCase testCase,
33-
String instructions,
34-
Mode mode
35-
) {
36-
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
27+
public EsqlSpecIT(String fileName, String groupName, String testName, Integer lineNumber, CsvTestCase testCase, String instructions) {
28+
super(fileName, groupName, testName, lineNumber, testCase, instructions);
3729
}
3830

3931
@Override

‎x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/EsqlSpecIT.java‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import org.elasticsearch.xpack.esql.planner.PhysicalSettings;
2020
import org.elasticsearch.xpack.esql.plugin.ComputeService;
2121
import org.elasticsearch.xpack.esql.qa.rest.EsqlSpecTestCase;
22-
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
2322
import org.junit.Before;
2423
import org.junit.ClassRule;
2524

@@ -37,16 +36,8 @@ protected String getTestRestCluster() {
3736
return cluster.getHttpAddresses();
3837
}
3938

40-
public EsqlSpecIT(
41-
String fileName,
42-
String groupName,
43-
String testName,
44-
Integer lineNumber,
45-
CsvTestCase testCase,
46-
String instructions,
47-
Mode mode
48-
) {
49-
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
39+
public EsqlSpecIT(String fileName, String groupName, String testName, Integer lineNumber, CsvTestCase testCase, String instructions) {
40+
super(fileName, groupName, testName, lineNumber, testCase, instructions);
5041
}
5142

5243
@Override

‎x-pack/plugin/esql/qa/server/single-node/src/javaRestTest/java/org/elasticsearch/xpack/esql/qa/single_node/GenerativeForkIT.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import org.elasticsearch.test.TestClustersThreadFilter;
1313
import org.elasticsearch.test.cluster.ElasticsearchCluster;
1414
import org.elasticsearch.xpack.esql.CsvSpecReader;
15-
import org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.Mode;
1615
import org.elasticsearch.xpack.esql.qa.rest.generative.GenerativeForkRestTest;
1716
import org.junit.ClassRule;
1817

@@ -32,10 +31,9 @@ public GenerativeForkIT(
3231
String testName,
3332
Integer lineNumber,
3433
CsvSpecReader.CsvTestCase testCase,
35-
String instructions,
36-
Mode mode
34+
String instructions
3735
) {
38-
super(fileName, groupName, testName, lineNumber, testCase, instructions, mode);
36+
super(fileName, groupName, testName, lineNumber, testCase, instructions);
3937
}
4038

4139
@Override

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/EsqlSpecTestCase.java‎

Lines changed: 20 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,6 @@
6868
import static org.elasticsearch.xpack.esql.CsvTestUtils.ExpectedResults;
6969
import static org.elasticsearch.xpack.esql.CsvTestUtils.isEnabled;
7070
import static org.elasticsearch.xpack.esql.CsvTestUtils.loadCsvSpecValues;
71-
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.availableDatasetsForEs;
7271
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.createInferenceEndpoints;
7372
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.deleteInferenceEndpoints;
7473
import static org.elasticsearch.xpack.esql.CsvTestsDataLoader.loadDataSetIntoEs;
@@ -79,6 +78,7 @@
7978
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SEMANTIC_TEXT_FIELD_CAPS;
8079
import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.SOURCE_FIELD_MAPPING;
8180
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.assertNotPartial;
81+
import static org.elasticsearch.xpack.esql.qa.rest.RestEsqlTestCase.hasCapabilities;
8282

8383
// This test can run very long in serverless configurations
8484
@TimeoutSuite(millis = 30 * TimeUnits.MINUTE)
@@ -101,19 +101,7 @@ public abstract class EsqlSpecTestCase extends ESRestTestCase {
101101
public static List<Object[]> readScriptSpec() throws Exception {
102102
List<URL> urls = classpathResources("/*.csv-spec");
103103
assertTrue("Not enough specs found " + urls, urls.size() > 0);
104-
List<Object[]> specs = SpecReader.readScriptSpec(urls, specParser());
105-
106-
int len = specs.get(0).length;
107-
List<Object[]> testcases = new ArrayList<>();
108-
for (var spec : specs) {
109-
for (Mode mode : Mode.values()) {
110-
Object[] obj = new Object[len + 1];
111-
System.arraycopy(spec, 0, obj, 0, len);
112-
obj[len] = mode;
113-
testcases.add(obj);
114-
}
115-
}
116-
return testcases;
104+
return SpecReader.readScriptSpec(urls, specParser());
117105
}
118106

119107
protected EsqlSpecTestCase(
@@ -122,34 +110,38 @@ protected EsqlSpecTestCase(
122110
String testName,
123111
Integer lineNumber,
124112
CsvTestCase testCase,
125-
String instructions,
126-
Mode mode
113+
String instructions
127114
) {
128115
this.fileName = fileName;
129116
this.groupName = groupName;
130117
this.testName = testName;
131118
this.lineNumber = lineNumber;
132119
this.testCase = testCase;
133120
this.instructions = instructions;
134-
this.mode = mode;
121+
this.mode = randomFrom(Mode.values());
135122
}
136123

124+
private static boolean dataLoaded = false;
125+
137126
@Before
138127
public void setup() throws IOException {
139-
if (supportsInferenceTestService()) {
140-
createInferenceEndpoints(adminClient());
141-
}
142-
143128
boolean supportsLookup = supportsIndexModeLookup();
144129
boolean supportsSourceMapping = supportsSourceFieldMapping();
145-
if (indexExists(availableDatasetsForEs(client(), supportsLookup, supportsSourceMapping).iterator().next().indexName()) == false) {
146-
loadDataSetIntoEs(client(), supportsLookup, supportsSourceMapping);
130+
boolean supportsInferenceTestService = supportsInferenceTestService();
131+
if (dataLoaded == false) {
132+
if (supportsInferenceTestService) {
133+
createInferenceEndpoints(adminClient());
134+
}
135+
136+
loadDataSetIntoEs(client(), supportsLookup, supportsSourceMapping, supportsInferenceTestService);
137+
dataLoaded = true;
147138
}
148139
}
149140

150141
@AfterClass
151142
public static void wipeTestData() throws IOException {
152143
try {
144+
dataLoaded = false;
153145
adminClient().performRequest(new Request("DELETE", "/*"));
154146
} catch (ResponseException e) {
155147
// 404 here just means we had no indexes
@@ -159,7 +151,6 @@ public static void wipeTestData() throws IOException {
159151
}
160152

161153
deleteInferenceEndpoints(adminClient());
162-
163154
}
164155

165156
public boolean logResults() {
@@ -211,38 +202,6 @@ protected static void checkCapabilities(RestClient client, TestFeatureService te
211202
}
212203
}
213204

214-
protected static boolean hasCapabilities(List<String> requiredCapabilities) throws IOException {
215-
return hasCapabilities(adminClient(), requiredCapabilities);
216-
}
217-
218-
public static boolean hasCapabilities(RestClient client, List<String> requiredCapabilities) throws IOException {
219-
if (requiredCapabilities.isEmpty()) {
220-
return true;
221-
}
222-
try {
223-
if (clusterHasCapability(client, "POST", "/_query", List.of(), requiredCapabilities).orElse(false)) {
224-
return true;
225-
}
226-
LOGGER.info("capabilities API returned false, we might be in a mixed version cluster so falling back to cluster features");
227-
} catch (ResponseException e) {
228-
if (e.getResponse().getStatusLine().getStatusCode() / 100 == 4) {
229-
/*
230-
* The node we're testing against is too old for the capabilities
231-
* API which means it has to be pretty old. Very old capabilities
232-
* are ALSO present in the features API, so we can check them instead.
233-
*
234-
* It's kind of weird that we check for *any* 400, but that's required
235-
* because old versions of Elasticsearch return 400, not the expected
236-
* 404.
237-
*/
238-
LOGGER.info("capabilities API failed, falling back to cluster features");
239-
} else {
240-
throw e;
241-
}
242-
}
243-
return false;
244-
}
245-
246205
protected boolean supportsInferenceTestService() {
247206
return true;
248207
}
@@ -271,7 +230,9 @@ protected final void doTest(String query) throws Throwable {
271230
builder.tables(tables());
272231
}
273232

274-
Map<?, ?> prevTooks = supportsTook() ? tooks() : null;
233+
boolean checkTook = supportsTook() && rarely();
234+
235+
Map<?, ?> prevTooks = checkTook ? tooks() : null;
275236
Map<String, Object> answer = RestEsqlTestCase.runEsql(
276237
builder.query(query),
277238
testCase.assertWarnings(deduplicateExactWarnings()),
@@ -296,7 +257,7 @@ protected final void doTest(String query) throws Throwable {
296257

297258
assertResults(expectedColumnsWithValues, actualColumns, actualValues, testCase.ignoreOrder, logger);
298259

299-
if (supportsTook()) {
260+
if (checkTook) {
300261
LOGGER.info("checking took incremented from {}", prevTooks);
301262
long took = ((Number) answer.get("took")).longValue();
302263
int prevTookHisto = ((Number) prevTooks.remove(tookKey(took))).intValue();
@@ -413,15 +374,14 @@ protected boolean preserveClusterUponCompletion() {
413374
return true;
414375
}
415376

416-
@Before
417377
@After
418378
public void assertRequestBreakerEmptyAfterTests() throws Exception {
419379
assertRequestBreakerEmpty();
420380
}
421381

422382
public static void assertRequestBreakerEmpty() throws Exception {
423383
assertBusy(() -> {
424-
HttpEntity entity = adminClient().performRequest(new Request("GET", "/_nodes/stats")).getEntity();
384+
HttpEntity entity = adminClient().performRequest(new Request("GET", "/_nodes/stats?metric=breaker")).getEntity();
425385
Map<?, ?> stats = XContentHelper.convertToMap(XContentType.JSON.xContent(), entity.getContent(), false);
426386
Map<?, ?> nodes = (Map<?, ?>) stats.get("nodes");
427387

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.client.RequestOptions;
1717
import org.elasticsearch.client.Response;
1818
import org.elasticsearch.client.ResponseException;
19+
import org.elasticsearch.client.RestClient;
1920
import org.elasticsearch.client.WarningsHandler;
2021
import org.elasticsearch.common.bytes.BytesArray;
2122
import org.elasticsearch.common.io.Streams;
@@ -52,6 +53,7 @@
5253
import java.util.Locale;
5354
import java.util.Map;
5455
import java.util.Set;
56+
import java.util.concurrent.ConcurrentHashMap;
5557
import java.util.function.IntFunction;
5658

5759
import static java.util.Collections.emptySet;
@@ -1334,8 +1336,8 @@ public static Map<String, Object> runEsqlAsync(
13341336
checkKeepOnCompletion(requestObject, json, keepOnCompletion);
13351337
String id = (String) json.get("id");
13361338

1337-
var supportsAsyncHeaders = clusterHasCapability("POST", "/_query", List.of(), List.of("async_query_status_headers")).orElse(false);
1338-
var supportsSuggestedCast = clusterHasCapability("POST", "/_query", List.of(), List.of("suggested_cast")).orElse(false);
1339+
var supportsAsyncHeaders = hasCapabilities(client(), List.of("async_query_status_headers"));
1340+
var supportsSuggestedCast = hasCapabilities(client(), List.of("suggested_cast"));
13391341

13401342
if (id == null) {
13411343
// no id returned from an async call, must have completed immediately and without keep_on_completion
@@ -1409,13 +1411,27 @@ public static Map<String, Object> runEsqlAsync(
14091411
private static void prepareProfileLogger(RequestObjectBuilder requestObject, @Nullable ProfileLogger profileLogger) throws IOException {
14101412
if (profileLogger != null) {
14111413
profileLogger.clearProfile();
1412-
var isProfileSafe = clusterHasCapability("POST", "/_query", List.of(), List.of("fixed_profile_serialization")).orElse(false);
1414+
var isProfileSafe = hasCapabilities(client(), List.of("fixed_profile_serialization"));
14131415
if (isProfileSafe) {
14141416
requestObject.profile(true);
14151417
}
14161418
}
14171419
}
14181420

1421+
private static final Map<List<String>, Boolean> capabilities = new ConcurrentHashMap<>();
1422+
1423+
public static boolean hasCapabilities(RestClient client, List<String> requiredCapabilities) throws IOException {
1424+
if (requiredCapabilities.isEmpty()) {
1425+
return true;
1426+
}
1427+
Boolean cap = capabilities.get(requiredCapabilities);
1428+
if (cap == null) {
1429+
cap = clusterHasCapability(client, "POST", "/_query", List.of(), requiredCapabilities).orElse(false);
1430+
capabilities.put(requiredCapabilities, cap);
1431+
}
1432+
return cap;
1433+
}
1434+
14191435
private static Object removeOriginalTypesAndSuggestedCast(Object response) {
14201436
if (response instanceof ArrayList<?> columns) {
14211437
var newColumns = new ArrayList<>();

0 commit comments

Comments
 (0)