Fix helix-rest memory leak

apache · Nov 6, 2024 · 441c532 · 441c532
1 parent 8c94db9
commit 441c532
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
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
@@ -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.consume(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
@@ -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;
@@ -34,8 +36,10 @@
 import org.apache.http.StatusLine;
 import org.apache.http.client.ClientProtocolException;
 import org.apache.http.client.HttpClient;
+import org.apache.http.client.methods.CloseableHttpResponse;
 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 +173,30 @@ 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
+    // TODO: add a mock rest server
+    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 < 10; 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);