Merge pull request #292 from artem-ag/optimize-response-writers

Optimize single & batch query response writers to reduce memory allocation overhead

Author: oliemansm

graphql-java-kickstart/src/main/java/graphql/kickstart/execution/GraphQLObjectMapper.java

@@ -111,6 +111,19 @@ public void serializeResultAsJson(Writer writer, ExecutionResult executionResult
     getJacksonMapper().writeValue(writer, createResultFromExecutionResult(executionResult));
   }
 
+  /**
+   * Serializes result as bytes in UTF-8 encoding instead of string.
+   *
+   * @param executionResult query execution result to serialize.
+   * @return result serialized into Json representation in UTF-8 encoding, converted into {@code
+   *     byte[]}.
+   */
+  @SneakyThrows
+  public byte[] serializeResultAsBytes(ExecutionResult executionResult) {
+    return getJacksonMapper()
+        .writeValueAsBytes(createResultFromExecutionResult(executionResult));
+  }
+
   public boolean areErrorsPresent(ExecutionResult executionResult) {
     return graphQLErrorHandlerSupplier.get().errorsPresent(executionResult.getErrors());
   }

graphql-java-servlet/src/main/java/graphql/kickstart/servlet/BatchedQueryResponseWriter.java

@@ -4,7 +4,7 @@
 import graphql.kickstart.execution.GraphQLObjectMapper;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
-import java.util.Iterator;
+import java.util.ArrayList;
 import java.util.List;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
@@ -14,32 +14,48 @@
 @Slf4j
 @RequiredArgsConstructor
 class BatchedQueryResponseWriter implements QueryResponseWriter {
-
   private final List<ExecutionResult> results;
   private final GraphQLObjectMapper graphQLObjectMapper;
 
   @Override
   public void write(HttpServletRequest request, HttpServletResponse response) throws IOException {
+    response.setCharacterEncoding(StandardCharsets.UTF_8.name());
     response.setContentType(HttpRequestHandler.APPLICATION_JSON_UTF8);
     response.setStatus(HttpRequestHandler.STATUS_OK);
 
-    Iterator<ExecutionResult> executionInputIterator = results.iterator();
-    StringBuilder responseBuilder = new StringBuilder();
-    responseBuilder.append('[');
-    while (executionInputIterator.hasNext()) {
-      responseBuilder
-          .append(graphQLObjectMapper.serializeResultAsJson(executionInputIterator.next()));
-      if (executionInputIterator.hasNext()) {
-        responseBuilder.append(',');
+    // Use direct serialization to byte arrays and avoid any string concatenation to save multiple
+    // GiB of memory allocation during large response processing.
+    List<byte[]> serializedResults = new ArrayList<>(2 * results.size() + 1);
+
+    if (results.size() > 0) {
+      serializedResults.add("[".getBytes(StandardCharsets.UTF_8));
+    } else {
+      serializedResults.add("[]".getBytes(StandardCharsets.UTF_8));
+    }
+    long totalLength = serializedResults.get(0).length;
+
+    // '[', ',' and ']' are all 1 byte in UTF-8.
+    for (int i = 0; i < results.size(); i++) {
+      byte[] currentResult = graphQLObjectMapper.serializeResultAsBytes(results.get(i));
+      serializedResults.add(currentResult);
+
+      if (i != results.size() - 1) {
+        serializedResults.add(",".getBytes(StandardCharsets.UTF_8));
+      } else {
+        serializedResults.add("]".getBytes(StandardCharsets.UTF_8));
       }
+      totalLength += currentResult.length + 1; // result.length + ',' or ']'
     }
-    responseBuilder.append(']');
 
-    String responseContent = responseBuilder.toString();
-    byte[] contentBytes = responseContent.getBytes(StandardCharsets.UTF_8);
+    if (totalLength > Integer.MAX_VALUE) {
+      throw new IllegalStateException(
+              "Response size exceed 2GiB. Query will fail. Seen size: " + totalLength);
+    }
+    response.setContentLength((int) totalLength);
 
-    response.setContentLength(contentBytes.length);
-    response.getOutputStream().write(contentBytes);
+    for (byte[] result : serializedResults) {
+      response.getOutputStream().write(result);
+    }
   }
 
 }

graphql-java-servlet/src/main/java/graphql/kickstart/servlet/SingleQueryResponseWriter.java

@@ -19,10 +19,9 @@ public void write(HttpServletRequest request, HttpServletResponse response) thro
     response.setContentType(HttpRequestHandler.APPLICATION_JSON_UTF8);
     response.setStatus(HttpRequestHandler.STATUS_OK);
     response.setCharacterEncoding(StandardCharsets.UTF_8.name());
-    String responseContent = graphQLObjectMapper.serializeResultAsJson(result);
-    byte[] contentBytes = responseContent.getBytes(StandardCharsets.UTF_8);
+
+    byte[] contentBytes = graphQLObjectMapper.serializeResultAsBytes(result);
     response.setContentLength(contentBytes.length);
     response.getOutputStream().write(contentBytes);
   }
-
 }

graphql-java-servlet/src/test/groovy/graphql/kickstart/servlet/BatchedQueryResponseWriterTest.groovy

@@ -0,0 +1,55 @@
+package graphql.kickstart.servlet
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import graphql.ExecutionResultImpl
+import graphql.kickstart.execution.GraphQLObjectMapper
+import spock.lang.Specification
+import spock.lang.Unroll
+
+import javax.servlet.ServletOutputStream
+import javax.servlet.http.HttpServletRequest
+import javax.servlet.http.HttpServletResponse
+import java.nio.charset.StandardCharsets
+
+class BatchedQueryResponseWriterTest extends Specification {
+
+    @Unroll
+    def "should write utf8 results into the response with content #result"() {
+        given:
+        def byteArrayOutputStream = new ByteArrayOutputStream()
+        def graphQLObjectMapperMock = GraphQLObjectMapper.newBuilder().withObjectMapperProvider({ new ObjectMapper() }).build()
+        graphQLObjectMapperMock.getJacksonMapper() >> new ObjectMapper()
+
+        def requestMock = Mock(HttpServletRequest)
+        def responseMock = Mock(HttpServletResponse)
+        def servletOutputStreamMock = Mock(ServletOutputStream)
+
+        responseMock.getOutputStream() >> servletOutputStreamMock
+
+        1 * responseMock.setContentLength(expectedContentLengh)
+        1 * responseMock.setCharacterEncoding(StandardCharsets.UTF_8.name())
+        (1.._) * servletOutputStreamMock.write(_) >> { value ->
+            byteArrayOutputStream.write((byte[])(value[0]))
+        }
+
+        def executionResultList = new ArrayList()
+        for (LinkedHashMap<Object, Object> value : result) {
+            executionResultList.add(new ExecutionResultImpl(value, []))
+        }
+
+        def writer = new BatchedQueryResponseWriter(executionResultList, graphQLObjectMapperMock)
+
+        when:
+        writer.write(requestMock, responseMock)
+
+        then:
+        byteArrayOutputStream.toString(StandardCharsets.UTF_8.name()) == expectedResponseContent
+
+        where:
+        result                      || expectedContentLengh | expectedResponseContent
+        [[testValue: "abcde"]]      || 32                   | """[{"data":{"testValue":"abcde"}}]"""
+        [[testValue: "äöüüöß"]]     || 39                   | """[{"data":{"testValue":"äöüüöß"}}]"""
+        []                          || 2                    | """[]"""
+        [[k1: "äöüüöß"], [k2: "a"]] || 52                   | """[{"data":{"k1":"äöüüöß"}},{"data":{"k2":"a"}}]"""
+    }
+}