Merge pull request #2279 from newrelic/w3c_sampled

jbedell-newrelic · web-flow · commit 8d3800997b32 · 2025-03-20T16:59:56.000-04:00
Added sampling options when an inbound traceparent exists
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/HeadersUtil.java b/newrelic-agent/src/main/java/com/newrelic/agent/HeadersUtil.java
@@ -261,7 +261,7 @@ public static void parseAndAcceptDistributedTraceHeaders(Transaction tx, Inbound
             W3CTracePayload w3CTracePayload = W3CTracePayload.parseHeaders(tx, traceParent, traceState);
             if (w3CTracePayload != null) {
                 if (w3CTracePayload.getPayload() != null) {
-                    tx.acceptDistributedTracePayload(w3CTracePayload.getPayload());
+                    tx.acceptDistributedTracePayload(w3CTracePayload.getPayload(), w3CTracePayload.getTraceParent());
                 }
                 if (w3CTracePayload.getTraceParent() != null) {
                     tx.getSpanProxy().setInitiatingW3CTraceParent(w3CTracePayload.getTraceParent());
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/Transaction.java b/newrelic-agent/src/main/java/com/newrelic/agent/Transaction.java
@@ -60,6 +60,7 @@
 import com.newrelic.agent.tracing.DistributedTraceService;
 import com.newrelic.agent.tracing.DistributedTraceServiceImpl;
 import com.newrelic.agent.tracing.SpanProxy;
+import com.newrelic.agent.tracing.W3CTraceParent;
 import com.newrelic.agent.transaction.PriorityTransactionName;
 import com.newrelic.agent.transaction.TransactionCache;
 import com.newrelic.agent.transaction.TransactionCounts;
@@ -314,16 +315,37 @@ public boolean acceptDistributedTracePayload(String payload) {
         }
     }
 
-    public boolean acceptDistributedTracePayload(DistributedTracePayload payload) {
-        if (getAgentConfig().getDistributedTracingConfig().isEnabled()) {
+    public boolean acceptDistributedTracePayload(DistributedTracePayload payload, W3CTraceParent parent) {
+        DistributedTracingConfig dtConfig = getAgentConfig().getDistributedTracingConfig();
+        if (dtConfig.isEnabled()) {
             long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(
                     System.nanoTime() - this.getTransactionTimer().getStartTimeInNanos());
             long txnStartTimeSinceEpochInMillis = System.currentTimeMillis() - elapsedMillis;
             spanProxy.get().setTimestamp(txnStartTimeSinceEpochInMillis);
             boolean accepted = spanProxy.get().acceptDistributedTracePayload(payload);
             if (accepted) {
                 this.transportDurationInMillis = spanProxy.get().getTransportDurationInMillis();
-                this.setPriorityIfNotNull(spanProxy.get().getInboundDistributedTracePayload().priority);
+                if (parent != null) {
+                    if (parent.sampled()) { // traceparent exists and sampled is 1
+                        if (DistributedTracingConfig.SAMPLE_ALWAYS_ON.equals(dtConfig.getRemoteParentSampled())) {
+                            this.setPriorityIfNotNull(2.0f);
+                        } else if (DistributedTracingConfig.SAMPLE_ALWAYS_OFF.equals(dtConfig.getRemoteParentSampled())) {
+                            this.setPriorityIfNotNull(0.0f);
+                        } else {
+                            this.setPriorityIfNotNull(spanProxy.get().getInboundDistributedTracePayload().priority);
+                        }
+                    } else { // traceparent exists and sampled is 0
+                        if (DistributedTracingConfig.SAMPLE_ALWAYS_ON.equals(dtConfig.getRemoteParentNotSampled())) {
+                            this.setPriorityIfNotNull(2.0f);
+                        } else if (DistributedTracingConfig.SAMPLE_ALWAYS_OFF.equals(dtConfig.getRemoteParentNotSampled())) {
+                            this.setPriorityIfNotNull(0.0f);
+                        } else {
+                            this.setPriorityIfNotNull(spanProxy.get().getInboundDistributedTracePayload().priority);
+                        }
+                    }
+                } else {
+                    this.setPriorityIfNotNull(spanProxy.get().getInboundDistributedTracePayload().priority);
+                }
             }
             return accepted;
         } else {
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/TransactionApiImpl.java b/newrelic-agent/src/main/java/com/newrelic/agent/TransactionApiImpl.java
@@ -509,7 +509,7 @@ public void acceptDistributedTracePayload(String payload) {
     public void acceptDistributedTracePayload(DistributedTracePayload payload) {
         Transaction tx = getTransactionIfExists();
         if (tx != null) {
-            tx.acceptDistributedTracePayload(payload);
+            tx.acceptDistributedTracePayload(payload, null);
         }
     }
 
diff --git a/newrelic-agent/src/main/java/com/newrelic/agent/config/DistributedTracingConfig.java b/newrelic-agent/src/main/java/com/newrelic/agent/config/DistributedTracingConfig.java
@@ -21,12 +21,20 @@ public class DistributedTracingConfig extends BaseConfig {
     public static final String DISTRIBUTED_TRACING_ENABLED = SYSTEM_PROPERTY_ROOT + ENABLED;
     public static final String ENABLED_ENV_KEY = "NEW_RELIC_DISTRIBUTED_TRACING_ENABLED";
     public static final String EXCLUDE_NEWRELIC_HEADER = "exclude_newrelic_header";
+    public static final String SAMPLER = "sampler";
+    public static final String REMOTE_PARENT_SAMPLED = "remote_parent_sampled";
+    public static final String REMOTE_PARENT_NOT_SAMPLED = "remote_parent_not_sampled";
+    public static final String SAMPLE_ALWAYS_ON = "always_on";
+    public static final String SAMPLE_ALWAYS_OFF = "always_off";
+    public static final String SAMPLE_DEFAULT = "default";
 
     private final boolean enabled;
     private final String trustedAccountKey;
     private final String accountId;
     private final String primaryApplicationId;
     private final boolean includeNewRelicHeader;
+    private final String remoteParentSampled;
+    private final String remoteParentNotSampled;
 
     DistributedTracingConfig(Map<String, Object> props) {
         super(props, SYSTEM_PROPERTY_ROOT);
@@ -35,6 +43,10 @@ public class DistributedTracingConfig extends BaseConfig {
         this.accountId = getProperty(ACCOUNT_ID);
         this.primaryApplicationId = getProperty(PRIMARY_APPLICATION_ID);
         this.includeNewRelicHeader = !getProperty(EXCLUDE_NEWRELIC_HEADER, false);
+
+        BaseConfig samplerConfig = new BaseConfig(nestedProps(SAMPLER), SYSTEM_PROPERTY_ROOT+"."+SAMPLER);
+        this.remoteParentSampled = samplerConfig.getProperty(REMOTE_PARENT_SAMPLED, SAMPLE_DEFAULT);
+        this.remoteParentNotSampled = samplerConfig.getProperty(REMOTE_PARENT_NOT_SAMPLED, SAMPLE_DEFAULT);
     }
 
     public String getTrustedAccountKey() {
@@ -56,4 +68,12 @@ public String getAccountId() {
     public boolean isIncludeNewRelicHeader() {
         return includeNewRelicHeader;
     }
+
+    public String getRemoteParentSampled() {
+        return remoteParentSampled;
+    }
+
+    public String getRemoteParentNotSampled() {
+        return remoteParentNotSampled;
+    }
 }
diff --git a/newrelic-agent/src/main/resources/newrelic.yml b/newrelic-agent/src/main/resources/newrelic.yml
@@ -311,6 +311,21 @@ common: &default_settings
     # Default is false.
     exclude_newrelic_header: false
 
+    # Agent version 8.20.0+ utilize these settings for controlling how to sample when we have a remote parent
+    # involved. That is, if there is a valid traceparent with a value in the sampled flag on the inbound headers.
+    # If the inbound traceparent is marked as sampled then the remote_parent_sampled setting will be used.
+    # If the inbound traceparent is marked as NOT sampled then the remote_parent_not_sampled setting will be used.
+    # Otherwise New Relic's standard sampling will be used.
+    # Note: Reservoir limits are still in effect on top of these settings, so spans may be dropped even if they
+    # are marked for sampling, if the reservoir reaches its limit (max_samples_stored).
+    # Possible values:
+    #   default: New Relic's standard sampling rules will be used.
+    #   always_on: Always sample spans in this case.
+    #   always_off: Always skip sampling in this case.
+    sampler:
+      remote_parent_sampled: default
+      remote_parent_not_sampled: default
+
   # New Relic's distributed tracing UI uses Span events to display traces across different services.
   # Span events capture attributes that describe execution context and provide linking metadata.
   # Span events require distributed tracing to be enabled.
diff --git a/newrelic-agent/src/test/java/com/newrelic/agent/HeadersUtilTest.java b/newrelic-agent/src/test/java/com/newrelic/agent/HeadersUtilTest.java
@@ -18,12 +18,14 @@
 import com.newrelic.agent.tracers.ClassMethodSignature;
 import com.newrelic.agent.tracers.DefaultTracer;
 import com.newrelic.agent.tracers.Tracer;
+import com.newrelic.agent.transaction.TransactionTimer;
 import com.newrelic.agent.util.MockDistributedTraceService;
 import com.newrelic.api.agent.HeaderType;
 import com.newrelic.api.agent.InboundHeaders;
 import com.newrelic.test.marker.RequiresFork;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.mockito.Mockito;
 
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
@@ -105,6 +107,97 @@ public void testGetSyntheticsInfoHeader() {
                 (ImmutableMap.of("X-NewRelic-Synthetics-Info", synthInfoValue), HeaderType.HTTP)));
     }
 
+    @Test
+    public void testInboundParentSampledTrueConfigAlwaysOn() {
+        Transaction tx = setupAndCreateTx(DistributedTracingConfig.SAMPLE_ALWAYS_ON, DistributedTracingConfig.SAMPLE_DEFAULT);
+        InboundHeaders inboundHeaders = createInboundHeaders(ImmutableMap.of(
+                "traceparent", "01-0123456789abcdef0123456789abcdef-0123456789abcdef-01", // last entry is the sampled flag = true
+                "tracestate", "trustyrusty@nr=0-0-709288-8599547-f85f42fd82a4cf1d-164d3b4b0d09cb05164d3b4b0d09cb05--0.5-1563574856827"
+        ), HeaderType.HTTP);
+
+        HeadersUtil.parseAndAcceptDistributedTraceHeaders(tx, inboundHeaders);
+        assertTrue(tx.getSpanProxy().getInitiatingW3CTraceParent().sampled());
+        assertTrue(tx.getPriority() == 2.0f);
+
+        Transaction.clearTransaction();
+    }
+
+    @Test
+    public void testInboundParentSampledTrueConfigAlwaysOff() {
+        Transaction tx = setupAndCreateTx(DistributedTracingConfig.SAMPLE_ALWAYS_OFF, DistributedTracingConfig.SAMPLE_DEFAULT);
+        InboundHeaders inboundHeaders = createInboundHeaders(ImmutableMap.of(
+                "traceparent", "01-0123456789abcdef0123456789abcdef-0123456789abcdef-01", // last entry is the sampled flag = true
+                "tracestate", "trustyrusty@nr=0-0-709288-8599547-f85f42fd82a4cf1d-164d3b4b0d09cb05164d3b4b0d09cb05--0.5-1563574856827"
+        ), HeaderType.HTTP);
+
+        HeadersUtil.parseAndAcceptDistributedTraceHeaders(tx, inboundHeaders);
+        assertTrue(tx.getSpanProxy().getInitiatingW3CTraceParent().sampled());
+        assertTrue(tx.getPriority() == 0.0f);
+
+        Transaction.clearTransaction();
+    }
+
+    @Test
+    public void testInboundParentSampledFalseConfigAlwaysOn() {
+        Transaction tx = setupAndCreateTx(DistributedTracingConfig.SAMPLE_DEFAULT, DistributedTracingConfig.SAMPLE_ALWAYS_ON);
+        InboundHeaders inboundHeaders = createInboundHeaders(ImmutableMap.of(
+                "traceparent", "01-0123456789abcdef0123456789abcdef-0123456789abcdef-00", // last entry is the sampled flag = true
+                "tracestate", "trustyrusty@nr=0-0-709288-8599547-f85f42fd82a4cf1d-164d3b4b0d09cb05164d3b4b0d09cb05--0.5-1563574856827"
+        ), HeaderType.HTTP);
+
+        HeadersUtil.parseAndAcceptDistributedTraceHeaders(tx, inboundHeaders);
+        assertFalse(tx.getSpanProxy().getInitiatingW3CTraceParent().sampled());
+        assertTrue(tx.getPriority() == 2.0f);
+
+        Transaction.clearTransaction();
+    }
+
+    @Test
+    public void testInboundParentSampledFalseConfigAlwaysOff() {
+        Transaction tx = setupAndCreateTx(DistributedTracingConfig.SAMPLE_DEFAULT, DistributedTracingConfig.SAMPLE_ALWAYS_OFF);
+        InboundHeaders inboundHeaders = createInboundHeaders(ImmutableMap.of(
+                "traceparent", "01-0123456789abcdef0123456789abcdef-0123456789abcdef-00", // last entry is the sampled flag = true
+                "tracestate", "trustyrusty@nr=0-0-709288-8599547-f85f42fd82a4cf1d-164d3b4b0d09cb05164d3b4b0d09cb05--0.5-1563574856827"
+        ), HeaderType.HTTP);
+
+        HeadersUtil.parseAndAcceptDistributedTraceHeaders(tx, inboundHeaders);
+        assertFalse(tx.getSpanProxy().getInitiatingW3CTraceParent().sampled());
+        assertTrue(tx.getPriority() == 0.0f);
+
+        Transaction.clearTransaction();
+    }
+
+    private Transaction setupAndCreateTx(String remoteParentSampled, String remoteParentNotSampled) {
+        System.out.println("Setting up config: "+remoteParentSampled+"; "+remoteParentNotSampled);
+        ConfigService mockConfigService = new MockConfigService(AgentConfigImpl.createAgentConfig(
+                ImmutableMap.of(
+                        AgentConfigImpl.APP_NAME,
+                        "Unit Test",
+                        AgentConfigImpl.DISTRIBUTED_TRACING,
+                        ImmutableMap.of(
+                                DistributedTracingConfig.ENABLED, true,
+                                DistributedTracingConfig.SAMPLER,
+                                ImmutableMap.of(
+                                        DistributedTracingConfig.REMOTE_PARENT_SAMPLED, remoteParentSampled,
+                                        DistributedTracingConfig.REMOTE_PARENT_NOT_SAMPLED, remoteParentNotSampled)
+                        ),
+                        AgentConfigImpl.SPAN_EVENTS,
+                        ImmutableMap.of(
+                                SpanEventsConfig.ENABLED, true,
+                                SpanEventsConfig.COLLECT_SPAN_EVENTS, true)
+                )));
+        MockServiceManager mockServiceManager = new MockServiceManager(mockConfigService);
+        mockServiceManager.setDistributedTraceService(new MockDistributedTraceService());
+        ServiceFactory.setServiceManager(mockServiceManager);
+        Transaction tx = Mockito.spy(Transaction.getTransaction());
+        TransactionTimer timer = Mockito.mock(TransactionTimer.class);
+        Mockito.when(timer.getStartTimeInNanos()).thenReturn(System.nanoTime());
+        Mockito.when(tx.getTransactionTimer()).thenReturn(timer);
+        tx.setPriorityIfNotNull(0.5f);  // something > 0 and < 1
+
+        return tx;
+    }
+
     private InboundHeaders createInboundHeaders(final Map<String, String> map, final HeaderType type) {
         return new InboundHeaders() {
             @Override
diff --git a/newrelic-agent/src/test/java/com/newrelic/agent/tracers/DefaultTracerClmTest.java b/newrelic-agent/src/test/java/com/newrelic/agent/tracers/DefaultTracerClmTest.java
@@ -888,7 +888,7 @@ public void testTransactionABParenting() {
         DefaultTracer span5Tracer = new OtherRootTracer(txB, new ClassMethodSignature("class",
                 "span5", "()V"), null, DefaultTracer.NULL_METRIC_NAME_FORMATTER);
         txB.getTransactionActivity().tracerStarted(span5Tracer);
-        span5Tracer.getTransaction().acceptDistributedTracePayload(payload);
+        span5Tracer.getTransaction().acceptDistributedTracePayload(payload, null);
 
         txB.setTransactionName(com.newrelic.api.agent.TransactionNamePriority.CUSTOM_HIGH, true, "Transaction B");
         txB.setThrowable(new Throwable(), TransactionErrorPriority.API, false);
diff --git a/newrelic-agent/src/test/java/com/newrelic/agent/tracers/DefaultTracerTest.java b/newrelic-agent/src/test/java/com/newrelic/agent/tracers/DefaultTracerTest.java
@@ -924,7 +924,7 @@ public void testTransactionABParenting() {
         DefaultTracer span5Tracer = new OtherRootTracer(txB, new ClassMethodSignature("class",
                 "span5", "()V"), null, DefaultTracer.NULL_METRIC_NAME_FORMATTER);
         txB.getTransactionActivity().tracerStarted(span5Tracer);
-        span5Tracer.getTransaction().acceptDistributedTracePayload(payload);
+        span5Tracer.getTransaction().acceptDistributedTracePayload(payload, null);
 
         txB.setTransactionName(com.newrelic.api.agent.TransactionNamePriority.CUSTOM_HIGH, true, "Transaction B");
         txB.setThrowable(new Throwable(), TransactionErrorPriority.API, false);
diff --git a/newrelic-agent/src/test/java/com/newrelic/agent/tracing/W3CTraceStateSupportTest.java b/newrelic-agent/src/test/java/com/newrelic/agent/tracing/W3CTraceStateSupportTest.java
@@ -120,7 +120,7 @@ public void testParseHeadersNoPriorityOrSampledTriggersResamplingDecision() {
                 Collections.singletonList("02-12341234123412341234123412341234-4321432143214321-01"),
                 Arrays.asList("190@nr=0-0-709288-8599547-f85f42fd82a4cf1d-164d3b4b0d09cb05164d3b4b0d09cb05---1563574856827")).getPayload();
         assertNotNull(inboundPayload);
-        transaction.acceptDistributedTracePayload(inboundPayload);
+        transaction.acceptDistributedTracePayload(inboundPayload, null);
 
         transaction.createDistributedTracePayload("meatball!");
         String[] outboundPayload = new W3CTraceStateHeader(true, true).create(transaction.getSpanProxy()).split("-");
@@ -159,7 +159,7 @@ public void testParseHeadersDup() {
                 Arrays.asList("190@nr=0-0-709288-8599547-f85f42fd82a4cf1d-164d3b4b0d09cb05164d3b4b0d09cb05-1-0.789-1563574856827", "congo@=0-qzx932-abc123",
                         "congo@=0-very-qzx932-abc123")).getPayload();
         assertNotNull(inboundPayload);
-        transaction.acceptDistributedTracePayload(inboundPayload);
+        transaction.acceptDistributedTracePayload(inboundPayload, null);
 
         String newSpanId = TransactionGuidFactory.generate16CharGuid();
         transaction.createDistributedTracePayload(newSpanId);

Original file line number	Diff line number	Diff line change
`@@ -261,7 +261,7 @@ public static void parseAndAcceptDistributedTraceHeaders(Transaction tx, Inbound`
`261`	`261`	`W3CTracePayload w3CTracePayload = W3CTracePayload.parseHeaders(tx, traceParent, traceState);`
`262`	`262`	`if (w3CTracePayload != null) {`
`263`	`263`	`if (w3CTracePayload.getPayload() != null) {`
`264`		`- tx.acceptDistributedTracePayload(w3CTracePayload.getPayload());`
	`264`	`+ tx.acceptDistributedTracePayload(w3CTracePayload.getPayload(), w3CTracePayload.getTraceParent());`
`265`	`265`	`}`
`266`	`266`	`if (w3CTracePayload.getTraceParent() != null) {`
`267`	`267`	`tx.getSpanProxy().setInitiatingW3CTraceParent(w3CTracePayload.getTraceParent());`
Original file line number	Diff line number	Diff line change
`@@ -509,7 +509,7 @@ public void acceptDistributedTracePayload(String payload) {`
`509`	`509`	`public void acceptDistributedTracePayload(DistributedTracePayload payload) {`
`510`	`510`	`Transaction tx = getTransactionIfExists();`
`511`	`511`	`if (tx != null) {`
`512`		`- tx.acceptDistributedTracePayload(payload);`
	`512`	`+ tx.acceptDistributedTracePayload(payload, null);`
`513`	`513`	`}`
`514`	`514`	`}`
`515`	`515`