From e8d5247cf65fd37077ab35495506e66c4d37d5a5 Mon Sep 17 00:00:00 2001
From: He-Pin <hepin1989@gmail.com>
Date: Sat, 28 Dec 2024 22:29:50 +0800
Subject: [PATCH] fix: Fix ByteBuf leak in remoting transport.

---
 .../remote/testconductor/RemoteConnection.scala | 11 +++++++++--
 .../remote/transport/netty/TcpSupport.scala     | 17 ++++++++++++-----
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/multi-node-testkit/src/main/scala/org/apache/pekko/remote/testconductor/RemoteConnection.scala b/multi-node-testkit/src/main/scala/org/apache/pekko/remote/testconductor/RemoteConnection.scala
index 99801878bf9..8b857ff9c3a 100644
--- a/multi-node-testkit/src/main/scala/org/apache/pekko/remote/testconductor/RemoteConnection.scala
+++ b/multi-node-testkit/src/main/scala/org/apache/pekko/remote/testconductor/RemoteConnection.scala
@@ -56,8 +56,15 @@ 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) =>
+        throw new RuntimeException("Error while decoding protobuf message", 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
index 8d8fab344a1..cac1dc9b865 100644
--- 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
@@ -14,10 +14,8 @@
 package org.apache.pekko.remote.transport.netty
 
 import java.net.InetSocketAddress
-
 import scala.annotation.nowarn
 import scala.concurrent.{ Future, Promise }
-
 import org.apache.pekko
 import pekko.actor.Address
 import pekko.event.LoggingAdapter
@@ -25,11 +23,12 @@ import pekko.remote.transport.AssociationHandle
 import pekko.remote.transport.AssociationHandle.{ Disassociated, HandleEvent, HandleEventListener, InboundPayload }
 import pekko.remote.transport.Transport.AssociationEventListener
 import pekko.util.ByteString
-
 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 +55,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 = {