Improve memory usage by serializing the response object directly

endpoints-framework/src/main/java/com/google/api/server/spi/ServletInitializationParameters.java

@@ -38,6 +38,7 @@ public abstract class ServletInitializationParameters {
   private static final String ILLEGAL_ARGUMENT_BACKEND_ERROR = "illegalArgumentIsBackendError";
   private static final String EXCEPTION_COMPATIBILITY = "enableExceptionCompatibility";
   private static final String PRETTY_PRINT = "prettyPrint";
+  private static final String ADD_CONTENT_LENGTH = "addContentLength";
 
   private static final Splitter CSV_SPLITTER = Splitter.on(',').omitEmptyStrings().trimResults();
   private static final Joiner CSV_JOINER = Joiner.on(',').skipNulls();
@@ -83,13 +84,22 @@ public String apply(Class<?> clazz) {
    */
   public abstract boolean isPrettyPrintEnabled();
 
+  /**
+   * Returns if the Content-Length header should be set on response. Should be disabled when running
+   * on App Engine, as Content-Length header is discarded by front-end servers. If enabled, has a
+   * small negative impact on CPU usage and latency.
+   *
+   */
+  public abstract boolean isAddContentLength();
+
   public static Builder builder() {
     return new AutoValue_ServletInitializationParameters.Builder()
         .setServletRestricted(true)
         .setClientIdWhitelistEnabled(true)
         .setIllegalArgumentBackendError(false)
         .setExceptionCompatibilityEnabled(true)
-        .setPrettyPrintEnabled(true);
+        .setPrettyPrintEnabled(true)
+        .setAddContentLength(false);
   }
 
   /**
@@ -160,6 +170,11 @@ public Builder setRestricted(boolean servletRestricted) {
      */
     public abstract Builder setPrettyPrintEnabled(boolean prettyPrint);
 
+    /**
+     * Sets if the content length header should be set. Defaults to {@code false}.
+     */
+    public abstract Builder setAddContentLength(boolean addContentLength);
+
     abstract ServletInitializationParameters autoBuild();
 
     public ServletInitializationParameters build() {
@@ -203,6 +218,10 @@ public static ServletInitializationParameters fromServletConfig(
       if (prettyPrint != null) {
         builder.setPrettyPrintEnabled(parseBoolean(prettyPrint, PRETTY_PRINT));
       }
+      String addContentLength = config.getInitParameter(ADD_CONTENT_LENGTH);
+      if (addContentLength != null) {
+        builder.setAddContentLength(parseBoolean(addContentLength, ADD_CONTENT_LENGTH));
+      }
     }
     return builder.build();
   }
@@ -238,6 +257,7 @@ public ImmutableMap<String, String> asMap() {
         .put(ILLEGAL_ARGUMENT_BACKEND_ERROR, Boolean.toString(isIllegalArgumentBackendError()))
         .put(EXCEPTION_COMPATIBILITY, Boolean.toString(isExceptionCompatibilityEnabled()))
         .put(PRETTY_PRINT, Boolean.toString(isPrettyPrintEnabled()))
+        .put(ADD_CONTENT_LENGTH, Boolean.toString(isAddContentLength()))
         .build();
   }
 }

endpoints-framework/src/main/java/com/google/api/server/spi/handlers/EndpointsMethodHandler.java

@@ -91,6 +91,7 @@ protected ResultWriter createResultWriter(EndpointsContext context,
       ApiSerializationConfig serializationConfig) {
     return new RestResponseResultWriter(context.getResponse(), serializationConfig,
         StandardParameters.shouldPrettyPrint(context),
+        initParameters.isAddContentLength(),
         initParameters.isExceptionCompatibilityEnabled());
   }
 

endpoints-framework/src/main/java/com/google/api/server/spi/response/RestResponseResultWriter.java

@@ -16,14 +16,11 @@
 package com.google.api.server.spi.response;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.node.ArrayNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.api.server.spi.ObjectMapperUtil;
 import com.google.api.server.spi.ServiceException;
 import com.google.api.server.spi.config.model.ApiSerializationConfig;
 import com.google.common.base.Strings;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 
 import java.io.IOException;
 
@@ -39,8 +36,8 @@ public class RestResponseResultWriter extends ServletResponseResultWriter {
 
   public RestResponseResultWriter(
       HttpServletResponse servletResponse, ApiSerializationConfig serializationConfig,
-      boolean prettyPrint, boolean enableExceptionCompatibility) {
-    super(servletResponse, serializationConfig, prettyPrint);
+      boolean prettyPrint, boolean addContentLength, boolean enableExceptionCompatibility) {
+    super(servletResponse, serializationConfig, prettyPrint, addContentLength);
     this.enableExceptionCompatibility = enableExceptionCompatibility;
     this.objectMapper = ObjectMapperUtil.createStandardObjectMapper(serializationConfig);
   }
@@ -70,8 +67,7 @@ public void writeError(ServiceException e) throws IOException {
         e.getReason() : errorMap.getReason(e.getStatusCode());
     String domain = !Strings.isNullOrEmpty(e.getDomain()) ?
         e.getDomain() : errorMap.getDomain(e.getStatusCode());
-    write(code, e.getHeaders(),
-        writeValueAsString(createError(code, reason, domain, e.getMessage())));
+    write(code, e.getHeaders(), createError(code, reason, domain, e.getMessage()));
   }
 
   private Object createError(int code, String reason, String domain, String message) {

endpoints-framework/src/main/java/com/google/api/server/spi/response/ServletResponseResultWriter.java

@@ -23,23 +23,28 @@
 import com.google.api.server.spi.types.DateAndTime;
 import com.google.api.server.spi.types.SimpleDate;
 import com.google.appengine.api.datastore.Blob;
+import com.google.common.io.ByteStreams;
+import com.google.common.io.CountingOutputStream;
 
 import com.fasterxml.jackson.core.JsonGenerator;
-import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.Version;
 import com.fasterxml.jackson.databind.JsonSerializer;
 import com.fasterxml.jackson.databind.ObjectWriter;
 import com.fasterxml.jackson.databind.SerializerProvider;
 import com.fasterxml.jackson.databind.module.SimpleModule;
 
+import java.io.BufferedOutputStream;
+import java.io.BufferedWriter;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.Map;
 import java.util.Set;
 
+import javax.servlet.ServletOutputStream;
 import javax.servlet.http.HttpServletResponse;
 
 /**
@@ -68,15 +73,16 @@ public class ServletResponseResultWriter implements ResultWriter {
 
   private final HttpServletResponse servletResponse;
   private final ObjectWriter objectWriter;
+  private final boolean addContentLength;
 
   public ServletResponseResultWriter(
       HttpServletResponse servletResponse, ApiSerializationConfig serializationConfig) {
-    this(servletResponse, serializationConfig, false /* prettyPrint */);
+    this(servletResponse, serializationConfig, false /* prettyPrint */, false /* addContentLength */);
   }
 
   public ServletResponseResultWriter(
       HttpServletResponse servletResponse, ApiSerializationConfig serializationConfig,
-      boolean prettyPrint) {
+      boolean prettyPrint, boolean addContentLength) {
     this.servletResponse = servletResponse;
     Set<SimpleModule> modules = new LinkedHashSet<>();
     modules.addAll(WRITER_MODULES);
@@ -90,27 +96,26 @@ public ServletResponseResultWriter(
       objectWriter = objectWriter.with(new EndpointsPrettyPrinter());
     }
     this.objectWriter = objectWriter;
+    this.addContentLength = addContentLength;
   }
 
   @Override
   public void write(Object response) throws IOException {
     if (response == null) {
       write(HttpServletResponse.SC_NO_CONTENT, null, null);
     } else {
-      write(HttpServletResponse.SC_OK, null,
-          writeValueAsString(ResponseUtil.wrapCollection(response)));
+      write(HttpServletResponse.SC_OK, null, ResponseUtil.wrapCollection(response));
     }
   }
 
   @Override
   public void writeError(ServiceException e) throws IOException {
     Map<String, String> errors = new HashMap<>();
     errors.put(Constant.ERROR_MESSAGE, e.getMessage());
-    write(e.getStatusCode(), e.getHeaders(),
-        writeValueAsString(errors));
+    write(e.getStatusCode(), e.getHeaders(), errors);
   }
 
-  protected void write(int status, Map<String, String> headers, String content) throws IOException {
+  protected void write(int status, Map<String, String> headers, Object content) throws IOException {
     // write response status code
     servletResponse.setStatus(status);
 
@@ -124,8 +129,12 @@ protected void write(int status, Map<String, String> headers, String content) th
     // write response body
     if (content != null) {
       servletResponse.setContentType(SystemService.MIME_JSON);
-      servletResponse.setContentLength(content.getBytes("UTF-8").length);
-      servletResponse.getWriter().write(content);
+      if (addContentLength) {
+        CountingOutputStream counter = new CountingOutputStream(ByteStreams.nullOutputStream());
+        objectWriter.writeValue(counter, content);
+        servletResponse.setContentLength((int) counter.getCount());
+      }
+      objectWriter.writeValue(servletResponse.getOutputStream(), content);
     }
   }
 
@@ -203,13 +212,4 @@ public void serialize(Blob value, JsonGenerator jgen, SerializerProvider provide
     return writeBlobAsBase64Module;
   }
 
-  // Writes a value as a JSON string and translates Jackson exceptions into IOException.
-  protected String writeValueAsString(Object value)
-      throws IOException {
-    try {
-      return objectWriter.writeValueAsString(value);
-    } catch (JsonProcessingException e) {
-      throw new IOException(e);
-    }
-  }
 }

endpoints-framework/src/test/java/com/google/api/server/spi/EndpointsServletTest.java

@@ -113,6 +113,33 @@ public void echo() throws IOException {
     assertThat(actual.get("x").asInt()).isEqualTo(1);
   }
 
+  @Test
+  public void contentLengthHeaderNull() throws IOException {
+    req.setRequestURI("/_ah/api/test/v2/echo");
+    req.setMethod("POST");
+    req.setParameter("x", "1");
+
+    servlet.service(req, resp);
+
+    assertThat(resp.getHeader("Content-Length")).isNull();
+  }
+
+  @Test
+  public void contentLengthHeaderPresent() throws IOException, ServletException {
+    MockServletConfig config = new MockServletConfig();
+    config.addInitParameter("services", TestApi.class.getName());
+    config.addInitParameter("addContentLength", "true");
+    servlet.init(config);
+
+    req.setRequestURI("/_ah/api/test/v2/echo");
+    req.setMethod("POST");
+    req.setParameter("x", "1");
+
+    servlet.service(req, resp);
+
+    assertThat(resp.getHeader("Content-Length")).isNotNull();
+  }
+
   @Test
   public void methodOverride() throws IOException {
     req.setRequestURI("/_ah/api/test/v2/increment");