fix: Fix ByteBuf leak in remoting transport.

apache · Dec 28, 2024 · e9fa693 · e9fa693
1 parent eb22d8d
commit e9fa693
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/...-node-testkit/src/main/scala/org/apache/pekko/remote/testconductor/RemoteConnection.scala b/...-node-testkit/src/main/scala/org/apache/pekko/remote/testconductor/RemoteConnection.scala
@@ -56,8 +56,14 @@ private[pekko] class ProtobufEncoder extends MessageToMessageEncoder[Message] {
 private[pekko] class ProtobufDecoder(prototype: Message) extends MessageToMessageDecoder[ByteBuf] {
 
   override def decode(ctx: ChannelHandlerContext, msg: ByteBuf, out: java.util.List[AnyRef]): Unit = {
-    val bytes = ByteBufUtil.getBytes(msg)
-    out.add(prototype.getParserForType.parseFrom(bytes))
+    try {
+      val bytes = ByteBufUtil.getBytes(msg)
+      out.add(prototype.getParserForType.parseFrom(bytes))
+    } catch {
+      case NonFatal(e) => ctx.pipeline().fireExceptionCaught(e)
+    } finally {
+      msg.release()
+    }
   }
 }
 

diff --git a/remote/src/main/scala/org/apache/pekko/remote/transport/netty/TcpSupport.scala b/remote/src/main/scala/org/apache/pekko/remote/transport/netty/TcpSupport.scala
@@ -30,6 +30,8 @@ import io.netty.buffer.{ ByteBuf, ByteBufUtil, Unpooled }
 import io.netty.channel.{ Channel, ChannelHandlerContext }
 import io.netty.util.AttributeKey
 
+import scala.util.control.NonFatal
+
 private[remote] object TcpHandlers {
   private val LISTENER = AttributeKey.valueOf[HandleEventListener]("listener")
 }
@@ -56,8 +58,16 @@ private[remote] trait TcpHandlers extends CommonHandlers {
   }
 
   override def onMessage(ctx: ChannelHandlerContext, msg: ByteBuf): Unit = {
-    val bytes: Array[Byte] = ByteBufUtil.getBytes(msg)
-    if (bytes.length > 0) notifyListener(ctx.channel(), InboundPayload(ByteString(bytes)))
+    try {
+      val bytes: Array[Byte] = ByteBufUtil.getBytes(msg)
+      if (bytes.length > 0) notifyListener(ctx.channel(), InboundPayload(ByteString(bytes)))
+    } catch {
+      case NonFatal(e) =>
+        log.warning("Error while handling message from [{}]: {}", ctx.channel().remoteAddress(), e)
+        onException(ctx, e)
+    } finally {
+      msg.release()
+    }
   }
 
   override def onException(ctx: ChannelHandlerContext, e: Throwable): Unit = {