fix(network): fix potential ByteBuf LEAK in fetch (#3027)

Signed-off-by: Robin Han <hanxvdovehx@gmail.com>

Author: superhx

core/src/main/scala/kafka/network/SocketServer.scala

@@ -414,6 +414,10 @@ object SocketServer {
 
   val ListenerReconfigurableConfigs: Set[String] = Set(SocketServerConfigs.MAX_CONNECTIONS_CONFIG, SocketServerConfigs.MAX_CONNECTION_CREATION_RATE_CONFIG)
 
+  // AutoMQ inject start
+  val MaxInflightRequestsPerConnection = 64;
+  // AutoMQ inject end
+
   def closeSocket(
     channel: SocketChannel,
     logging: Logging
@@ -1073,7 +1077,7 @@ private[kafka] class Processor(
               if (channel != null && channel.isMuted) {
                 val unmute = if (channelContext == null) {
                   true
-                } else if (channelContext.nextCorrelationId.size() < 8 && channelContext.clearQueueFull()) {
+                } else if (channelContext.nextCorrelationId.size() < MaxInflightRequestsPerConnection && channelContext.clearQueueFull()) {
                   true
                 } else {
                   false
@@ -1200,7 +1204,7 @@ private[kafka] class Processor(
                 // AutoMQ will pipeline the requests to accelerate the performance and also keep the request order.
 
                 // Mute the channel if the inflight requests exceed the threshold.
-                if (channelContext.nextCorrelationId.size() >= SocketServerConfigs.MAX_INFLIGHT_REQUESTS_PER_CONNECTION && !channel.isMuted) {
+                if (channelContext.nextCorrelationId.size() >= MaxInflightRequestsPerConnection && !channel.isMuted) {
                   if (isTraceEnabled) {
                     trace(s"Mute channel ${channel.id} because the inflight requests exceed the threshold, inflight count is ${channelContext.nextCorrelationId.size()}.")
                   }
@@ -1248,7 +1252,7 @@ private[kafka] class Processor(
           if (channel.isMuted) {
             val unmute = if (channelContext == null) {
               true
-            } else if (channelContext.nextCorrelationId.size() < 8 && channelContext.clearQueueFull()) {
+            } else if (channelContext.nextCorrelationId.size() < MaxInflightRequestsPerConnection && channelContext.clearQueueFull()) {
               if (isTraceEnabled) {
                 trace(s"Unmute channel ${send.destinationId} because the inflight requests are below the threshold.")
               }
@@ -1289,7 +1293,7 @@ private[kafka] class Processor(
           }
           remove
         })
-        channelContexts.remove(connectionId)
+        removeChannelContext(connectionId)
         // the channel has been closed by the selector but the quotas still need to be updated
         connectionQuotas.dec(listenerName, InetAddress.getByName(remoteHost))
       } catch {
@@ -1330,7 +1334,7 @@ private[kafka] class Processor(
         }
         remove
       })
-      channelContexts.remove(connectionId)
+      removeChannelContext(connectionId)
       //      inflightResponses.remove(connectionId).foreach(updateRequestMetrics)
       // AutoMQ for Kafka inject end
     }
@@ -1442,6 +1446,26 @@ private[kafka] class Processor(
     // AutoMQ for Kafka inject end
   }
 
+  // AutoMQ inject start
+  private def removeChannelContext(connectionId: String): Unit = {
+    val channelContext = channelContexts.remove(connectionId)
+    if (channelContext == null) {
+      return
+    }
+    channelContext.synchronized {
+      channelContext.nextCorrelationId.clear()
+      channelContext.responses.forEach((_, response) => {
+        response match {
+          case sendResponse: SendResponse =>
+            sendResponse.responseSend.release()
+          case _ =>
+        }
+      })
+      channelContext.responses.clear()
+    }
+  }
+  // AutoMQ inject end
+
   private def dequeueResponse(): RequestChannel.Response = {
     val response = responseQueue.poll()
     if (response != null)

core/src/main/scala/kafka/server/FetchSession.scala

@@ -548,6 +548,16 @@ class IncrementalFetchContext(private val time: Time,
       if (session.epoch != expectedEpoch) {
         info(s"Incremental fetch session ${session.id} expected epoch $expectedEpoch, but " +
           s"got ${session.epoch}.  Possible duplicate request.")
+        // AutoMQ inject start
+        // The fetch will return empty data, so we need to release the fetched records.
+        updates.forEach((_, response) => {
+          response.records() match {
+            case r: PooledResource =>
+              r.release()
+            case _ =>
+          }
+        })
+        // AutoMQ inject end
         FetchResponse.of(Errors.INVALID_FETCH_SESSION_EPOCH, 0, session.id, new FetchSession.RESP_MAP)
       } else {
         // Iterate over the update list using PartitionIterator. This will prune updates which don't need to be sent

server/src/main/java/org/apache/kafka/network/SocketServerConfigs.java

@@ -172,10 +172,6 @@ public class SocketServerConfigs {
     public static final int NUM_NETWORK_THREADS_DEFAULT = 3;
     public static final String NUM_NETWORK_THREADS_DOC = "The number of threads that the server uses for receiving requests from the network and sending responses to the network. Noted: each listener (except for controller listener) creates its own thread pool.";
 
-    // AutoMQ inject start
-    public static final int MAX_INFLIGHT_REQUESTS_PER_CONNECTION = 64;
-    // AutoMQ inject end
-
     public static final ConfigDef CONFIG_DEF =  new ConfigDef()
             .define(LISTENERS_CONFIG, STRING, LISTENERS_DEFAULT, HIGH, LISTENERS_DOC)
             .define(ADVERTISED_LISTENERS_CONFIG, STRING, null, HIGH, ADVERTISED_LISTENERS_DOC)