[Issue #45] [pulsar-client-kafka-compat] Handled Kafka record headers…

pulsar-client-kafka-compat/pulsar-client-kafka/src/main/java/org/apache/kafka/clients/constants/MessageConstants.java

@@ -0,0 +1,28 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.kafka.clients.constants;
+
+public final class MessageConstants {
+    private MessageConstants() {
+
+    }
+
+    public static final String KAFKA_MESSAGE_HEADER_PREFIX = "kafka.header.";
+}

pulsar-client-kafka-compat/pulsar-client-kafka/src/main/java/org/apache/kafka/clients/consumer/PulsarKafkaConsumer.java

@@ -43,10 +43,15 @@
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.codec.DecoderException;
+import org.apache.commons.codec.binary.Hex;
+import org.apache.kafka.clients.constants.MessageConstants;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.common.serialization.Deserializer;
 import org.apache.kafka.common.serialization.StringDeserializer;
@@ -209,11 +214,11 @@ private PulsarKafkaConsumer(ConsumerConfig consumerConfig, Schema<K> keySchema,
 
     private SubscriptionInitialPosition getStrategy(final String strategy) {
         switch(strategy) {
-    	    case "earliest":
-    	        return SubscriptionInitialPosition.Earliest;
-    	    default:
+            case "earliest":
+                return SubscriptionInitialPosition.Earliest;
+            default:
                 return SubscriptionInitialPosition.Latest;
-    	}
+        }
     }
 
     @Override
@@ -386,8 +391,8 @@ public ConsumerRecords<K, V> poll(long timeoutMillis) {
 
                 TopicPartition tp = new TopicPartition(topic, partition);
                 if (lastReceivedOffset.get(tp) == null && !unpolledPartitions.contains(tp)) {
-                	log.info("When polling offsets, invalid offsets were detected. Resetting topic partition {}", tp);
-                	resetOffsets(tp);
+                    log.info("When polling offsets, invalid offsets were detected. Resetting topic partition {}", tp);
+                    resetOffsets(tp);
                 }
 
                 K key = getKey(topic, msg);
@@ -404,10 +409,33 @@ public ConsumerRecords<K, V> poll(long timeoutMillis) {
                     timestamp = msg.getEventTime();
                     timestampType = TimestampType.CREATE_TIME;
                 }
-
-                ConsumerRecord<K, V> consumerRecord = new ConsumerRecord<>(topic, partition, offset, timestamp,
-                        timestampType, -1, msg.hasKey() ? msg.getKey().length() : 0, msg.getData().length, key, value);
-
+                ConsumerRecord<K, V> consumerRecord;
+                if (msg.getProperties() != null) {
+                    Headers headers = new RecordHeaders();
+                    msg.getProperties().forEach((k, v) -> {
+                        // Kafka Specific Headers from the record
+                        if (k.startsWith(MessageConstants.KAFKA_MESSAGE_HEADER_PREFIX)) {
+                            String originalKey = k.replace(MessageConstants.KAFKA_MESSAGE_HEADER_PREFIX, "");
+                            try {
+                                headers.add(originalKey, Hex.decodeHex(v));
+                            } catch (DecoderException e) {
+                                log.warn("Corrupted Header Key : {}", originalKey);
+                                throw new RuntimeException(e);
+                            }
+                        } else {
+                            // Default headers injected by the PulsarKafkaProducerClient
+                            // For ex:KafkaMessageRouter.PARTITION_ID
+                            headers.add(k, v.getBytes());
+                        }
+                    });
+                    consumerRecord = new ConsumerRecord<>(topic, partition, offset, timestamp, timestampType, -1L,
+                            msg.hasKey() ? msg.getKey().length() : 0, msg.getData() != null ? msg.getData().length : 0,
+                            key, value, headers);
+                } else {
+                    consumerRecord = new ConsumerRecord<>(topic, partition, offset, timestamp, timestampType, -1L,
+                            msg.hasKey() ? msg.getKey().length() : 0, msg.getData() != null ? msg.getData().length : 0,
+                            key, value);
+                }
                 records.computeIfAbsent(tp, k -> new ArrayList<>()).add(consumerRecord);
 
                 // Update last offset seen by application
@@ -520,10 +548,12 @@ private CompletableFuture<Void> doCommitOffsets(Map<TopicPartition, OffsetAndMet
 
             lastCommittedOffset.put(tp, offsetAndMetadata);
             MessageId msgId = MessageIdUtils.getMessageId(offsetAndMetadata.offset());
-            if (consumer instanceof MultiTopicsConsumerImpl) {
-                msgId = new TopicMessageIdImpl(topicPartition.topic(), tp.topic(), msgId);
+            if (consumer != null) {
+                if (consumer instanceof MultiTopicsConsumerImpl) {
+                    msgId = new TopicMessageIdImpl(topicPartition.topic(), tp.topic(), msgId);
+                }
+                futures.add(consumer.acknowledgeCumulativeAsync(msgId));
             }
-            futures.add(consumer.acknowledgeCumulativeAsync(msgId));
         });
 
         return FutureUtil.waitForAll(futures);
@@ -663,7 +693,7 @@ public long position(TopicPartition topicPartition, Duration duration) {
     }
 
     private SubscriptionInitialPosition resetOffsets(final TopicPartition partition) {
-    	log.info("Resetting partition {} and seeking to {} position", partition, strategy);
+        log.info("Resetting partition {} and seeking to {} position", partition, strategy);
         if (strategy == SubscriptionInitialPosition.Earliest) {
             seekToBeginning(Collections.singleton(partition));
         } else {

pulsar-client-kafka-compat/pulsar-client-kafka/src/main/java/org/apache/kafka/clients/producer/PulsarKafkaProducer.java

@@ -18,6 +18,7 @@
  */
 package org.apache.kafka.clients.producer;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Base64;
@@ -36,7 +37,9 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
-
+import lombok.Getter;
+import org.apache.commons.codec.binary.Hex;
+import org.apache.kafka.clients.constants.MessageConstants;
 import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.Cluster;
@@ -64,10 +67,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import com.google.common.annotations.VisibleForTesting;
-
-import lombok.Getter;
-
 public class PulsarKafkaProducer<K, V> implements Producer<K, V> {
 
     private static final Logger log = LoggerFactory.getLogger(PulsarKafkaProducer.class);
@@ -387,6 +386,17 @@ private int buildMessage(TypedMessageBuilder<byte[]> builder, ProducerRecord<K,
             builder.property(KafkaMessageRouter.PARTITION_ID, Integer.toString(partition));
         }
 
+        // If record contains Headers,
+        // Prefix that with kafka.header.{keyName} to differentiate that these headers are from Kafka Record
+        // Encode the value with Hex
+        if (record.headers() != null) {
+            record.headers().forEach(header -> {
+                String key = MessageConstants.KAFKA_MESSAGE_HEADER_PREFIX + header.key();
+                builder.property(key, Hex.encodeHexString(header.value()));
+                log.debug("Formatted Kafka Specific Headers Before : {}, After : {}", header.key(), key);
+            });
+        }
+
         return value.length;
     }
 

pulsar-client-kafka-compat/pulsar-client-kafka/src/test/java/org/apache/kafka/clients/consumer/PulsarKafkaConsumerTest.java

@@ -0,0 +1,109 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.kafka.clients.consumer;
+
+import static org.mockito.ArgumentMatchers.anyCollection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+import java.util.concurrent.CompletableFuture;
+import org.apache.commons.codec.binary.Hex;
+import org.apache.kafka.clients.constants.MessageConstants;
+import org.apache.kafka.common.serialization.IntegerDeserializer;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.pulsar.client.api.ClientBuilder;
+import org.apache.pulsar.client.api.Consumer;
+import org.apache.pulsar.client.api.Message;
+import org.apache.pulsar.client.api.MessageId;
+import org.apache.pulsar.client.api.PulsarClient;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.impl.MessageImpl;
+import org.apache.pulsar.client.impl.PulsarClientImpl;
+import org.apache.pulsar.client.kafka.compat.KafkaMessageRouter;
+import org.apache.pulsar.common.api.proto.MessageMetadata;
+import org.apache.pulsar.common.naming.TopicName;
+import org.mockito.Mockito;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class PulsarKafkaConsumerTest {
+
+    @Test
+    public void testPulsarKafkaConsumerWithHeaders_noAck() throws Exception {
+        Consumer consumer = Mockito.mock(Consumer.class);
+        String topic = "topic";
+
+        Mockito.when(Mockito.mock(TopicName.class).getPartitionedTopicName()).thenReturn(topic);
+        Mockito.doReturn("topic").when(consumer).getTopic();
+
+        MessageMetadata messageMetadata = new MessageMetadata();
+        messageMetadata.setPublishTime(System.currentTimeMillis());
+
+        Map<String, String> headerMap = new HashMap<>();
+        String kafkaHeaderKey = MessageConstants.KAFKA_MESSAGE_HEADER_PREFIX + "header1";
+        String kafkaHeaderValue = Hex.encodeHexString(kafkaHeaderKey.getBytes());
+        headerMap.put(kafkaHeaderKey, kafkaHeaderValue);
+        headerMap.put(KafkaMessageRouter.PARTITION_ID, "0");
+        Message<byte[]> msg = new MessageImpl<>(
+                topic,
+                "1:1",
+                headerMap,
+                "string".getBytes(),
+                Schema.BYTES,
+                messageMetadata
+        );
+
+        PulsarClient mockClient = Mockito.mock(PulsarClient.class);
+        PulsarClientImpl mockClientImpl = Mockito.mock(PulsarClientImpl.class);
+
+        CompletableFuture<Integer> mockNoOfPartitionFuture = CompletableFuture.completedFuture(1);
+
+        ClientBuilder mockClientBuilder = Mockito.mock(ClientBuilder.class);
+        Mockito.doReturn(mockClientBuilder).when(mockClientBuilder).serviceUrl(Mockito.anyString());
+        Mockito.doReturn(mockClient).when(mockClientBuilder).build();
+
+        Mockito.when(mockClientImpl.getNumberOfPartitions(Mockito.anyString())).thenReturn(mockNoOfPartitionFuture);
+
+        Properties properties = new Properties();
+
+        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, IntegerDeserializer.class.getName());
+        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
+        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, Collections.singletonList("pulsar://localhost:6650"));
+        properties.put(ConsumerConfig.GROUP_ID_CONFIG, "my-subscription-name");
+
+        PulsarKafkaConsumer<Integer, String> pulsarKafkaConsumer =
+                new PulsarKafkaConsumer<>(properties, new IntegerDeserializer(), new StringDeserializer());
+
+        PulsarKafkaConsumer<Integer, String> pulsarKafkaConsumerSpy = Mockito.spy(pulsarKafkaConsumer);
+
+        Mockito.doNothing().when(pulsarKafkaConsumerSpy).seekToEnd(anyCollection());
+
+        pulsarKafkaConsumerSpy.received(consumer, msg);
+        pulsarKafkaConsumerSpy.poll(100);
+        pulsarKafkaConsumerSpy.close();
+
+        Assert.assertEquals(kafkaHeaderValue, msg.getProperty(kafkaHeaderKey));
+        Mockito.verify(pulsarKafkaConsumerSpy).seekToEnd(anyCollection());
+        Mockito.verify(consumer, Mockito.times(0)).acknowledgeCumulativeAsync(Mockito.any(MessageId.class));
+    }
+
+
+}
+