Browse Source

Avoid convert to string when parse resp in heap attack (#105109)

We've seen cases of OOM errors in the test runner process, which occur 
when we convert a response to a JSON string and then parse it. We can
directly parse from its input stream to avoid these OOM errors.
Nhat Nguyen 1 year ago
parent
commit
86c1fa2a6c

+ 18 - 24
test/external-modules/esql-heap-attack/src/javaRestTest/java/org/elasticsearch/xpack/esql/heap_attack/HeapAttackIT.java

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.heap_attack;
 
-import org.apache.http.HttpEntity;
 import org.apache.http.HttpHost;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.util.EntityUtils;
@@ -21,7 +20,6 @@ import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
-import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.test.ListMatcher;
 import org.elasticsearch.test.cluster.ElasticsearchCluster;
@@ -31,7 +29,6 @@ import org.elasticsearch.threadpool.Scheduler;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xcontent.XContentBuilder;
-import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xcontent.json.JsonXContent;
 import org.junit.After;
 import org.junit.Before;
@@ -87,7 +84,7 @@ public class HeapAttackIT extends ESRestTestCase {
     public void testSortByManyLongsSuccess() throws IOException {
         initManyLongs();
         Response response = sortByManyLongs(2000);
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(response.getEntity()), false);
+        Map<?, ?> map = responseAsMap(response);
         ListMatcher columns = matchesList().item(matchesMap().entry("name", "a").entry("type", "long"))
             .item(matchesMap().entry("name", "b").entry("type", "long"));
         ListMatcher values = matchesList();
@@ -109,7 +106,7 @@ public class HeapAttackIT extends ESRestTestCase {
 
     private void assertCircuitBreaks(ThrowingRunnable r) throws IOException {
         ResponseException e = expectThrows(ResponseException.class, r);
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(e.getResponse().getEntity()), false);
+        Map<?, ?> map = responseAsMap(e.getResponse());
         logger.info("expected circuit breaker {}", map);
         assertMap(
             map,
@@ -133,11 +130,8 @@ public class HeapAttackIT extends ESRestTestCase {
      */
     public void testGroupOnSomeLongs() throws IOException {
         initManyLongs();
-        Map<?, ?> map = XContentHelper.convertToMap(
-            JsonXContent.jsonXContent,
-            EntityUtils.toString(groupOnManyLongs(200).getEntity()),
-            false
-        );
+        Response resp = groupOnManyLongs(200);
+        Map<?, ?> map = responseAsMap(resp);
         ListMatcher columns = matchesList().item(matchesMap().entry("name", "MAX(a)").entry("type", "long"));
         ListMatcher values = matchesList().item(List.of(9));
         assertMap(map, matchesMap().entry("columns", columns).entry("values", values));
@@ -148,11 +142,8 @@ public class HeapAttackIT extends ESRestTestCase {
      */
     public void testGroupOnManyLongs() throws IOException {
         initManyLongs();
-        Map<?, ?> map = XContentHelper.convertToMap(
-            JsonXContent.jsonXContent,
-            EntityUtils.toString(groupOnManyLongs(5000).getEntity()),
-            false
-        );
+        Response resp = groupOnManyLongs(5000);
+        Map<?, ?> map = responseAsMap(resp);
         ListMatcher columns = matchesList().item(matchesMap().entry("name", "MAX(a)").entry("type", "long"));
         ListMatcher values = matchesList().item(List.of(9));
         assertMap(map, matchesMap().entry("columns", columns).entry("values", values));
@@ -180,7 +171,8 @@ public class HeapAttackIT extends ESRestTestCase {
 
     public void testSmallConcat() throws IOException {
         initSingleDocIndex();
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(concat(2).getEntity()), false);
+        Response resp = concat(2);
+        Map<?, ?> map = responseAsMap(resp);
         ListMatcher columns = matchesList().item(matchesMap().entry("name", "a").entry("type", "long"))
             .item(matchesMap().entry("name", "str").entry("type", "keyword"));
         ListMatcher values = matchesList().item(List.of(1, "1".repeat(100)));
@@ -190,7 +182,7 @@ public class HeapAttackIT extends ESRestTestCase {
     public void testHugeConcat() throws IOException {
         initSingleDocIndex();
         ResponseException e = expectThrows(ResponseException.class, () -> concat(10));
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(e.getResponse().getEntity()), false);
+        Map<?, ?> map = responseAsMap(e.getResponse());
         logger.info("expected request rejected {}", map);
         assertMap(
             map,
@@ -216,7 +208,8 @@ public class HeapAttackIT extends ESRestTestCase {
      */
     public void testManyConcat() throws IOException {
         initManyLongs();
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(manyConcat(300).getEntity()), false);
+        Response resp = manyConcat(300);
+        Map<?, ?> map = responseAsMap(resp);
         ListMatcher columns = matchesList();
         for (int s = 0; s < 300; s++) {
             columns = columns.item(matchesMap().entry("name", "str" + s).entry("type", "keyword"));
@@ -267,7 +260,8 @@ public class HeapAttackIT extends ESRestTestCase {
 
     public void testManyEval() throws IOException {
         initManyLongs();
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(manyEval(1).getEntity()), false);
+        Response resp = manyEval(1);
+        Map<?, ?> map = responseAsMap(resp);
         ListMatcher columns = matchesList();
         columns = columns.item(matchesMap().entry("name", "a").entry("type", "long"));
         columns = columns.item(matchesMap().entry("name", "b").entry("type", "long"));
@@ -369,7 +363,7 @@ public class HeapAttackIT extends ESRestTestCase {
      */
     private void fetchManyBigFields(int docs) throws IOException {
         Response response = query("{\"query\": \"FROM manybigfields | SORT f000 | LIMIT " + docs + "\"}", "columns");
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(response.getEntity()), false);
+        Map<?, ?> map = responseAsMap(response);
         ListMatcher columns = matchesList();
         for (int f = 0; f < 1000; f++) {
             columns = columns.item(matchesMap().entry("name", "f" + String.format(Locale.ROOT, "%03d", f)).entry("type", "keyword"));
@@ -381,7 +375,7 @@ public class HeapAttackIT extends ESRestTestCase {
         int fieldValues = 100;
         initMvLongsIndex(1, 3, fieldValues);
         Response response = aggMvLongs(3);
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(response.getEntity()), false);
+        Map<?, ?> map = responseAsMap(response);
         ListMatcher columns = matchesList().item(matchesMap().entry("name", "MAX(f00)").entry("type", "long"))
             .item(matchesMap().entry("name", "f00").entry("type", "long"))
             .item(matchesMap().entry("name", "f01").entry("type", "long"))
@@ -406,7 +400,7 @@ public class HeapAttackIT extends ESRestTestCase {
         int fields = 100;
         initMvLongsIndex(100, fields, 1000);
         Response response = fetchMvLongs();
-        Map<?, ?> map = XContentHelper.convertToMap(JsonXContent.jsonXContent, EntityUtils.toString(response.getEntity()), false);
+        Map<?, ?> map = responseAsMap(response);
         ListMatcher columns = matchesList();
         for (int f = 0; f < fields; f++) {
             columns = columns.item(matchesMap().entry("name", String.format(Locale.ROOT, "f%02d", f)).entry("type", "long"));
@@ -570,8 +564,8 @@ public class HeapAttackIT extends ESRestTestCase {
             return;
         }
         assertBusy(() -> {
-            HttpEntity entity = adminClient().performRequest(new Request("GET", "/_nodes/stats")).getEntity();
-            Map<?, ?> stats = XContentHelper.convertToMap(XContentType.JSON.xContent(), entity.getContent(), false);
+            Response response = adminClient().performRequest(new Request("GET", "/_nodes/stats"));
+            Map<?, ?> stats = responseAsMap(response);
             Map<?, ?> nodes = (Map<?, ?>) stats.get("nodes");
             for (Object n : nodes.values()) {
                 Map<?, ?> node = (Map<?, ?>) n;