Issue 134: produce optimize (#150)

optimize for #134 including the following aspects:

recordToEntry paralleled
publish message to bookie asynchronous
In our environment, produce delay is about one-fifth of the previous

Author: ltamber

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

@@ -19,7 +19,8 @@
 import static io.streamnative.pulsar.handlers.kop.KafkaProtocolHandler.ListenerType.SSL;
 import static io.streamnative.pulsar.handlers.kop.KafkaProtocolHandler.getKopBrokerUrl;
 import static io.streamnative.pulsar.handlers.kop.KafkaProtocolHandler.getListenerPort;
-import static io.streamnative.pulsar.handlers.kop.MessagePublishContext.publishMessages;
+import static io.streamnative.pulsar.handlers.kop.MessagePublishContext.MESSAGE_BATCHED;
+import static io.streamnative.pulsar.handlers.kop.utils.MessageRecordUtils.recordsToByteBuf;
 import static io.streamnative.pulsar.handlers.kop.utils.TopicNameUtils.getKafkaTopicNameFromPulsarTopicname;
 import static io.streamnative.pulsar.handlers.kop.utils.TopicNameUtils.pulsarTopicName;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -29,6 +30,7 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Queues;
+import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
 import io.streamnative.pulsar.handlers.kop.coordinator.group.GroupCoordinator;
 import io.streamnative.pulsar.handlers.kop.coordinator.group.GroupMetadata.GroupOverview;
@@ -38,6 +40,7 @@
 import io.streamnative.pulsar.handlers.kop.utils.MessageIdUtils;
 import io.streamnative.pulsar.handlers.kop.utils.OffsetFinder;
 import io.streamnative.pulsar.handlers.kop.utils.SaslUtils;
+
 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.net.URI;
@@ -60,6 +63,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import javax.naming.AuthenticationException;
+
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.mledger.AsyncCallbacks;
@@ -496,92 +500,81 @@ protected void handleTopicMetadataRequest(KafkaHeaderAndRequest metadataHar,
     // whether the head of queue is running.
     private AtomicBoolean isHeadRequestRun = new AtomicBoolean(false);
 
-    private void handleProducerRequestInternal() {
-        // the first request that success set to running, get running.
-        if (produceRequestsQueue.isEmpty() || !isHeadRequestRun.compareAndSet(false, true)) {
-            // the head of queue is already running, when head complete, it will peek the following request to run.
-            if (log.isDebugEnabled()) {
-                log.debug(" Produce messages not entered. queue.size: {}, head isHeadRequestRun: {}",
-                    produceRequestsQueue.size(), isHeadRequestRun.get());
-            }
-            return;
-        }
-
-        Pair<KafkaHeaderAndRequest, CompletableFuture<AbstractResponse>> head = produceRequestsQueue.peek();
-        KafkaHeaderAndRequest produceHar = head.getKey();
-        CompletableFuture<AbstractResponse> resultFuture = head.getValue();
-        ProduceRequest produceRequest = (ProduceRequest) produceHar.getRequest();
-
-        // Ignore request.acks() and request.timeout(), which related to kafka replication in this broker.
-        Map<TopicPartition, CompletableFuture<PartitionResponse>> responsesFutures = new HashMap<>();
+    private ConcurrentHashMap<TopicName, Queue<Pair<CompletableFuture<ByteBuf>, CompletableFuture<PartitionResponse>>>>
+            transQueue = new ConcurrentHashMap<>();
 
-        final int responsesSize = produceRequest.partitionRecordsOrFail().size();
-
-        // TODO: handle un-exist topic:
-        //     nonExistingTopicResponses += topicPartition -> new PartitionResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION)
-        for (Map.Entry<TopicPartition, ? extends Records> entry : produceRequest.partitionRecordsOrFail().entrySet()) {
-            TopicPartition topicPartition = entry.getKey();
-
-            CompletableFuture<PartitionResponse> partitionResponse = new CompletableFuture<>();
-            responsesFutures.put(topicPartition, partitionResponse);
+    private void publishMessages(MemoryRecords records,
+                                 TopicName topic,
+                                 CompletableFuture<PartitionResponse> future,
+                                 CompletableFuture<ByteBuf> transFuture) {
+        // get records size.
+        AtomicInteger size = new AtomicInteger(0);
+        records.records().forEach(record -> size.incrementAndGet());
+        int rec = size.get();
 
-            if (log.isDebugEnabled()) {
-                log.debug("[{}] Request {}: Produce messages for topic {} partition {}, request size: {} ",
-                    ctx.channel(), produceHar.getHeader(),
-                    topicPartition.topic(), topicPartition.partition(), responsesSize);
-            }
+        if (log.isDebugEnabled()) {
+            log.debug("publishMessages for topic partition: {} , records size is {} ", topic.toString(), size.get());
+        }
 
-            TopicName topicName = pulsarTopicName(topicPartition, namespace);
+        if (MESSAGE_BATCHED) {
+            pulsarService.getExecutor().submit(() -> {
+                ByteBuf buf = recordsToByteBuf(records, rec);
+                transFuture.complete(buf);
+            });
 
-            topicManager.getTopic(topicName.toString()).whenComplete((persistentTopic, exception) -> {
-                if (exception != null || persistentTopic == null) {
-                    log.warn("[{}] Request {}: Failed to getOrCreateTopic {}. "
-                            + "Topic is in loading status, return LEADER_NOT_AVAILABLE. exception:",
-                        ctx.channel(), produceHar.getHeader(), topicName, exception);
-                    partitionResponse.complete(new PartitionResponse(Errors.LEADER_NOT_AVAILABLE));
+            transFuture.whenComplete((headerAndPayload, ex) -> {
+                if (ex != null) {
+                    log.error("record to bytebuf error: ", ex);
+                    future.complete(new PartitionResponse(Errors.KAFKA_STORAGE_ERROR));
                 } else {
-                    CompletableFuture<PersistentTopic> topicFuture = new CompletableFuture<>();
-                    topicFuture.complete(persistentTopic);
-                    publishMessages((MemoryRecords) entry.getValue(), persistentTopic, partitionResponse);
+                    doPublishMessages(topic);
                 }
             });
         }
+    }
 
-        CompletableFuture.allOf(responsesFutures.values().toArray(new CompletableFuture<?>[responsesSize]))
-            .whenComplete((ignore, ex) -> {
-                // all ex has translated to PartitionResponse with Errors.KAFKA_STORAGE_ERROR
-                Map<TopicPartition, PartitionResponse> responses = new ConcurrentHashMap<>();
-                for (Map.Entry<TopicPartition, CompletableFuture<PartitionResponse>> entry:
-                    responsesFutures.entrySet()) {
-                    responses.put(entry.getKey(), entry.getValue().join());
-                }
+    private void doPublishMessages(TopicName topic) {
+        Queue<Pair<CompletableFuture<ByteBuf>, CompletableFuture<PartitionResponse>>> topicQueue =
+                transQueue.get(topic);
 
+        // loop from first responseFuture.
+        while (topicQueue != null && topicQueue.peek() != null
+                && topicQueue.peek().getLeft().isDone() && isActive.get()) {
+            CompletableFuture<Long> offsetFuture = new CompletableFuture<>();
+            Pair<CompletableFuture<ByteBuf>, CompletableFuture<PartitionResponse>> result = topicQueue.remove();
+            try {
                 if (log.isDebugEnabled()) {
-                    log.debug("[{}] Request {}: Complete handle produce.",
-                        ctx.channel(), produceHar.toString());
+                    log.debug("[{}] topic", topic.toString());
                 }
-                resultFuture.complete(new ProduceResponse(responses));
-            });
-
-        // trigger following request to run.
-        resultFuture.whenComplete((response, throwable) -> {
-            if (throwable != null) {
-                log.warn("Error produce message for {}.", produceHar.getHeader(), throwable);
-            }
+                ByteBuf headerAndPayload = result.getLeft().get();
+                topicManager.getTopic(topic.toString()).whenComplete((persistentTopic, throwable) -> {
+                    if (throwable != null || persistentTopic == null) {
+                        log.warn("[{}] Request {}: Failed to getOrCreateTopic {}. "
+                                        + "Topic is in loading status, return LEADER_NOT_AVAILABLE. exception:",
+                                ctx.channel(), topic.toString(), throwable);
+                        result.getRight().complete(new PartitionResponse(Errors.LEADER_NOT_AVAILABLE));
+                    } else {
+                        persistentTopic.publishMessage(
+                                headerAndPayload,
+                                MessagePublishContext.get(
+                                        offsetFuture, persistentTopic, System.nanoTime()));
+                    }
+                });
 
-            if (log.isDebugEnabled()) {
-                log.debug("Produce messages complete. trigger next. queue.size: {}, head isHeadRequestRun: {}",
-                    produceRequestsQueue.size(), isHeadRequestRun.get());
+                offsetFuture.whenComplete((offset, ex) -> {
+                    if (ex != null) {
+                        log.error("publishMessages for topic partition: {} failed when write.",
+                                topic.toString(), ex);
+                        result.getRight().complete(new PartitionResponse(Errors.KAFKA_STORAGE_ERROR));
+                    } else {
+                        result.getRight().complete(new PartitionResponse(Errors.NONE));
+                    }
+                });
+            } catch (Exception e) {
+                // should not comes here.
+                log.error("error to get Response ByteBuf:", e);
             }
-
-            boolean compare = isHeadRequestRun.compareAndSet(true, false);
-            checkState(compare, "Head should be running when completed head.");
-            // remove completed request.
-            produceRequestsQueue.remove();
-
-            // trigger another run.
-            handleProducerRequestInternal();
-        });
+        }
     }
 
     protected void handleProduceRequest(KafkaHeaderAndRequest produceHar,
@@ -596,13 +589,59 @@ protected void handleProduceRequest(KafkaHeaderAndRequest produceHar,
             return;
         }
 
-        if (log.isDebugEnabled()) {
-            log.debug(" new produce request comes: {}, isHeadRequestRun: {}",
-                produceRequestsQueue.size(), isHeadRequestRun.get());
+        Map<TopicPartition, CompletableFuture<PartitionResponse>> responsesFutures = new HashMap<>();
+
+        final int responsesSize = produceRequest.partitionRecordsOrFail().size();
+
+        // TODO: handle un-exist topic:
+        //     nonExistingTopicResponses += topicPartition -> new PartitionResponse(Errors.UNKNOWN_TOPIC_OR_PARTITION)
+        for (Map.Entry<TopicPartition, ? extends Records> entry : produceRequest.partitionRecordsOrFail().entrySet()) {
+            TopicPartition topicPartition = entry.getKey();
+
+            CompletableFuture<PartitionResponse> partitionResponse = new CompletableFuture<>();
+            responsesFutures.put(topicPartition, partitionResponse);
+
+            if (log.isDebugEnabled()) {
+                log.debug("[{}] Request {}: Produce messages for topic {} partition {}, request size: {} ",
+                        ctx.channel(), produceHar.getHeader(),
+                        topicPartition.topic(), topicPartition.partition(), responsesSize);
+            }
+
+            TopicName topicName = pulsarTopicName(topicPartition, namespace);
+
+            CompletableFuture<ByteBuf> transFuture = new CompletableFuture<>();
+            //put queue
+            transQueue.compute(topicName, (key, queue) -> {
+                if (queue == null) {
+                    Queue<Pair<CompletableFuture<ByteBuf>, CompletableFuture<PartitionResponse>>> newQueue =
+                            Queues.newConcurrentLinkedQueue();
+                    newQueue.add(Pair.of(transFuture, partitionResponse));
+                    return newQueue;
+                } else {
+                    queue.add(Pair.of(transFuture, partitionResponse));
+                    return queue;
+                }
+            });
+
+            topicManager.getTopic(topicName.toString());
+            publishMessages((MemoryRecords) entry.getValue(), topicName, partitionResponse, transFuture);
         }
-        produceRequestsQueue.add(Pair.of(produceHar, resultFuture));
 
-        handleProducerRequestInternal();
+        CompletableFuture.allOf(responsesFutures.values().toArray(new CompletableFuture<?>[responsesSize]))
+                .whenComplete((ignore, ex) -> {
+                    // all ex has translated to PartitionResponse with Errors.KAFKA_STORAGE_ERROR
+                    Map<TopicPartition, PartitionResponse> responses = new ConcurrentHashMap<>();
+                    for (Map.Entry<TopicPartition, CompletableFuture<PartitionResponse>> entry :
+                            responsesFutures.entrySet()) {
+                        responses.put(entry.getKey(), entry.getValue().join());
+                    }
+
+                    if (log.isDebugEnabled()) {
+                        log.debug("[{}] Request {}: Complete handle produce.",
+                                ctx.channel(), produceHar.toString());
+                    }
+                    resultFuture.complete(new ProduceResponse(responses));
+                });
     }
 
     protected void handleFindCoordinatorRequest(KafkaHeaderAndRequest findCoordinator,

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

@@ -22,6 +22,7 @@
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.broker.PulsarServerException;
@@ -218,8 +219,8 @@ private void lookupBroker(String topicName,
                     long waitTimeMs = backoff.next();
 
                     if (backoff.isMandatoryStopMade()) {
-                        log.warn("[{}] getBroker for topic failed, retried too many times, return null. throwable: ",
-                            requestHandler.ctx.channel(), topicName, waitTimeMs, th);
+                        log.warn("[{}] getBroker for topic {} failed, retried too many times {}, return null."
+                                        + " throwable: ", requestHandler.ctx.channel(), topicName, waitTimeMs, th);
                         retFuture.complete(null);
                     } else {
                         log.warn("[{}] getBroker for topic failed, will retry in {} ms. throwable: ",
@@ -260,58 +261,55 @@ public CompletableFuture<PersistentTopic> getTopic(String topicName) {
             rwLock.readLock().unlock();
         }
 
-        return topics.computeIfAbsent(topicName,
-            t -> {
-                getTopicBroker(t).whenCompleteAsync((ignore, th) -> {
-                    if (th != null || ignore == null) {
-                        log.warn("[{}] Failed getTopicBroker {}, return null PersistentTopic. throwable: ",
+        return topics.computeIfAbsent(topicName, t -> {
+            getTopicBroker(t).whenCompleteAsync((ignore, th) -> {
+                if (th != null || ignore == null) {
+                    log.warn("[{}] Failed getTopicBroker {}, return null PersistentTopic. throwable: ",
                             requestHandler.ctx.channel(), t, th);
 
-                        // get topic broker returns null. topic should be removed from LookupCache.
-                        if (ignore == null) {
-                            removeLookupCache(topicName);
-                        }
+                    // get topic broker returns null. topic should be removed from LookupCache.
+                    if (ignore == null) {
+                        removeLookupCache(topicName);
+                    }
+
+                    topicCompletableFuture.complete(null);
+                    return;
+                }
+
+                if (log.isDebugEnabled()) {
+                    log.debug("[{}] getTopicBroker for {} in KafkaTopicManager. brokerAddress: {}",
+                            requestHandler.ctx.channel(), t, ignore);
+                }
 
+                brokerService.getTopic(t, true).whenComplete((t2, throwable) -> {
+                    if (throwable != null) {
+                        log.error("[{}] Failed to getTopic {}. exception:",
+                                requestHandler.ctx.channel(), t, throwable);
+                        // failed to getTopic from current broker, remove cache, which added in getTopicBroker.
+                        removeLookupCache(t);
                         topicCompletableFuture.complete(null);
                         return;
                     }
 
-                    if (log.isDebugEnabled()) {
-                        log.debug("[{}] getTopicBroker for {} in KafkaTopicManager. brokerAddress: {}",
-                            requestHandler.ctx.channel(), t, ignore);
+                    try {
+                        if (t2.isPresent()) {
+                            PersistentTopic persistentTopic = (PersistentTopic) t2.get();
+                            references.putIfAbsent(t, registerInPersistentTopic(persistentTopic));
+                            topicCompletableFuture.complete(persistentTopic);
+                        } else {
+                            log.error("[{}]Get empty topic for name {}",
+                                    requestHandler.ctx.channel(), t);
+                            topicCompletableFuture.complete(null);
+                        }
+                    } catch (Exception e) {
+                        log.error("[{}] Failed to get client in registerInPersistentTopic {}. exception:",
+                                requestHandler.ctx.channel(), t, e);
+                        topicCompletableFuture.complete(null);
                     }
-
-                    brokerService
-                        .getTopic(t, true)
-                        .whenComplete((t2, throwable) -> {
-                            if (throwable != null) {
-                                log.error("[{}] Failed to getTopic {}. exception:",
-                                    requestHandler.ctx.channel(), t, throwable);
-                                // failed to getTopic from current broker, remove cache, which added in getTopicBroker.
-                                removeLookupCache(t);
-                                topicCompletableFuture.complete(null);
-                                return;
-                            }
-
-                            try {
-                                if (t2.isPresent()) {
-                                    PersistentTopic persistentTopic = (PersistentTopic) t2.get();
-                                    references.putIfAbsent(t, registerInPersistentTopic(persistentTopic));
-                                    topicCompletableFuture.complete(persistentTopic);
-                                } else {
-                                    log.error("[{}]Get empty topic for name {}",
-                                        requestHandler.ctx.channel(), t);
-                                    topicCompletableFuture.complete(null);
-                                }
-                            } catch (Exception e) {
-                                log.error("[{}] Failed to get client in registerInPersistentTopic {}. exception:",
-                                    requestHandler.ctx.channel(), t, e);
-                                topicCompletableFuture.complete(null);
-                            }
-                        });
                 });
-                return topicCompletableFuture;
             });
+            return topicCompletableFuture;
+        });
     }
 
     // when channel close, release all the topics reference in persistentTopic