[improvement] Used more Pooled ByteBufs

Author: Enrico Olivelli

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/KafkaCommandDecoder.java

@@ -584,7 +584,7 @@ protected abstract void channelPrepare(ChannelHandlerContext ctx,
     protected abstract void
     handleCreatePartitions(KafkaHeaderAndRequest kafkaHeaderAndRequest, CompletableFuture<AbstractResponse> response);
 
-    public static class KafkaHeaderAndRequest implements Closeable {
+    public static class KafkaHeaderAndRequest {
 
         private static final String DEFAULT_CLIENT_HOST = "";
 
@@ -593,6 +593,8 @@ public static class KafkaHeaderAndRequest implements Closeable {
         private final ByteBuf buffer;
         private final SocketAddress remoteAddress;
 
+        private final AtomicBoolean released = new AtomicBoolean();
+
         public KafkaHeaderAndRequest(RequestHeader header,
                               AbstractRequest request,
                               ByteBuf buffer,
@@ -604,6 +606,9 @@ public KafkaHeaderAndRequest(RequestHeader header,
         }
 
         public ByteBuf getBuffer() {
+            if (released.get()) {
+                throw new IllegalStateException("Already released");
+            }
             return buffer;
         }
 
@@ -637,9 +642,17 @@ public String toString() {
                 this.header, this.request, this.remoteAddress);
         }
 
-        @Override
+        public void bufferReleased() {
+            if (!released.compareAndSet(false, true)) {
+                throw new IllegalStateException("Already released");
+            }
+        }
+
         public void close() {
-            this.buffer.release();
+           if (!released.compareAndSet(false, true)) {
+               return;
+           }
+           ReferenceCountUtil.safeRelease(this.buffer);
         }
     }
 

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/PendingRequest.java

@@ -13,6 +13,7 @@
  */
 package io.streamnative.pulsar.handlers.kop.coordinator.transaction;
 
+import io.netty.buffer.ByteBuf;
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
@@ -41,7 +42,7 @@ public PendingRequest(final ApiKeys apiKeys,
         this.responseConsumerHandler = responseConsumerHandler;
     }
 
-    public ByteBuffer serialize() {
+    public ByteBuf serialize() {
         return KopResponseUtils.serializeRequest(requestHeader, request);
     }
 

kafka-impl/src/main/java/io/streamnative/pulsar/handlers/kop/coordinator/transaction/TransactionMarkerChannelHandler.java

@@ -73,7 +73,7 @@ public TransactionMarkerChannelHandler(
     private void enqueueRequest(ChannelHandlerContext channel, PendingRequest pendingRequest) {
         final long correlationId = pendingRequest.getCorrelationId();
         pendingRequestMap.put(correlationId, pendingRequest);
-        channel.writeAndFlush(Unpooled.wrappedBuffer(pendingRequest.serialize())).addListener(writeFuture -> {
+        channel.writeAndFlush(pendingRequest.serialize()).addListener(writeFuture -> {
             if (!writeFuture.isSuccess()) {
                 pendingRequest.completeExceptionally(writeFuture.cause());
                 pendingRequestMap.remove(correlationId);

kafka-impl/src/main/java/org/apache/kafka/common/requests/KopResponseUtils.java

@@ -46,8 +46,8 @@ private static ByteBuf serializeWithHeader(AbstractResponse response, ResponseHe
         return serialize(header.data(), header.headerVersion(), response.data(), version);
     }
 
-    public static ByteBuffer serializeRequest(RequestHeader requestHeader, AbstractRequest request) {
-        return RequestUtils.serialize(requestHeader.data(), requestHeader.headerVersion(),
+    public static ByteBuf serializeRequest(RequestHeader requestHeader, AbstractRequest request) {
+        return serialize(requestHeader.data(), requestHeader.headerVersion(),
                 request.data(), request.version());
     }
 

tests/src/test/java/io/streamnative/pulsar/handlers/kop/KafkaCommonTestUtils.java

@@ -115,15 +115,13 @@ public static KafkaCommandDecoder.KafkaHeaderAndRequest buildRequest(AbstractReq
         RequestHeader mockHeader = new RequestHeader(builder.apiKey(), request.version(), "dummy", 1233);
 
 
-        ByteBuffer serializedRequest = KopResponseUtils.serializeRequest(mockHeader, request);
+        ByteBuf byteBuf = KopResponseUtils.serializeRequest(mockHeader, request);
 
-        ByteBuf byteBuf = Unpooled.copiedBuffer(serializedRequest);
-
-        RequestHeader header = RequestHeader.parse(serializedRequest);
+        RequestHeader header = RequestHeader.parse(byteBuf.nioBuffer());
 
         ApiKeys apiKey = header.apiKey();
         short apiVersion = header.apiVersion();
-        AbstractRequest body = AbstractRequest.parseRequest(apiKey, apiVersion, serializedRequest).request;
+        AbstractRequest body = AbstractRequest.parseRequest(apiKey, apiVersion, byteBuf.nioBuffer()).request;
         return new KafkaCommandDecoder.KafkaHeaderAndRequest(header, body, byteBuf, serviceAddress);
     }
 }

tests/src/test/java/io/streamnative/pulsar/handlers/kop/KafkaRequestHandlerTest.java

@@ -188,17 +188,16 @@ public void testByteBufToRequest() {
             correlationId);
 
         // 1. serialize request into ByteBuf
-        ByteBuffer serializedRequest = KopResponseUtils.serializeRequest(header, apiVersionsRequest);
-        int size = serializedRequest.remaining();
-        ByteBuf inputBuf = Unpooled.buffer(size);
-        inputBuf.writeBytes(serializedRequest);
+        ByteBuf serializedRequest = KopResponseUtils.serializeRequest(header, apiVersionsRequest);
 
         // 2. turn Bytebuf into KafkaHeaderAndRequest.
-        KafkaHeaderAndRequest request = handler.byteBufToRequest(inputBuf);
+        KafkaHeaderAndRequest request = handler.byteBufToRequest(serializedRequest, null);
 
         // 3. verify byteBufToRequest works well.
         assertEquals(request.getHeader().data(), header.data());
         assertTrue(request.getRequest() instanceof ApiVersionsRequest);
+
+        request.close();
     }