Avoid closing response stream while downloading instance logs (#14730)

apache · Jan 17, 2025 · 44b07b0 · 44b07b0
1 parent d1ac83e
commit 44b07b0
Showing 1 changed file with 27 additions and 25 deletions.
diff --git a/...roller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerLogger.java b/...roller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerLogger.java
@@ -26,6 +26,7 @@
 import io.swagger.annotations.SecurityDefinition;
 import io.swagger.annotations.SwaggerDefinition;
 import java.io.IOException;
+import java.io.InputStream;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Arrays;
@@ -47,9 +48,9 @@
 import javax.ws.rs.core.HttpHeaders;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
+import javax.ws.rs.core.StreamingOutput;
 import javax.ws.rs.core.UriBuilder;
 import org.apache.commons.collections4.MapUtils;
-import org.apache.commons.io.IOUtils;
 import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse;
 import org.apache.hc.core5.http.HttpVersion;
 import org.apache.hc.core5.http.io.support.ClassicRequestBuilder;
@@ -212,33 +213,34 @@ public Response downloadLogFileFromInstance(
       @ApiParam(value = "Instance Name", required = true) @PathParam("instanceName") String instanceName,
       @ApiParam(value = "Log file path", required = true) @QueryParam("filePath") String filePath,
       @Context Map<String, String> headers) {
-    try {
-      URI uri = UriBuilder.fromUri(getInstanceBaseUri(instanceName)).path("/loggers/download")
-          .queryParam("filePath", filePath).build();
-      ClassicRequestBuilder requestBuilder = ClassicRequestBuilder.get(uri).setVersion(HttpVersion.HTTP_1_1);
-      if (MapUtils.isNotEmpty(headers)) {
-        for (Map.Entry<String, String> header : headers.entrySet()) {
-          requestBuilder.addHeader(header.getKey(), header.getValue());
-        }
+    URI uri = UriBuilder.fromUri(getInstanceBaseUri(instanceName)).path("/loggers/download")
+        .queryParam("filePath", filePath).build();
+    ClassicRequestBuilder requestBuilder = ClassicRequestBuilder.get(uri).setVersion(HttpVersion.HTTP_1_1);
+    if (MapUtils.isNotEmpty(headers)) {
+      for (Map.Entry<String, String> header : headers.entrySet()) {
+        requestBuilder.addHeader(header.getKey(), header.getValue());
       }
-      if (authorization != null) {
-        requestBuilder.addHeader(HttpHeaders.AUTHORIZATION, authorization);
-      }
-      try (CloseableHttpResponse httpResponse = _fileUploadDownloadClient.getHttpClient()
-          .execute(requestBuilder.build())) {
-        if (httpResponse.getCode() >= 400) {
-          throw new WebApplicationException(IOUtils.toString(httpResponse.getEntity().getContent(), "UTF-8"),
-              Response.Status.fromStatusCode(httpResponse.getCode()));
+    }
+    if (authorization != null) {
+      requestBuilder.addHeader(HttpHeaders.AUTHORIZATION, authorization);
+    }
+
+    StreamingOutput streamingOutput = output -> {
+      try (CloseableHttpResponse response = _fileUploadDownloadClient.getHttpClient().execute(requestBuilder.build());
+          InputStream inputStream = response.getEntity().getContent()) {
+        // Stream the data using a buffer
+        byte[] buffer = new byte[1024];
+        int bytesRead;
+        while ((bytesRead = inputStream.read(buffer)) != -1) {
+          output.write(buffer, 0, bytesRead);
         }
-        Response.ResponseBuilder builder = Response.ok();
-        builder.entity(httpResponse.getEntity().getContent());
-        builder.contentLocation(uri);
-        builder.header(HttpHeaders.CONTENT_LENGTH, httpResponse.getEntity().getContentLength());
-        return builder.build();
+        output.flush();
       }
-    } catch (IOException e) {
-      throw new WebApplicationException(e, Response.Status.INTERNAL_SERVER_ERROR);
-    }
+    };
+    Response.ResponseBuilder builder = Response.ok();
+    builder.entity(streamingOutput);
+    builder.contentLocation(uri);
+    return builder.build();
   }
 
   private String getInstanceBaseUri(String instanceName) {