add support for http response compression

opensearch-project · Jan 26, 2025 · 93a9d76 · 93a9d76
1 parent b256431
commit 93a9d76
Showing 1 changed file with 29 additions and 3 deletions.
diff --git a/mr/src/main/java/org/opensearch/hadoop/rest/commonshttp/CommonsHttpTransport.java b/mr/src/main/java/org/opensearch/hadoop/rest/commonshttp/CommonsHttpTransport.java
@@ -108,6 +108,8 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
+import java.util.zip.GZIPInputStream;
 import java.util.zip.GZIPOutputStream;
 
 /**
@@ -152,12 +154,14 @@ public class CommonsHttpTransport implements Transport, StatsAware {
     private static class ResponseInputStream extends DelegatingInputStream implements ReusableInputStream {
 
         private final HttpMethod method;
+        private final boolean compressed;
         private final boolean reusable;
 
         public ResponseInputStream(HttpMethod http) throws IOException {
-            super(http.getResponseBodyAsStream());
+            super(getStream(http));
             this.method = http;
-            reusable = (delegate() instanceof ByteArrayInputStream);
+            this.compressed = hasCompressedResponseBody(http);
+            this.reusable = (http.getResponseBodyAsStream() instanceof ByteArrayInputStream);
         }
 
         @Override
@@ -173,7 +177,7 @@ public boolean equals(Object obj) {
         @Override
         public InputStream copy() {
             try {
-                return (reusable ? method.getResponseBodyAsStream() : null);
+                return (reusable ? getStream(method, compressed) : null);
             } catch (IOException ex) {
                 throw new OpenSearchHadoopIllegalStateException(ex);
             }
@@ -190,6 +194,23 @@ public void close() throws IOException {
             }
             method.releaseConnection();
         }
+
+        static private InputStream getStream(HttpMethod http) throws IOException {
+            return getStream(http, hasCompressedResponseBody(http));
+        }
+
+        static private InputStream getStream(HttpMethod http, boolean compressed) throws IOException {
+            InputStream responseBodyStream = http.getResponseBodyAsStream();
+            if (compressed) {
+                responseBodyStream = new GZIPInputStream(responseBodyStream);
+            }
+            return responseBodyStream;
+        }
+
+        static private boolean hasCompressedResponseBody(HttpMethod http) {
+            Header contentEncoding = http.getResponseHeader("Content-Encoding");
+            return (contentEncoding != null) && Objects.equals(contentEncoding.getValue(), "gzip");
+        }
     }
 
     private class SocketTrackingConnectionManager extends SimpleHttpConnectionManager {
@@ -692,6 +713,11 @@ public Response execute(Request request) throws IOException {
             entityMethod.setContentChunked(false);
         }
 
+        if (settings.getHttpCompression()) {
+            log.debug("Requesting compressed response");
+            http.setRequestHeader("Accept-Encoding", "gzip");
+        }
+
         headers.applyTo(http);
 
         // We don't want a token added from a proxy user to collide with the