fix: Add ability to add headers

Author: shubhendumadhukar

pom.xml

@@ -73,7 +73,7 @@
         <dependency>
             <groupId>org.apache.kafka</groupId>
             <artifactId>kafka-clients</artifactId>
-            <version>0.9.0.1</version>
+            <version>3.0.0</version>
             <exclusions>
                 <!-- XXX Transitively includes log4j 1.2.15 which has bad metadata
                      http://stackoverflow.com/a/9047963 -->

src/main/java/co/signal/kafkameter/KafkaProducerSampler.java

@@ -18,6 +18,8 @@
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.Properties;
 
 import org.apache.kafka.clients.producer.KafkaProducer;
@@ -33,6 +35,7 @@
 import org.apache.jmeter.samplers.SampleResult;
 import org.apache.jorphan.logging.LoggingManager;
 import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.header.internals.RecordHeader;
 import org.apache.log.Logger;
 
 import co.signal.handlebars.CustomHandlebars;
@@ -121,6 +124,11 @@ public class KafkaProducerSampler extends AbstractJavaSamplerClient {
    */
   private static final String PARAMETER_KAFKA_PARTITION = "kafka_partition";
 
+  /**
+   * Parameter for setting the headers. It is optional.
+   */
+  private static final String PARAMETER_KAFKA_HEADERS = "kafka_headers";
+
   // private Producer<Long, byte[]> producer;
 
   private KafkaProducer<String, String> producer;
@@ -178,6 +186,7 @@ public Arguments getDefaultParameters() {
     defaultParameters.addArgument(PARAMETER_KAFKA_USE_SSL, "${PARAMETER_KAFKA_USE_SSL}");
     defaultParameters.addArgument(PARAMETER_KAFKA_COMPRESSION_TYPE, null);
     defaultParameters.addArgument(PARAMETER_KAFKA_PARTITION, null);
+    defaultParameters.addArgument(PARAMETER_KAFKA_HEADERS, null);
     return defaultParameters;
   }
 
@@ -197,6 +206,19 @@ public SampleResult runTest(JavaSamplerContext context) {
     } catch (IOException e1) {
       e1.printStackTrace();
     }
+    Map<String, String> headers = new LinkedHashMap<String, String>();
+    try {
+      String headerString = handlebars.compileInline(context.getParameter(PARAMETER_KAFKA_HEADERS)).apply("");
+      if (!Strings.isNullOrEmpty(headerString)) {
+        String[] pairs = headerString.split("&");
+        for (String pair : pairs) {
+          int idx = pair.indexOf("=");
+          headers.put(pair.substring(0, idx), pair.substring(idx + 1));
+        }
+      }
+    } catch (IOException e1) {
+      e1.printStackTrace();
+    }
     sampleResultStart(result, message);
 
     final ProducerRecord<String, String> producerRecord;
@@ -207,6 +229,9 @@ public SampleResult runTest(JavaSamplerContext context) {
       final int partitionNumber = Integer.parseInt(partitionString);
       producerRecord = new ProducerRecord<String, String>(topic, partitionNumber, key, message);
     }
+    for (Map.Entry<String, String> entry : headers.entrySet()) {
+      producerRecord.headers().add(new RecordHeader(entry.getKey(), entry.getValue().getBytes()));
+    }
 
     try {
       producer.send(producerRecord);
@@ -301,4 +326,4 @@ private String getStackTrace(Exception exception) {
     exception.printStackTrace(new PrintWriter(stringWriter));
     return stringWriter.toString();
   }
-}
+}