Added Support and Handling For Generic Headers (#931)

* Added Support and Handling For Generic Headers

* Added an exception path in case unsupported header type is encountered

* Fix a comment

---------

Co-authored-by: Shrinand Thakkar <sthakkar@sthakkar-mn2.linkedin.biz>

Author: shrinandthakkar

build.gradle

@@ -154,7 +154,6 @@ project(':datastream-common') {
     compile "com.linkedin.pegasus:restli-server:$pegasusVersion"
     compile "com.intellij:annotations:$intellijAnnotationsVersion"
     compile "com.google.guava:guava:$guavaVersion"
-    compile "com.linkedin.kafka:kafka-clients:$kafkaVersion"
     compile "com.fasterxml.jackson.core:jackson-annotations:$jacksonVersion"
     compile "com.fasterxml.jackson.core:jackson-core:$jacksonVersion"
     compile "com.fasterxml.jackson.core:jackson-databind:$jacksonVersion"
@@ -172,6 +171,7 @@ project(':datastream-server-api') {
 project(':datastream-utils') {
   dependencies {
     compile project(':datastream-common')
+    compile "com.linkedin.kafka:kafka-clients:$kafkaVersion"
     compile "org.apache.helix:zookeeper-api:$helixZkclientVersion"
     compile "com.linkedin.zookeeper:zookeeper:$zookeeperVersion"
     compile "com.google.guava:guava:$guavaVersion"

datastream-common/src/main/java/com/linkedin/datastream/common/BrooklinEnvelope.java

@@ -11,7 +11,6 @@
 
 import org.apache.avro.reflect.Nullable;
 import org.apache.commons.lang.Validate;
-import org.apache.kafka.common.header.Headers;
 
 
 /**
@@ -30,7 +29,7 @@ public class BrooklinEnvelope {
 
   private Map<String, String> _metadata;
 
-  private Headers _headers;
+  private Object _headers;
 
   /**
    * Construct a BrooklinEnvelope using record key, value, and metadata
@@ -59,11 +58,11 @@ public BrooklinEnvelope(@Nullable Object key, @Nullable Object value, @Nullable
    * @param key The record key (e.g. primary key)
    * @param value The new record value
    * @param previousValue The old record value
-   * @param headers Kafka headers to associate with the change event
+   * @param headers Generic headers to associate with the event
    * @param metadata Additional metadata to associate with the change event
    */
   public BrooklinEnvelope(@Nullable Object key, @Nullable Object value, @Nullable Object previousValue,
-      @Nullable Headers headers, Map<String, String> metadata) {
+      @Nullable Object headers, Map<String, String> metadata) {
     Validate.notNull(metadata, "metadata cannot be null");
     setKey(key);
     setValue(value);
@@ -93,16 +92,16 @@ public void setPreviousValue(Object previousValue) {
   }
 
   /**
-   * Get the Kafka headers
+   * Get the Generic headers
    */
-  public Headers getHeaders() {
+  public Object getHeaders() {
     return _headers;
   }
 
   /**
-   * Set the Kafka headers
+   * Set the Generic headers
    */
-  public void setHeaders(Headers headers) {
+  public void setHeaders(Object headers) {
     _headers = headers;
   }
 

datastream-kafka-connector/src/test/java/com/linkedin/datastream/connectors/kafka/mirrormaker/KafkaMirrorMakerConnectorTestUtils.java

@@ -16,6 +16,7 @@
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.testng.Assert;
 
@@ -52,19 +53,23 @@ static Properties getKafkaProducerProperties(DatastreamEmbeddedZookeeperKafkaClu
 
   static void produceEvents(String topic, int destinationPartition, int numEvents,
       DatastreamEmbeddedZookeeperKafkaCluster kafkaCluster) {
-    produceEventsToPartition(topic, destinationPartition, numEvents, kafkaCluster);
+    produceEventsToPartition(topic, destinationPartition, numEvents, kafkaCluster, null);
   }
 
   static void produceEvents(String topic, int numEvents, DatastreamEmbeddedZookeeperKafkaCluster kafkaCluster) {
-    produceEventsToPartition(topic, null, numEvents, kafkaCluster);
+    produceEventsToPartition(topic, null, numEvents, kafkaCluster, null);
+  }
+
+  static void produceEventsWithHeaders(String topic, int numEvents, DatastreamEmbeddedZookeeperKafkaCluster kafkaCluster, Headers headers) {
+    produceEventsToPartition(topic, null, numEvents, kafkaCluster, headers);
   }
 
   static void produceEventsToPartition(String topic, Integer destinationPartition, int numEvents,
-      DatastreamEmbeddedZookeeperKafkaCluster kafkaCluster) {
+      DatastreamEmbeddedZookeeperKafkaCluster kafkaCluster, Headers headers) {
     try (Producer<byte[], byte[]> producer = new KafkaProducer<>(getKafkaProducerProperties(kafkaCluster))) {
       for (int i = 0; i < numEvents; i++) {
         producer.send(new ProducerRecord<>(topic, destinationPartition, ("key-" + i).getBytes(Charsets.UTF_8),
-            ("value-" + i).getBytes(Charsets.UTF_8)), (metadata, exception) -> {
+            ("value-" + i).getBytes(Charsets.UTF_8), headers), (metadata, exception) -> {
           if (exception != null) {
             throw new RuntimeException("Failed to send message.", exception);
           }

datastream-kafka-connector/src/test/java/com/linkedin/datastream/connectors/kafka/mirrormaker/TestKafkaConsumerOffsets.java

@@ -61,7 +61,7 @@ public void testConsumerOffsetsDiagEndpoint() {
 
     // produce messages to each topic partition
     topics.forEach(topic -> IntStream.range(0, PARTITION_COUNT).forEach(partition ->
-        KafkaMirrorMakerConnectorTestUtils.produceEventsToPartition(topic, partition, PARTITION_MESSAGE_COUNT, _kafkaCluster)));
+        KafkaMirrorMakerConnectorTestUtils.produceEventsToPartition(topic, partition, PARTITION_MESSAGE_COUNT, _kafkaCluster, null)));
 
     connector.onAssignmentChange(Collections.singletonList(task));
 

datastream-kafka-connector/src/test/java/com/linkedin/datastream/connectors/kafka/mirrormaker/TestKafkaMirrorMakerConnectorTask.java

@@ -19,6 +19,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
+import org.apache.commons.io.Charsets;
 import org.apache.commons.lang.StringUtils;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -27,6 +28,8 @@
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.InterruptException;
 import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
@@ -153,6 +156,55 @@ public List<BrooklinMetricInfo> getMetricInfos() {
         "did not shut down on time");
   }
 
+  @Test
+  public void testConsumeFromSingleTopicWithHeaders() throws Exception {
+    String connectorName = "testConsumeFromSingleTopicWithHeaders";
+
+    String yummyTopic = "YummyPizza";
+
+    createTopic(_adminClient, yummyTopic);
+
+    // create a datastream to consume from topics ending in "Pizza"
+    Datastream datastream = KafkaMirrorMakerConnectorTestUtils.createDatastream("pizzaStream", _broker, "\\w+Pizza");
+
+    DatastreamTaskImpl task = new DatastreamTaskImpl(Collections.singletonList(datastream));
+    MockDatastreamEventProducer datastreamProducer = new MockDatastreamEventProducer();
+    task.setEventProducer(datastreamProducer);
+
+    KafkaMirrorMakerConnectorTask connectorTask =
+        KafkaMirrorMakerConnectorTestUtils.createKafkaMirrorMakerConnectorTask(task, connectorName);
+    KafkaMirrorMakerConnectorTestUtils.runKafkaMirrorMakerConnectorTask(connectorTask);
+
+    Headers genericHeaders = new RecordHeaders().add("headerKey", "headerValue".getBytes(Charsets.UTF_8));
+    // produce an event to the topic with generic headers
+    KafkaMirrorMakerConnectorTestUtils.produceEventsWithHeaders(yummyTopic, 10, _kafkaCluster, genericHeaders);
+
+    if (!PollUtils.poll(() -> datastreamProducer.getEvents().size() == 10, POLL_PERIOD_MS, POLL_TIMEOUT_MS)) {
+      Assert.fail("did not transfer the msgs within timeout. transferred " + datastreamProducer.getEvents().size());
+    }
+
+    List<DatastreamProducerRecord> records = datastreamProducer.getEvents();
+    for (DatastreamProducerRecord record : records) {
+      String destinationTopic = record.getDestination().get();
+      Assert.assertTrue(destinationTopic.endsWith("Pizza"),
+          "Unexpected event consumed from Datastream and sent to topic: " + destinationTopic);
+      Assert.assertEquals(genericHeaders, record.getEvents().get(0).getHeaders());
+    }
+
+    // Verify that metrics created through DynamicMetricsManager match those returned by getMetricInfos() given the
+    // connector name of interest.
+    MetricsTestUtils.verifyMetrics(new MetricsAware() {
+      @Override
+      public List<BrooklinMetricInfo> getMetricInfos() {
+        return KafkaMirrorMakerConnectorTask.getMetricInfos(connectorName);
+      }
+    }, DynamicMetricsManager.getInstance(), s -> s.startsWith(connectorName));
+
+    connectorTask.stop();
+    Assert.assertTrue(connectorTask.awaitStop(CONNECTOR_AWAIT_STOP_TIMEOUT_MS, TimeUnit.MILLISECONDS),
+        "did not shut down on time");
+  }
+
   @Test
   public void testConsumeFromMultipleTopicsWithDestinationTopicPrefixMetadata() throws Exception {
     String yummyTopic = "YummyPizza";
@@ -227,7 +279,7 @@ public void testIdentityPartitioningEnabled() throws Exception {
     // produce an event half of the partitions
     Set<Integer> expectedPartitionsWithData = new HashSet<>();
     for (int i = 0; i < partitionCount; i += 2) {
-      KafkaMirrorMakerConnectorTestUtils.produceEventsToPartition(yummyTopic, i, 1, _kafkaCluster);
+      KafkaMirrorMakerConnectorTestUtils.produceEventsToPartition(yummyTopic, i, 1, _kafkaCluster, null);
       expectedPartitionsWithData.add(i);
     }
 

datastream-kafka/src/main/java/com/linkedin/datastream/kafka/KafkaTransportProvider.java

@@ -112,7 +112,14 @@ private ProducerRecord<byte[], byte[]> convertToProducerRecord(String topicName,
     Headers headers = null;
     if (event instanceof BrooklinEnvelope) {
       BrooklinEnvelope envelope = (BrooklinEnvelope) event;
-      headers = envelope.getHeaders();
+      if (envelope.getHeaders() != null) {
+        if (!(envelope.getHeaders() instanceof Headers)) {
+          throw new DatastreamRuntimeException(
+              String.format("Unsupported header encountered %s in kafka transport provider for record %s",
+                  envelope.getHeaders(), record));
+        }
+        headers = (Headers) envelope.getHeaders();
+      }
       if (envelope.key().isPresent() && envelope.key().get() instanceof byte[]) {
         keyValue = (byte[]) envelope.key().get();
       }