DRILL-8489: Sender memory leak when rpc encode exception

apache · Apr 17, 2024 · 55a9546 · 55a9546
1 parent af7cfcd
commit 55a9546
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/exec/rpc/src/main/java/org/apache/drill/exec/rpc/RpcEncoder.java b/exec/rpc/src/main/java/org/apache/drill/exec/rpc/RpcEncoder.java
@@ -21,11 +21,14 @@
 import io.netty.buffer.ByteBufOutputStream;
 import io.netty.buffer.CompositeByteBuf;
 import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.EncoderException;
 import io.netty.handler.codec.MessageToMessageEncoder;
 
 import java.io.OutputStream;
+import java.util.Arrays;
 import java.util.List;
 
+import io.netty.util.ReferenceCounted;
 import org.apache.drill.exec.proto.GeneralRPCProtos.CompleteRpcMessage;
 import org.apache.drill.exec.proto.GeneralRPCProtos.RpcHeader;
 
@@ -61,7 +64,7 @@ protected void encode(ChannelHandlerContext ctx, OutboundRpcMessage msg, List<Ob
       msg.release();
       return;
     }
-
+    boolean encodeSucess = false;
     try{
       if (RpcConstants.EXTRA_DEBUGGING) {
         logger.debug("Encoding outbound message {}", msg);
@@ -131,9 +134,14 @@ protected void encode(ChannelHandlerContext ctx, OutboundRpcMessage msg, List<Ob
       if (RpcConstants.EXTRA_DEBUGGING) {
         logger.debug("Sent message.  Ending writer index was {}.", buf.writerIndex());
       }
+      encodeSucess = true;
     } finally {
       // make sure to release Rpc Messages underlying byte buffers.
-      //msg.release();
+
+      //if msg instanceof ReferenceCounted,netty can release msg.
+      if (!encodeSucess && !(msg instanceof ReferenceCounted)) {
+        msg.release();
+      }
     }
   }