Merge pull request #2059 from atsushieno/avoid-buffered-response-stream

Default GitHubConnectorResponse to streamed body instead of in-memory buffer

Author: bitwiseman

pom.xml

@@ -611,7 +611,7 @@
                 </configuration>
               </execution>
               <execution>
-                <id>httpclient-test</id>
+                <id>httpclient-test-tracing</id>
                 <phase>integration-test</phase>
                 <goals>
                   <goal>test</goal>

src/main/java/org/kohsuke/github/GitHubClient.java

@@ -629,6 +629,7 @@ private void logResponseBody(@Nonnull final GitHubConnectorResponse response) {
         LOGGER.log(FINEST, () -> {
             String body;
             try {
+                response.setBodyStreamRereadable();
                 body = GitHubResponse.getBodyAsString(response);
             } catch (Throwable e) {
                 body = "Error reading response body";

src/main/java/org/kohsuke/github/connector/GitHubConnectorResponse.java

@@ -7,15 +7,25 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.io.InputStream;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
 import java.util.zip.GZIPInputStream;
 
 import javax.annotation.CheckForNull;
 import javax.annotation.Nonnull;
 
+import static java.net.HttpURLConnection.HTTP_OK;
+
 /**
  * Response information supplied when a response is received and before the body is processed.
  * <p>
+ * During a request to GitHub, {@link GitHubConnector#send(GitHubConnectorRequest)} returns a
+ * {@link GitHubConnectorResponse}. This is processed to create a GitHubResponse.
+ * <p>
  * Instances of this class are closed once the response is done being processed. This means that {@link #bodyStream()}
  * will not be readable after a call is completed.
  *
@@ -35,6 +45,11 @@ public abstract class GitHubConnectorResponse implements Closeable {
     private final GitHubConnectorRequest request;
     @Nonnull
     private final Map<String, List<String>> headers;
+    private boolean bodyStreamCalled = false;
+    private InputStream bodyStream = null;
+    private byte[] bodyBytes = null;
+    private boolean isClosed = false;
+    private boolean isBodyStreamRereadable;
 
     /**
      * GitHubConnectorResponse constructor
@@ -58,6 +73,7 @@ protected GitHubConnectorResponse(@Nonnull GitHubConnectorRequest request,
             caseInsensitiveMap.put(entry.getKey(), Collections.unmodifiableList(new ArrayList<>(entry.getValue())));
         }
         this.headers = Collections.unmodifiableMap(caseInsensitiveMap);
+        this.isBodyStreamRereadable = false;
     }
 
     /**
@@ -79,17 +95,72 @@ public String header(String name) {
     /**
      * The response body as an {@link InputStream}.
      *
+     * When {@link #isBodyStreamRereadable} is false, {@link #bodyStream()} can only be called once and the returned
+     * stream should be assumed to be read-once and not resetable. This is the default behavior for HTTP_OK responses
+     * and significantly reduces memory usage.
+     *
+     * When {@link #isBodyStreamRereadable} is true, {@link #bodyStream()} can be called be called multiple times. The
+     * full stream data is read into a byte array during the first call. Each call returns a new stream backed by the
+     * same byte array. This uses more memory, but is required to enable rereading the body stream during trace logging,
+     * debugging, and error responses.
+     *
      * @return the response body
      * @throws IOException
      *             if response stream is null or an I/O Exception occurs.
      */
     @Nonnull
-    public abstract InputStream bodyStream() throws IOException;
+    public InputStream bodyStream() throws IOException {
+        synchronized (this) {
+            if (isClosed) {
+                throw new IOException("Response is closed");
+            }
+
+            if (bodyStreamCalled) {
+                if (!isBodyStreamRereadable()) {
+                    throw new IOException("Response body not rereadable");
+                }
+            } else {
+                bodyStream = wrapStream(rawBodyStream());
+                bodyStreamCalled = true;
+            }
+
+            if (bodyStream == null) {
+                throw new IOException("Response body missing, stream null");
+            } else if (!isBodyStreamRereadable()) {
+                return bodyStream;
+            }
+
+            // Load rereadable byte array
+            if (bodyBytes == null) {
+                bodyBytes = IOUtils.toByteArray(bodyStream);
+                // Close the raw body stream after successfully reading
+                IOUtils.closeQuietly(bodyStream);
+            }
+
+            return new ByteArrayInputStream(bodyBytes);
+        }
+    }
 
     /**
-     * Gets the {@link GitHubConnectorRequest} for this response.
+     * Get the raw implementation specific body stream for this response.
+     *
+     * This method will only be called once to completion. If an exception is thrown by this method, it may be called
+     * multiple times.
      *
-     * @return the {@link GitHubConnectorRequest} for this response.
+     * The stream returned from this method will be closed when the response is closed or sooner. Inheriting classes do
+     * not need to close it.
+     *
+     * @return the stream for the raw response
+     * @throws IOException
+     *             if an I/O Exception occurs.
+     */
+    @CheckForNull
+    protected abstract InputStream rawBodyStream() throws IOException;
+
+    /**
+     * Gets the {@link GitHubConnector} for this response.
+     *
+     * @return the {@link GitHubConnector} for this response.
      */
     @Nonnull
     public GitHubConnectorRequest request() {
@@ -116,6 +187,56 @@ public Map<String, List<String>> allHeaders() {
         return headers;
     }
 
+    /**
+     * The body stream rereadable state.
+     *
+     * Body stream defaults to read once for HTTP_OK responses (to reduce memory usage). For non-HTTP_OK responses, body
+     * stream is switched to rereadable (in-memory byte array) for error processing.
+     *
+     * Calling {@link #setBodyStreamRereadable()} will force {@link #isBodyStreamRereadable} to be true for this
+     * response regardless of {@link #statusCode} value.
+     *
+     * @return true when body stream is rereadable.
+     */
+    public boolean isBodyStreamRereadable() {
+        synchronized (this) {
+            return isBodyStreamRereadable || statusCode != HTTP_OK;
+        }
+    }
+
+    /**
+     * Force body stream to rereadable regardless of status code.
+     *
+     * Calling {@link #setBodyStreamRereadable()} will force {@link #isBodyStreamRereadable} to be true for this
+     * response regardless of {@link #statusCode} value.
+     *
+     * This is required to support body value logging during low-level tracing but should be avoided in general since it
+     * consumes significantly more memory.
+     *
+     * Will throw runtime exception if a non-rereadable body stream has already been returned from
+     * {@link #bodyStream()}.
+     */
+    public void setBodyStreamRereadable() {
+        synchronized (this) {
+            if (bodyStreamCalled && !isBodyStreamRereadable()) {
+                throw new RuntimeException("bodyStream() already called in read-once mode");
+            }
+            isBodyStreamRereadable = true;
+        }
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void close() throws IOException {
+        synchronized (this) {
+            IOUtils.closeQuietly(bodyStream);
+            isClosed = true;
+            this.bodyBytes = null;
+        }
+    }
+
     /**
      * Handles wrapping the body stream if indicated by the "Content-Encoding" header.
      *
@@ -155,13 +276,12 @@ public final int parseInt(String name) throws NumberFormatException {
 
     /**
      * A ByteArrayResponse class
+     *
+     * @deprecated Inherit directly from {@link GitHubConnectorResponse}.
      */
+    @Deprecated
     public abstract static class ByteArrayResponse extends GitHubConnectorResponse {
 
-        private boolean inputStreamRead = false;
-        private byte[] inputBytes = null;
-        private boolean isClosed = false;
-
         /**
          * Constructor for ByteArray Response
          *
@@ -177,52 +297,5 @@ protected ByteArrayResponse(@Nonnull GitHubConnectorRequest request,
                 @Nonnull Map<String, List<String>> headers) {
             super(request, statusCode, headers);
         }
-
-        /**
-         * {@inheritDoc}
-         */
-        @Override
-        @Nonnull
-        public InputStream bodyStream() throws IOException {
-            if (isClosed) {
-                throw new IOException("Response is closed");
-            }
-            synchronized (this) {
-                if (!inputStreamRead) {
-                    InputStream rawStream = rawBodyStream();
-                    try (InputStream stream = wrapStream(rawStream)) {
-                        if (stream != null) {
-                            inputBytes = IOUtils.toByteArray(stream);
-                        }
-                    }
-                    inputStreamRead = true;
-                }
-            }
-
-            if (inputBytes == null) {
-                throw new IOException("Response body missing, stream null");
-            }
-
-            return new ByteArrayInputStream(inputBytes);
-        }
-
-        /**
-         * Get the raw implementation specific body stream for this response.
-         *
-         * This method will only be called once to completion. If an exception is thrown, it may be called multiple
-         * times.
-         *
-         * @return the stream for the raw response
-         * @throws IOException
-         *             if an I/O Exception occurs.
-         */
-        @CheckForNull
-        protected abstract InputStream rawBodyStream() throws IOException;
-
-        @Override
-        public void close() throws IOException {
-            isClosed = true;
-            this.inputBytes = null;
-        }
     }
 }

src/main/java/org/kohsuke/github/extras/HttpClientGitHubConnector.java

@@ -93,7 +93,7 @@ public GitHubConnectorResponse send(GitHubConnectorRequest connectorRequest) thr
      *
      * Implementation specific to {@link HttpResponse}.
      */
-    private static class HttpClientGitHubConnectorResponse extends GitHubConnectorResponse.ByteArrayResponse {
+    private static class HttpClientGitHubConnectorResponse extends GitHubConnectorResponse {
 
         @Nonnull
         private final HttpResponse<InputStream> response;
@@ -113,7 +113,6 @@ protected InputStream rawBodyStream() throws IOException {
         @Override
         public void close() throws IOException {
             super.close();
-            IOUtils.closeQuietly(response.body());
         }
     }
 }

src/main/java/org/kohsuke/github/extras/okhttp3/OkHttpGitHubConnector.java

@@ -103,7 +103,7 @@ private List<ConnectionSpec> TlsConnectionSpecs() {
      *
      * Implementation specific to {@link okhttp3.Response}.
      */
-    private static class OkHttpGitHubConnectorResponse extends GitHubConnectorResponse.ByteArrayResponse {
+    private static class OkHttpGitHubConnectorResponse extends GitHubConnectorResponse {
 
         @Nonnull
         private final Response response;

src/test/java/org/kohsuke/github/GHWorkflowRunTest.java

@@ -9,6 +9,7 @@
 import org.kohsuke.github.GHWorkflowRun.Status;
 import org.kohsuke.github.function.InputStreamFunction;
 
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.time.Duration;
 import java.time.Instant;
@@ -18,6 +19,8 @@
 import java.util.Optional;
 import java.util.Scanner;
 import java.util.function.Function;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
@@ -393,8 +396,14 @@ public void testArtifacts() throws IOException {
         checkArtifactProperties(artifacts.get(0), "artifact1");
         checkArtifactProperties(artifacts.get(1), "artifact2");
 
+        Logger clientLogger = Logger.getLogger(GitHubClient.class.getName());
+
         // Test download from upload-artifact@v3 infrastructure
         String artifactContent = artifacts.get(0).download((is) -> {
+            // At finest log level, all body responses are byte arrays.
+            if (clientLogger.getLevel() != Level.FINEST) {
+                assertThat(is, not(isA(ByteArrayInputStream.class)));
+            }
             try (ZipInputStream zis = new ZipInputStream(is)) {
                 StringBuilder sb = new StringBuilder();