Fix the dead lock issue by not calling resolveAllRemainingChunks in write (#3162)

* Fix the dead lock issue by not calling resolveAllRemainingChunks in write

* refactor

---------

Co-authored-by: Sophie Guo <[email protected]>

Author: SophieGuo410

ambry-commons/src/main/java/com/github/ambry/commons/ByteBufferAsyncWritableChannel.java

@@ -20,6 +20,7 @@
 import java.nio.ByteBuffer;
 import java.nio.channels.ClosedChannelException;
 import java.util.Queue;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -100,14 +101,20 @@ public Future<Long> write(ByteBuf src, Callback<Long> callback) {
     if (src == null) {
       throw new IllegalArgumentException("Source buffer cannot be null");
     }
+    if (!isOpen()) {
+      src.release();
+      CompletableFuture<Long> failedFuture = new CompletableFuture<>();
+      failedFuture.completeExceptionally(new ClosedChannelException());
+      if (callback != null) {
+        callback.onCompletion(0L, new ClosedChannelException());
+      }
+      return failedFuture;
+    }
     ChunkData chunkData = new ChunkData(src, callback);
     chunks.add(chunkData);
     if (channelEventListener != null) {
       channelEventListener.onEvent(EventType.Write);
     }
-    if (!isOpen()) {
-      resolveAllRemainingChunks(new ClosedChannelException());
-    }
     return chunkData.future;
   }
 

ambry-commons/src/test/java/com/github/ambry/commons/ByteBufferAsyncWritableChannelTest.java

@@ -81,6 +81,40 @@ public void commonCaseTest() throws Exception {
     assertNull("There should have been no chunk returned", channel.getNextChunk(0));
   }
 
+  @Test
+  public void writeAfterCloseShouldFailWithClosedChannelException() throws Exception {
+    ByteBufferAsyncWritableChannel channel = new ByteBufferAsyncWritableChannel();
+
+    // Close the channel first
+    channel.close();
+    assertFalse("Channel should be closed", channel.isOpen());
+
+    // Prepare a dummy ByteBuf
+    ByteBuf src = ByteBufAllocator.DEFAULT.heapBuffer(50);
+    src.writeBytes(new byte[]{1, 2, 3});
+
+    WriteCallback callback = new WriteCallback(0);
+
+    // Perform write
+    Future<Long> future = channel.write(src, callback);
+
+    // The future should fail with ClosedChannelException
+    try {
+      future.get();
+      fail("Expected write to fail due to closed channel");
+    } catch (ExecutionException e) {
+      Throwable root = Utils.getRootCause(e);
+      assertTrue("Expected ClosedChannelException but got " + root, root instanceof ClosedChannelException);
+    }
+
+    // The callback should also have received a ClosedChannelException
+    assertTrue("Expected callback exception to be ClosedChannelException but got " + callback.exception,
+        callback.exception instanceof ClosedChannelException);
+
+    // The ByteBuf should have been released
+    assertEquals("Expected src ByteBuf to be released", 0, src.refCnt());
+  }
+
   @Test
   public void commonCaseTestForNettyByteBuf() throws Exception {
     for (boolean useCompositeByteBuf : Arrays.asList(false, true)) {