From 2b54fbb8c01dd74cecfebc412d1ea88a2cd7c677 Mon Sep 17 00:00:00 2001
From: Xiaxuan Gao <xiaxgao@linkedin.com>
Date: Fri, 8 Nov 2024 09:41:43 -0800
Subject: [PATCH] Fix helix-rest memory leak (#2960)

Fix helix-rest stoppable check memory leak
---
 .../rest/client/CustomRestClientImpl.java     |  1 +
 .../rest/client/TestCustomRestClient.java     | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/helix-rest/src/main/java/org/apache/helix/rest/client/CustomRestClientImpl.java b/helix-rest/src/main/java/org/apache/helix/rest/client/CustomRestClientImpl.java
index 09783c8dd1..db943b3ddb 100644
--- a/helix-rest/src/main/java/org/apache/helix/rest/client/CustomRestClientImpl.java
+++ b/helix-rest/src/main/java/org/apache/helix/rest/client/CustomRestClientImpl.java
@@ -202,6 +202,7 @@ protected HttpResponse post(String url, Map<String, Object> payloads) throws IOE
             LOG.warn("Received 302 but no Location header is present, stopping retries.");
             break;  // Break out if there is no valid redirect location
           }
+          EntityUtils.consumeQuietly(response.getEntity());
         } else {
           LOG.warn("Received non-200 and non-302 status code: {}, payloads: {}", status, payloads);
           return response;  // Return response without retry
diff --git a/helix-rest/src/test/java/org/apache/helix/rest/client/TestCustomRestClient.java b/helix-rest/src/test/java/org/apache/helix/rest/client/TestCustomRestClient.java
index 5a5fec04bb..3d81dc432a 100644
--- a/helix-rest/src/test/java/org/apache/helix/rest/client/TestCustomRestClient.java
+++ b/helix-rest/src/test/java/org/apache/helix/rest/client/TestCustomRestClient.java
@@ -20,10 +20,12 @@
  */
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -36,6 +38,7 @@
 import org.apache.http.client.HttpClient;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.conn.ConnectTimeoutException;
+import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.client.HttpClients;
 import org.junit.Assert;
 import org.mockito.Mock;
@@ -169,6 +172,29 @@ public void testPostRequestFormat() throws IOException {
     }
   }
 
+  @Test (description = "Validate if the post request has memory leak or no")
+  public void testMultiplePost() throws IOException {
+    // a popular echo server that echos all the inputs
+    final String echoServer = "https://httpbin.org/redirect-to?url=http://httpbin.org/post";
+    HttpClientBuilder httpClientBuilder = HttpClients.custom()
+        .evictExpiredConnections()
+        .setMaxConnPerRoute(1)
+        .evictIdleConnections(Duration.ofSeconds(30).toMillis(), TimeUnit.MILLISECONDS);
+    HttpClient httpClient = httpClientBuilder.build();
+    CustomRestClientImpl customRestClient = new CustomRestClientImpl(httpClient);
+    HttpResponse response;
+
+    for (int i = 0; i < 5; i++) {
+      response = customRestClient.post(echoServer, Collections.emptyMap());
+      if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
+        JsonNode json = customRestClient.getJsonObject(response);
+
+        Assert.assertEquals(json.get("headers").get("Accept").asText(), "application/json");
+        Assert.assertEquals(json.get("data").asText(), "{}");
+      }
+    }
+  }
+
   @Test
   public void testGetPartitionStoppableCheckWhenTimeout() throws IOException {
     MockCustomRestClient customRestClient = new MockCustomRestClient(_httpClient);