fix for bug: #3640

kumjiten · kumjiten · commit 303e408efcd7 · 2022-06-29T12:00:26.000+05:30
Signed-off-by: Jitendra Kumar &lt;jkjitend@amazon.com&gt;
diff --git a/client/rest/src/main/java/org/opensearch/client/RestClient.java b/client/rest/src/main/java/org/opensearch/client/RestClient.java
@@ -131,6 +131,7 @@ public class RestClient implements Closeable {
     private volatile NodeTuple<List<Node>> nodeTuple;
     private final WarningsHandler warningsHandler;
     private final boolean compressionEnabled;
+    private final boolean chunkedTransferEncodingEnabled;
 
     RestClient(
         CloseableHttpAsyncClient client,
@@ -141,6 +142,20 @@ public class RestClient implements Closeable {
         NodeSelector nodeSelector,
         boolean strictDeprecationMode,
         boolean compressionEnabled
+    ) {
+        this(client, defaultHeaders, nodes, pathPrefix, failureListener, nodeSelector, strictDeprecationMode, compressionEnabled, true);
+    }
+
+    RestClient(
+        CloseableHttpAsyncClient client,
+        Header[] defaultHeaders,
+        List<Node> nodes,
+        String pathPrefix,
+        FailureListener failureListener,
+        NodeSelector nodeSelector,
+        boolean strictDeprecationMode,
+        boolean compressionEnabled,
+        boolean chunkedTransferEncodingEnabled
     ) {
         this.client = client;
         this.defaultHeaders = Collections.unmodifiableList(Arrays.asList(defaultHeaders));
@@ -149,6 +164,7 @@ public class RestClient implements Closeable {
         this.nodeSelector = nodeSelector;
         this.warningsHandler = strictDeprecationMode ? WarningsHandler.STRICT : WarningsHandler.PERMISSIVE;
         this.compressionEnabled = compressionEnabled;
+        this.chunkedTransferEncodingEnabled = chunkedTransferEncodingEnabled;
         setNodes(nodes);
     }
 
@@ -583,36 +599,51 @@ private static void addSuppressedException(Exception suppressedException, Except
         }
     }
 
-    private static HttpRequestBase createHttpRequest(String method, URI uri, HttpEntity entity, boolean compressionEnabled) {
+    private static HttpRequestBase createHttpRequest(
+        String method,
+        URI uri,
+        HttpEntity entity,
+        boolean compressionEnabled,
+        boolean chunkedTransferEncodingEnabled
+    ) {
         switch (method.toUpperCase(Locale.ROOT)) {
             case HttpDeleteWithEntity.METHOD_NAME:
-                return addRequestBody(new HttpDeleteWithEntity(uri), entity, compressionEnabled);
+                return addRequestBody(new HttpDeleteWithEntity(uri), entity, compressionEnabled, chunkedTransferEncodingEnabled);
             case HttpGetWithEntity.METHOD_NAME:
-                return addRequestBody(new HttpGetWithEntity(uri), entity, compressionEnabled);
+                return addRequestBody(new HttpGetWithEntity(uri), entity, compressionEnabled, chunkedTransferEncodingEnabled);
             case HttpHead.METHOD_NAME:
-                return addRequestBody(new HttpHead(uri), entity, compressionEnabled);
+                return addRequestBody(new HttpHead(uri), entity, compressionEnabled, chunkedTransferEncodingEnabled);
             case HttpOptions.METHOD_NAME:
-                return addRequestBody(new HttpOptions(uri), entity, compressionEnabled);
+                return addRequestBody(new HttpOptions(uri), entity, compressionEnabled, chunkedTransferEncodingEnabled);
             case HttpPatch.METHOD_NAME:
-                return addRequestBody(new HttpPatch(uri), entity, compressionEnabled);
+                return addRequestBody(new HttpPatch(uri), entity, compressionEnabled, chunkedTransferEncodingEnabled);
             case HttpPost.METHOD_NAME:
                 HttpPost httpPost = new HttpPost(uri);
-                addRequestBody(httpPost, entity, compressionEnabled);
+                addRequestBody(httpPost, entity, compressionEnabled, chunkedTransferEncodingEnabled);
                 return httpPost;
             case HttpPut.METHOD_NAME:
-                return addRequestBody(new HttpPut(uri), entity, compressionEnabled);
+                return addRequestBody(new HttpPut(uri), entity, compressionEnabled, chunkedTransferEncodingEnabled);
             case HttpTrace.METHOD_NAME:
-                return addRequestBody(new HttpTrace(uri), entity, compressionEnabled);
+                return addRequestBody(new HttpTrace(uri), entity, compressionEnabled, chunkedTransferEncodingEnabled);
             default:
                 throw new UnsupportedOperationException("http method not supported: " + method);
         }
     }
 
-    private static HttpRequestBase addRequestBody(HttpRequestBase httpRequest, HttpEntity entity, boolean compressionEnabled) {
+    private static HttpRequestBase addRequestBody(
+        HttpRequestBase httpRequest,
+        HttpEntity entity,
+        boolean compressionEnabled,
+        boolean chunkedTransferEncodingEnabled
+    ) {
         if (entity != null) {
             if (httpRequest instanceof HttpEntityEnclosingRequestBase) {
                 if (compressionEnabled) {
-                    entity = new ContentCompressingEntity(entity);
+                    if (chunkedTransferEncodingEnabled) {
+                        entity = new ContentCompressingChunkedEntity(entity);
+                    } else {
+                        entity = new ContentCompressingEntity(entity);
+                    }
                 }
                 ((HttpEntityEnclosingRequestBase) httpRequest).setEntity(entity);
             } else {
@@ -782,7 +813,13 @@ private class InternalRequest {
             String ignoreString = params.remove("ignore");
             this.ignoreErrorCodes = getIgnoreErrorCodes(ignoreString, request.getMethod());
             URI uri = buildUri(pathPrefix, request.getEndpoint(), params);
-            this.httpRequest = createHttpRequest(request.getMethod(), uri, request.getEntity(), compressionEnabled);
+            this.httpRequest = createHttpRequest(
+                request.getMethod(),
+                uri,
+                request.getEntity(),
+                compressionEnabled,
+                chunkedTransferEncodingEnabled
+            );
             this.cancellable = Cancellable.fromRequest(httpRequest);
             setHeaders(httpRequest, request.getOptions().getHeaders());
             setRequestConfig(httpRequest, request.getOptions().getRequestConfig());
@@ -932,6 +969,30 @@ private static Exception extractAndWrapCause(Exception exception) {
         return new RuntimeException("error while performing request", exception);
     }
 
+    /**
+     * A gzip compressing entity that also implements {@code getContent()}.
+     */
+    public static class ContentCompressingChunkedEntity extends GzipCompressingEntity {
+
+        /**
+         * Creates a {@link ContentCompressingChunkedEntity} instance with the provided HTTP entity.
+         *
+         * @param entity the HTTP entity.
+         */
+        public ContentCompressingChunkedEntity(HttpEntity entity) {
+            super(entity);
+        }
+
+        @Override
+        public InputStream getContent() throws IOException {
+            ByteArrayInputOutputStream out = new ByteArrayInputOutputStream(1024);
+            try (GZIPOutputStream gzipOut = new GZIPOutputStream(out)) {
+                wrappedEntity.writeTo(gzipOut);
+            }
+            return out.asInput();
+        }
+    }
+
     /**
      * A gzip compressing entity that also implements {@code getContent()}.
      */
diff --git a/client/rest/src/main/java/org/opensearch/client/RestClientBuilder.java b/client/rest/src/main/java/org/opensearch/client/RestClientBuilder.java
@@ -84,6 +84,7 @@ public final class RestClientBuilder {
     private NodeSelector nodeSelector = NodeSelector.ANY;
     private boolean strictDeprecationMode = false;
     private boolean compressionEnabled = false;
+    private boolean chunkedTransferEncodingEnabled = true;
 
     /**
      * Creates a new builder instance and sets the hosts that the client will send requests to.
@@ -238,6 +239,16 @@ public RestClientBuilder setCompressionEnabled(boolean compressionEnabled) {
         return this;
     }
 
+    /**
+     * Whether the REST client should use Transfer-Encoding: chunked for compress requests"
+     *
+     * @param chunkedTransferEncodingEnabled flag for enabling Transfer-Encoding: chunked
+     */
+    public RestClientBuilder setChunkedTransferEncodingEnabled(boolean chunkedTransferEncodingEnabled) {
+        this.chunkedTransferEncodingEnabled = chunkedTransferEncodingEnabled;
+        return this;
+    }
+
     /**
      * Creates a new {@link RestClient} based on the provided configuration.
      */
@@ -256,7 +267,8 @@ public RestClient build() {
             failureListener,
             nodeSelector,
             strictDeprecationMode,
-            compressionEnabled
+            compressionEnabled,
+            chunkedTransferEncodingEnabled
         );
         httpClient.start();
         return restClient;
diff --git a/client/rest/src/test/java/org/opensearch/client/RestClientCompressionTests.java b/client/rest/src/test/java/org/opensearch/client/RestClientCompressionTests.java
@@ -92,7 +92,9 @@ public void handle(HttpExchange exchange) throws IOException {
         }
     }
 
-    /** Read all bytes of an input stream and close it. */
+    /**
+     * Read all bytes of an input stream and close it.
+     */
     private static byte[] readAll(InputStream in) throws IOException {
         byte[] buffer = new byte[1024];
         ByteArrayOutputStream bos = new ByteArrayOutputStream();
@@ -104,15 +106,16 @@ private static byte[] readAll(InputStream in) throws IOException {
         return bos.toByteArray();
     }
 
-    private RestClient createClient(boolean enableCompression) {
+    private RestClient createClient(boolean enableCompression, boolean chunkedEnabled) {
         InetSocketAddress address = httpServer.getAddress();
         return RestClient.builder(new HttpHost(address.getHostString(), address.getPort(), "http"))
             .setCompressionEnabled(enableCompression)
+            .setChunkedTransferEncodingEnabled(chunkedEnabled)
             .build();
     }
 
     public void testCompressingClientWithContentLengthSync() throws Exception {
-        RestClient restClient = createClient(true);
+        RestClient restClient = createClient(true, false);
 
         Request request = new Request("POST", "/");
         request.setEntity(new StringEntity("compressing client", ContentType.TEXT_PLAIN));
@@ -129,9 +132,7 @@ public void testCompressingClientWithContentLengthSync() throws Exception {
 
     public void testCompressingClientContentLengthAsync() throws Exception {
         InetSocketAddress address = httpServer.getAddress();
-        RestClient restClient = RestClient.builder(new HttpHost(address.getHostString(), address.getPort(), "http"))
-            .setCompressionEnabled(true)
-            .build();
+        RestClient restClient = createClient(true, false);
 
         Request request = new Request("POST", "/");
         request.setEntity(new StringEntity("compressing client", ContentType.TEXT_PLAIN));
