add kafka consumer distributed traces

Author: obenkenobi

Diff for: instrumentation/kafka-clients-spans-consumer-0.11.0.0/build.gradle

@@ -0,0 +1,19 @@
+
+dependencies {
+    implementation(project(":agent-bridge"))
+    implementation("org.apache.kafka:kafka-clients:0.11.0.0")
+}
+
+jar {
+    manifest { attributes 'Implementation-Title': 'com.newrelic.instrumentation.kafka-clients-spans-consumer-0.11.0.0', 'Enabled': 'false',
+            'Implementation-Title-Alias': 'kafka-clients-spans-consumer' }
+}
+
+verifyInstrumentation {
+    passes 'org.apache.kafka:kafka-clients:[0.11.0.0,)'
+}
+
+site {
+    title 'Kafka'
+    type 'Messaging'
+}

Diff for: instrumentation/kafka-clients-spans-consumer-0.11.0.0/src/main/java/com/nr/instrumentation/kafka/HeadersWrapper.java

@@ -0,0 +1,81 @@
+package com.nr.instrumentation.kafka;
+
+import com.newrelic.api.agent.HeaderType;
+import com.newrelic.api.agent.Headers;
+import org.apache.kafka.common.header.Header;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Objects;
+
+public class HeadersWrapper implements Headers {
+
+    private final org.apache.kafka.common.header.Headers delegate;
+
+    public HeadersWrapper(org.apache.kafka.common.header.Headers headers) {
+        this.delegate = headers;
+    }
+
+    @Override
+    public HeaderType getHeaderType() {
+        return HeaderType.MESSAGE;
+    }
+
+    @Override
+    public String getHeader(String name) {
+        String value = null;
+        Iterator<Header> iterator = delegate.headers(name).iterator();
+        if (iterator.hasNext()) {
+            byte[] bytes = iterator.next().value();
+            if (bytes != null) {
+                value = new String(bytes);
+            }
+        }
+        return value;
+    }
+
+    @Override
+    public Collection<String> getHeaders(String name) {
+        Collection<String> headers = new ArrayList<>();
+        Iterator<Header> iterator = delegate.headers(name).iterator();
+        while (iterator.hasNext()) {
+            byte[] bytes = iterator.next().value();
+            if (bytes != null) {
+                headers.add(new String(bytes));
+            }
+        }
+        return headers;
+    }
+
+    @Override
+    public void setHeader(String name, String value) {
+        delegate.remove(name);
+        delegate.add(name, value.getBytes());
+    }
+
+    @Override
+    public void addHeader(String name, String value) {
+        delegate.add(name, value.getBytes());
+    }
+
+    @Override
+    public Collection<String> getHeaderNames() {
+        Collection<String> headerNames = new HashSet<>();
+        for(Header header : delegate) {
+            headerNames.add(header.key());
+        }
+        return headerNames;
+    }
+
+    @Override
+    public boolean containsHeader(String name) {
+        for(Header header : delegate) {
+            if (Objects.equals(name,header.key())) {
+                return true;
+            }
+        }
+        return false;
+    }
+}

Diff for: instrumentation/kafka-clients-spans-consumer-0.11.0.0/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer_Instrumentation.java

@@ -0,0 +1,41 @@
+package org.apache.kafka.clients.consumer;
+
+import com.newrelic.agent.bridge.AgentBridge;
+import com.newrelic.api.agent.Headers;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.TransportType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.nr.instrumentation.kafka.HeadersWrapper;
+
+import java.time.Duration;
+import java.util.logging.Level;
+
+@Weave(originalName = "org.apache.kafka.clients.consumer.KafkaConsumer")
+public class KafkaConsumer_Instrumentation<K, V> {
+
+    public ConsumerRecords<K, V> poll(final Duration timeout) {
+        final ConsumerRecords<K, V> records = Weaver.callOriginal();
+        nrAcceptDtHeaders(records);
+        return records;
+    }
+
+    public ConsumerRecords<K, V> poll(final long timeoutMs) {
+        final ConsumerRecords<K, V> records = Weaver.callOriginal();
+        nrAcceptDtHeaders(records);
+        return records;
+    }
+
+    private void nrAcceptDtHeaders(ConsumerRecords<K, V> records) {
+        AgentBridge.getAgent().getLogger().log(Level.INFO, "nrAcceptDtHeaders");
+        if (AgentBridge.getAgent().getTransaction(false) != null) {
+            AgentBridge.getAgent().getLogger().log(Level.INFO, "nrAcceptDtHeaders getTransaction exists");
+            for (ConsumerRecord<?, ?> record : records) {
+                AgentBridge.getAgent().getLogger().log(Level.INFO, "nrAcceptDtHeaders acceptDistributedTraceHeaders");
+                Headers dtHeaders = new HeadersWrapper(record.headers());
+                NewRelic.getAgent().getTransaction().acceptDistributedTraceHeaders(TransportType.Kafka, dtHeaders);
+                break;
+            }
+        }
+    }
+}

Diff for: settings.gradle

@@ -255,6 +255,7 @@ include 'instrumentation:kafka-clients-node-metrics-2.3.0'
 include 'instrumentation:kafka-clients-node-metrics-3.7.0'
 include 'instrumentation:kafka-clients-node-metrics-3.9.0'
 include 'instrumentation:kafka-clients-spans-0.11.0.0'
+include 'instrumentation:kafka-clients-spans-consumer-0.11.0.0'
 include 'instrumentation:kafka-connect-metrics-1.0.0'
 include 'instrumentation:kafka-connect-spans-2.0.0'
 include 'instrumentation:kafka-connect-spans-3.3.0'