Prevent duplicate DT headers for grpc

instrumentation/grpc-1.22.0/src/main/java/com/nr/agent/instrumentation/grpc/OutboundHeadersWrapper.java

@@ -27,7 +27,9 @@ public HeaderType getHeaderType() {
     @Override
     public void setHeader(String key, String value) {
         if (GrpcConfig.disributedTracingEnabled) {
-            metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            if (!metadata.containsKey(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER))) {
+                metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            }
         }
     }
 }

instrumentation/grpc-1.22.0/src/test/java/GrpcDistributedTracingTest.java

@@ -0,0 +1,68 @@
+import app.TestClient;
+import app.TestServer;
+import com.newrelic.agent.introspec.InstrumentationTestConfig;
+import com.newrelic.agent.introspec.InstrumentationTestRunner;
+import com.newrelic.agent.introspec.Introspector;
+import com.newrelic.agent.introspec.TransactionEvent;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Collection;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = { "io.grpc", "com.nr.agent.instrumentation.grpc" }, configName="enable_dist_tracing.yml")
+public class GrpcDistributedTracingTest {
+
+    private static TestServer server;
+    private static TestClient client;
+
+    @BeforeClass
+    public static void before() throws Exception {
+        server = new TestServer();
+        server.start();
+        client = new TestClient("localhost", server.getPort());
+    }
+
+    @AfterClass
+    public static void after() throws InterruptedException {
+        if (client != null) {
+            client.shutdown();
+        }
+        if (server != null) {
+            server.stop();
+        }
+    }
+
+    @Test
+    public void testDTAsyncRequest() {
+        client.helloAsync("Async");
+
+        String clientTxName = "OtherTransaction/Custom/app.TestClient/helloAsync";
+        String serverTxName = "WebTransaction/gRPC/helloworld.Greeter/SayHello";
+
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+
+        assertEquals(2, introspector.getFinishedTransactionCount(30000));
+
+        Collection<TransactionEvent> clientTxns = introspector.getTransactionEvents(clientTxName);
+        assertEquals(1, clientTxns.size());
+        String clientGuid = "";
+        String clientTripId = "";
+        for (TransactionEvent txn: clientTxns) {
+            //make other assertions?
+            clientGuid = txn.getMyGuid();
+            clientTripId = txn.getTripId();
+        }
+
+        Collection<TransactionEvent> serverTxns = introspector.getTransactionEvents(serverTxName);
+        assertEquals(1, serverTxns.size());
+        for (TransactionEvent txn: serverTxns) {
+            assertEquals(clientGuid, txn.getParentId());
+            assertEquals(clientTripId, txn.getTripId());
+        }
+    }
+}

instrumentation/grpc-1.22.0/src/test/resources/enable_dist_tracing.yml

@@ -0,0 +1,3 @@
+common: &default_settings
+  distributed_tracing:
+    enabled: true

instrumentation/grpc-1.30.0/src/main/java/com/nr/agent/instrumentation/grpc/OutboundHeadersWrapper.java

@@ -27,7 +27,9 @@ public HeaderType getHeaderType() {
     @Override
     public void setHeader(String key, String value) {
         if (GrpcConfig.disributedTracingEnabled) {
-            metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            if (!metadata.containsKey(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER))) {
+                metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            }
         }
     }
 }

instrumentation/grpc-1.30.0/src/test/java/GrpcDistributedTracingTest.java

@@ -0,0 +1,67 @@
+import app.TestClient;
+import app.TestServer;
+import com.newrelic.agent.introspec.InstrumentationTestConfig;
+import com.newrelic.agent.introspec.InstrumentationTestRunner;
+import com.newrelic.agent.introspec.Introspector;
+import com.newrelic.agent.introspec.TransactionEvent;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Collection;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = { "io.grpc", "com.nr.agent.instrumentation.grpc" }, configName="enable_dist_tracing.yml")
+public class GrpcDistributedTracingTest {
+
+    private static TestServer server;
+    private static TestClient client;
+
+    @BeforeClass
+    public static void before() throws Exception {
+        server = new TestServer();
+        server.start();
+        client = new TestClient("localhost", server.getPort());
+    }
+
+    @AfterClass
+    public static void after() throws InterruptedException {
+        if (client != null) {
+            client.shutdown();
+        }
+        if (server != null) {
+            server.stop();
+        }
+    }
+
+    @Test
+    public void testDTAsyncRequest() {
+        client.helloAsync("Async");
+
+        String clientTxName = "OtherTransaction/Custom/app.TestClient/helloAsync";
+        String serverTxName = "WebTransaction/gRPC/helloworld.Greeter/SayHello";
+
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+
+        assertEquals(2, introspector.getFinishedTransactionCount(30000));
+
+        Collection<TransactionEvent> clientTxns = introspector.getTransactionEvents(clientTxName);
+        assertEquals(1, clientTxns.size());
+        String clientGuid = "";
+        String clientTripId = "";
+        for (TransactionEvent txn: clientTxns) {
+            clientGuid = txn.getMyGuid();
+            clientTripId = txn.getTripId();
+        }
+
+        Collection<TransactionEvent> serverTxns = introspector.getTransactionEvents(serverTxName);
+        assertEquals(1, serverTxns.size());
+        for (TransactionEvent txn: serverTxns) {
+            assertEquals(clientGuid, txn.getParentId());
+            assertEquals(clientTripId, txn.getTripId());
+        }
+    }
+}

instrumentation/grpc-1.30.0/src/test/resources/enable_dist_tracing.yml

@@ -0,0 +1,3 @@
+common: &default_settings
+  distributed_tracing:
+    enabled: true

instrumentation/grpc-1.4.0/src/main/java/com/nr/agent/instrumentation/grpc/OutboundHeadersWrapper.java

@@ -27,7 +27,9 @@ public HeaderType getHeaderType() {
     @Override
     public void setHeader(String key, String value) {
         if (GrpcConfig.disributedTracingEnabled) {
-            metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            if (!metadata.containsKey(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER))) {
+                metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            }
         }
     }
 }

instrumentation/grpc-1.4.0/src/test/java/GrpcDistributedTracingTest.java

@@ -0,0 +1,68 @@
+import app.TestClient;
+import app.TestServer;
+import com.newrelic.agent.introspec.InstrumentationTestConfig;
+import com.newrelic.agent.introspec.InstrumentationTestRunner;
+import com.newrelic.agent.introspec.Introspector;
+import com.newrelic.agent.introspec.TransactionEvent;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Collection;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = { "io.grpc", "com.nr.agent.instrumentation.grpc" }, configName="enable_dist_tracing.yml")
+public class GrpcDistributedTracingTest {
+
+    private static TestServer server;
+    private static TestClient client;
+
+    @BeforeClass
+    public static void before() throws Exception {
+        server = new TestServer();
+        server.start();
+        client = new TestClient("localhost", server.getPort());
+    }
+
+    @AfterClass
+    public static void after() throws InterruptedException {
+        if (client != null) {
+            client.shutdown();
+        }
+        if (server != null) {
+            server.stop();
+        }
+    }
+
+    @Test
+    public void testDTAsyncRequest() {
+        client.helloAsync("Async");
+
+        String clientTxName = "OtherTransaction/Custom/app.TestClient/helloAsync";
+        String serverTxName = "WebTransaction/gRPC/helloworld.Greeter/SayHello";
+
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+
+        assertEquals(2, introspector.getFinishedTransactionCount(30000));
+
+        Collection<TransactionEvent> clientTxns = introspector.getTransactionEvents(clientTxName);
+        assertEquals(1, clientTxns.size());
+        String clientGuid = "";
+        String clientTripId = "";
+        for (TransactionEvent txn: clientTxns) {
+            //make other assertions?
+            clientGuid = txn.getMyGuid();
+            clientTripId = txn.getTripId();
+        }
+
+        Collection<TransactionEvent> serverTxns = introspector.getTransactionEvents(serverTxName);
+        assertEquals(1, serverTxns.size());
+        for (TransactionEvent txn: serverTxns) {
+            assertEquals(clientGuid, txn.getParentId());
+            assertEquals(clientTripId, txn.getTripId());
+        }
+    }
+}

instrumentation/grpc-1.4.0/src/test/resources/enable_dist_tracing.yml

@@ -0,0 +1,3 @@
+common: &default_settings
+  distributed_tracing:
+    enabled: true

instrumentation/grpc-1.40.0/src/main/java/com/nr/agent/instrumentation/grpc/OutboundHeadersWrapper.java

@@ -27,7 +27,9 @@ public HeaderType getHeaderType() {
     @Override
     public void setHeader(String key, String value) {
         if (GrpcConfig.disributedTracingEnabled) {
-            metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            if (!metadata.containsKey(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER))) {
+                metadata.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
+            }
         }
     }
 }

instrumentation/grpc-1.40.0/src/test/java/GrpcDistributedTracingTest.java

@@ -0,0 +1,67 @@
+
+import app.TestClient;
+import app.TestServer;
+import com.newrelic.agent.introspec.*;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Ignore;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.Collection;
+
+import static org.junit.Assert.assertEquals;
+
+@RunWith(InstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = { "io.grpc", "com.nr.agent.instrumentation.grpc" }, configName="enable_dist_tracing.yml")
+public class GrpcDistributedTracingTest {
+
+    private static TestServer server;
+    private static TestClient client;
+
+    @BeforeClass
+    public static void before() throws Exception {
+        server = new TestServer();
+        server.start();
+        client = new TestClient("localhost", server.getPort());
+    }
+
+    @AfterClass
+    public static void after() throws InterruptedException {
+        if (client != null) {
+            client.shutdown();
+        }
+        if (server != null) {
+            server.stop();
+        }
+    }
+
+    @Test
+    public void testDTAsyncRequest() {
+        client.helloAsync("Async");
+
+        String clientTxName = "OtherTransaction/Custom/app.TestClient/helloAsync";
+        String serverTxName = "WebTransaction/gRPC/helloworld.Greeter/SayHello";
+
+        Introspector introspector = InstrumentationTestRunner.getIntrospector();
+
+        assertEquals(2, introspector.getFinishedTransactionCount(30000));
+
+        Collection<TransactionEvent> clientTxns = introspector.getTransactionEvents(clientTxName);
+        assertEquals(1, clientTxns.size());
+        String clientGuid = "";
+        String clientTripId = "";
+        for (TransactionEvent txn: clientTxns) {
+            //make other assertions?
+            clientGuid = txn.getMyGuid();
+            clientTripId = txn.getTripId();
+        }
+
+        Collection<TransactionEvent> serverTxns = introspector.getTransactionEvents(serverTxName);
+        assertEquals(1, serverTxns.size());
+        for (TransactionEvent txn: serverTxns) {
+            assertEquals(clientGuid, txn.getParentId());
+            assertEquals(clientTripId, txn.getTripId());
+        }
+    }
+}

instrumentation/grpc-1.40.0/src/test/resources/enable_dist_tracing.yml

@@ -0,0 +1,3 @@
+common: &default_settings
+  distributed_tracing:
+    enabled: true

newrelic-agent/src/main/java/com/newrelic/agent/tracing/W3CTraceParentParser.java

@@ -7,19 +7,30 @@
 
 package com.newrelic.agent.tracing;
 
+import com.newrelic.agent.Agent;
 import com.newrelic.agent.MetricNames;
 import com.newrelic.agent.service.ServiceFactory;
 
 import java.util.List;
+import java.util.logging.Level;
 
 public class W3CTraceParentParser {
 
     static W3CTraceParent parseHeaders(List<String> traceParentHeaders) {
-        if (traceParentHeaders.size() != 1) {
-            ServiceFactory.getStatsService().getMetricAggregator().incrementCounter(MetricNames.SUPPORTABILITY_TRACE_CONTEXT_INVALID_PARENT_HEADER_COUNT);
+        if (traceParentHeaders.isEmpty()) {
             return null;
         }
-
+        if (traceParentHeaders.size() > 1) {
+            ServiceFactory.getStatsService().getMetricAggregator().incrementCounter(MetricNames.SUPPORTABILITY_TRACE_CONTEXT_INVALID_PARENT_HEADER_COUNT);
+            Agent.LOG.log(Level.WARNING, "Multiple traceparent headers found on inbound request.");
+            // Multiple values ok if all are equal
+            String first = traceParentHeaders.get(0);
+            for (String header : traceParentHeaders) {
+                if (!header.equals(first)) {
+                   return null;
+                }
+            }
+        }
         String traceParentHeader = traceParentHeaders.get(0);
         return parseHeader(traceParentHeader);
     }

newrelic-agent/src/test/java/com/newrelic/agent/tracing/W3CTraceParentHeaderTest.java

@@ -69,9 +69,23 @@ public void testCreateTraceParentPayloadPadsAndLowerCasesTraceId() {
     }
 
     @Test
-    public void testParseMultiple_oopsTooMany() {
+    public void testParseMultiple_differentValues_shouldFail() {
+        W3CTraceParent result = W3CTraceParentParser.parseHeaders(
+                Arrays.asList("00-12345678123456781234567812345678-1234123412341234-01", "00-87654321876543218765432187654321-1234123412341234-01"));
+        assertNull(result);
+    }
+
+    @Test
+    public void testParseMultiple_sameValues_shouldPass() {
         W3CTraceParent result = W3CTraceParentParser.parseHeaders(
                 Arrays.asList("00-12345678123456781234567812345678-1234123412341234-01", "00-12345678123456781234567812345678-1234123412341234-01"));
+        W3CTraceParent expected = new W3CTraceParent("00", "12345678123456781234567812345678", "1234123412341234", 1);
+        assertEquals(expected, result);
+    }
+
+    @Test
+    public void testParseMultiple_empty_shouldFail() {
+        W3CTraceParent result = W3CTraceParentParser.parseHeaders(Arrays.asList());
         assertNull(result);
     }