Fix memory leak for batch delete

Author: Sophie Guo

ambry-frontend/src/main/java/com/github/ambry/frontend/s3/S3BatchDeleteHandler.java

@@ -187,14 +187,19 @@ private Callback<Long> parseRequestBodyAndDeleteCallback(RetainingAsyncWritableC
    */
   public S3MessagePayload.S3BatchDeleteObjects deserializeRequest(RetainingAsyncWritableChannel channel)
       throws RestServiceException {
+    ByteBuf byteBuffer = null;
     try {
-      ByteBuf byteBuffer = channel.consumeContentAsByteBuf();
+      byteBuffer = channel.consumeContentAsByteBuf();
       byte[] byteArray = new byte[byteBuffer.readableBytes()];
       byteBuffer.readBytes(byteArray);
       return new XmlMapper().readValue(byteArray, S3MessagePayload.S3BatchDeleteObjects.class);
     } catch (Exception e) {
       logger.trace("s3batchdelete failed to deserialize request");
       throw new RestServiceException("failed to deserialize", e, RestServiceErrorCode.BadRequest);
+    } finally {
+      if (byteBuffer != null) {
+        byteBuffer.release();
+      }
     }
   }
 }

ambry-frontend/src/test/java/com/github/ambry/frontend/S3BatchDeleteHandlerTest.java

@@ -14,6 +14,9 @@
  */
 package com.github.ambry.frontend;
 import com.github.ambry.commons.InMemNamedBlobDbFactory;
+import com.github.ambry.commons.RetainingAsyncWritableChannel;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
 import java.nio.charset.StandardCharsets;
 
 import com.codahale.metrics.MetricRegistry;
@@ -42,17 +45,22 @@
 import com.github.ambry.router.FutureResult;
 import com.github.ambry.router.InMemoryRouter;
 import com.github.ambry.router.ReadableStreamChannel;
+import com.github.ambry.utils.NettyByteBufLeakHelper;
 import com.github.ambry.utils.TestUtils;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.LinkedList;
 import java.util.Properties;
 import org.json.JSONObject;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mockito;
 
 import static com.github.ambry.frontend.s3.S3Constants.*;
 import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
 
 
 public class S3BatchDeleteHandlerTest {
@@ -67,6 +75,8 @@ public class S3BatchDeleteHandlerTest {
   private FrontendConfig frontendConfig;
   private NamedBlobPutHandler namedBlobPutHandler;
   private S3BatchDeleteHandler s3BatchDeleteHandler;
+  private final NettyByteBufLeakHelper nettyByteBufLeakHelper = new NettyByteBufLeakHelper();
+
 
   public S3BatchDeleteHandlerTest() throws Exception {
     account = ACCOUNT_SERVICE.createAndAddRandomAccount();
@@ -83,6 +93,16 @@ public S3BatchDeleteHandlerTest() throws Exception {
     performPutOperation(KEY_NAME_3, KEY_NAME, container, account);
   }
 
+  @Before
+  public void before() {
+    nettyByteBufLeakHelper.beforeTest();
+  }
+
+  @After
+  public void after() {
+    nettyByteBufLeakHelper.afterTest();
+  }
+
   @Test
   public void deleteObjectTest() throws Exception {
     String uri = String.format("/s3/%s/%s", account.getName(), container.getName());
@@ -123,6 +143,74 @@ public void deleteObjectTest() throws Exception {
     assertEquals("Mismatch on status", ResponseStatus.Ok, restResponseChannel.getStatus());
   }
 
+  @Test
+  public void testBufferReleaseInDeserializeRequest() throws Exception {
+    // Arrange
+    S3BatchDeleteHandler handler = new S3BatchDeleteHandler(null, null);
+    RetainingAsyncWritableChannel channel = mock(RetainingAsyncWritableChannel.class);
+    ByteBuf mockByteBuf = mock(ByteBuf.class);
+
+    String validXml = "<?xml version=\"1.0\" encoding=\"UTF-8\"?>" +
+        "<Delete xmlns=\"http://s3.amazonaws.com/doc/2006-03-01\">" +
+        "<Object>" +
+        "<Key>key-success</Key>" +
+        "</Object>" +
+        "<Object>" +
+        "<Key>key-error</Key>" +
+        "</Object>" +
+        "<Object>" +
+        "<Key>key-error2</Key>" +
+        "</Object>" +
+        "<Object>" +
+        "<Key>key-success-2</Key>" +
+        "</Object>" +
+        "</Delete>";
+
+    byte[] xmlBytes = validXml.getBytes();
+
+    when(channel.consumeContentAsByteBuf()).thenReturn(mockByteBuf);
+    when(mockByteBuf.readableBytes()).thenReturn(xmlBytes.length);
+    doAnswer(invocation -> {
+      byte[] buffer = invocation.getArgument(0);
+      System.arraycopy(xmlBytes, 0, buffer, 0, xmlBytes.length);
+      return null;
+    }).when(mockByteBuf).readBytes(Mockito.any(byte[].class));
+
+    S3MessagePayload.S3BatchDeleteObjects result = handler.deserializeRequest(channel);
+
+    assertNotNull(result);
+    verify(mockByteBuf, times(1)).release(); // Verify that release() was called
+  }
+
+  @Test
+  public void testDeserializeRequestWithRealChannel() throws Exception {
+    // Arrange
+    String validXml =
+        "<?xml version=\"1.0\" encoding=\"UTF-8\"?>" + "<Delete xmlns=\"http://s3.amazonaws.com/doc/2006-03-01\">"
+            + "<Object>" + "<Key>key-success</Key>" + "</Object>" + "<Object>" + "<Key>key-error</Key>" + "</Object>"
+            + "</Delete>";
+
+    ByteBuf byteBuf = Unpooled.wrappedBuffer(validXml.getBytes());
+    RetainingAsyncWritableChannel channel = new RetainingAsyncWritableChannel();
+    channel.write(byteBuf, (result, exception) -> {
+      if (exception != null) {
+        fail("Failed to write to channel: " + exception.getMessage());
+      }
+    });
+
+    S3BatchDeleteHandler handler = new S3BatchDeleteHandler(null, null);
+
+    // Act
+    S3MessagePayload.S3BatchDeleteObjects result = handler.deserializeRequest(channel);
+
+    // Assert
+    assertNotNull(result);
+    assertEquals(2, result.getObjects().size());
+    assertEquals("key-success", result.getObjects().get(0).getKey());
+    assertEquals("key-error", result.getObjects().get(1).getKey());
+  }
+
+
   @Test
   public void malformedXMLRequestTest() throws Exception {
     String uri = String.format("/s3/%s/%s", account.getName(), container.getName());