add new test

Author: m1a2st

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

@@ -148,6 +148,8 @@ class KafkaRequestHandler(
             if (originalRequest.callbackRequestCompleteTimeNanos.isEmpty)
               originalRequest.callbackRequestCompleteTimeNanos = Some(time.nanoseconds())
             threadCurrentRequest.remove()
+            if (!originalRequest.header.apiKey().forwardable)
+              originalRequest.releaseBuffer()
           }
 
         case request: RequestChannel.Request =>
@@ -163,9 +165,8 @@ class KafkaRequestHandler(
             case e: Throwable => error("Exception when handling request", e)
           } finally {
             threadCurrentRequest.remove()
-            if (!request.header.apiKey().forwardable) {
+            if (!request.header.apiKey().forwardable)
               request.releaseBuffer()
-            }
           }
 
         case RequestChannel.WakeupRequest => 

core/src/test/scala/kafka/server/KafkaRequestHandlerTest.scala

@@ -231,16 +231,23 @@ class KafkaRequestHandlerTest {
     })
   }
 
-  def makeRequest(time: Time, metrics: RequestChannelMetrics): RequestChannel.Request = {
+  def makeRequest(
+    time: Time, 
+    metrics: RequestChannelMetrics,
+    apiKeys: ApiKeys = ApiKeys.API_VERSIONS,
+    version: Short = 0,
+    buffer: ByteBuffer = ByteBuffer.allocate(0),
+    memoryPool: MemoryPool = mock(classOf[MemoryPool])
+  ): RequestChannel.Request = {
     // Make unsupported API versions request to avoid having to parse a real request
     val requestHeader = mock(classOf[RequestHeader])
-    when(requestHeader.apiKey()).thenReturn(ApiKeys.API_VERSIONS)
-    when(requestHeader.apiVersion()).thenReturn(0.toShort)
+    when(requestHeader.apiKey()).thenReturn(apiKeys)
+    when(requestHeader.apiVersion()).thenReturn(version)
 
     val context = new RequestContext(requestHeader, "0", mock(classOf[InetAddress]), new KafkaPrincipal("", ""),
       new ListenerName(""), SecurityProtocol.PLAINTEXT, mock(classOf[ClientInformation]), false)
     new RequestChannel.Request(0, context, time.nanoseconds(),
-      mock(classOf[MemoryPool]), ByteBuffer.allocate(0), metrics)
+      memoryPool, buffer, metrics)
   }
 
   def setupBrokerTopicMetrics(systemRemoteStorageEnabled: Boolean = true): BrokerTopicMetrics = {
@@ -699,4 +706,32 @@ class KafkaRequestHandlerTest {
     // cleanup
     brokerTopicStats.close()
   }
+
+  @Test
+  def testRequestBufferRelease(): Unit = {
+    val time = new MockTime()
+    val metrics = new RequestChannelMetrics(Collections.emptySet[ApiKeys])
+    val requestChannel = new RequestChannel(10, time, metrics)
+    val apiHandler = mock(classOf[ApiRequestHandler])
+    val memoryPool = mock(classOf[MemoryPool])
+    val buffer = ByteBuffer.allocate(1024)
+
+    val handler = new KafkaRequestHandler(0, 0, mock(classOf[Meter]), new AtomicInteger(1), requestChannel, apiHandler, time)
+
+    val request = makeRequest(time, metrics, ApiKeys.PRODUCE, 3, buffer, memoryPool)
+    requestChannel.sendRequest(request)
+
+    val shutdownThread = new Thread(() => {
+      try {
+        Thread.sleep(1000)
+        requestChannel.sendShutdownRequest()
+      } catch {
+        case _: InterruptedException =>
+      }
+    })
+
+    shutdownThread.start()
+    handler.run()
+    verify(memoryPool, times(1)).release(buffer)
+  }
 }