Introduce trace.upstream.services.tracking.enabled

DataDog · Jan 6, 2022 · 6e81a99 · 6e81a99
1 parent cb9df7f
commit 6e81a99
Show file tree

Hide file tree

Showing 22 changed files with 121 additions and 63 deletions.
diff --git a/dd-java-agent/instrumentation/opentelemetry/src/test/groovy/OpenTelemetryTest.groovy b/dd-java-agent/instrumentation/opentelemetry/src/test/groovy/OpenTelemetryTest.groovy
@@ -2,8 +2,6 @@ import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDId
 import datadog.trace.api.DDTags
 import datadog.trace.api.interceptor.MutableSpan
-import static datadog.trace.api.sampling.PrioritySampling.*
-import static datadog.trace.api.sampling.SamplingMechanism.*
 import datadog.trace.context.TraceScope
 import datadog.trace.core.propagation.ExtractedContext
 import io.grpc.Context
@@ -15,6 +13,9 @@ import io.opentelemetry.trace.Status
 import io.opentelemetry.trace.TracingContextUtils
 import spock.lang.Subject
 
+import static datadog.trace.api.sampling.PrioritySampling.*
+import static datadog.trace.api.sampling.SamplingMechanism.*
+
 class OpenTelemetryTest extends AgentTestRunner {
   @Subject
   def tracer = OpenTelemetry.tracerProvider.get("test-inst")

diff --git a/dd-java-agent/instrumentation/opentracing/api-0.31/src/test/groovy/OpenTracing31Test.groovy b/dd-java-agent/instrumentation/opentracing/api-0.31/src/test/groovy/OpenTracing31Test.groovy
@@ -2,8 +2,6 @@ import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDId
 import datadog.trace.api.DDTags
 import datadog.trace.api.interceptor.MutableSpan
-import static datadog.trace.api.sampling.PrioritySampling.*
-import static datadog.trace.api.sampling.SamplingMechanism.*
 import datadog.trace.context.TraceScope
 import datadog.trace.core.DDSpan
 import datadog.trace.core.propagation.ExtractedContext
@@ -18,6 +16,8 @@ import io.opentracing.util.GlobalTracer
 import spock.lang.Subject
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.api.sampling.PrioritySampling.*
+import static datadog.trace.api.sampling.SamplingMechanism.*
 
 class OpenTracing31Test extends AgentTestRunner {
 

diff --git a/dd-java-agent/instrumentation/opentracing/api-0.32/src/test/groovy/OpenTracing32Test.groovy b/dd-java-agent/instrumentation/opentracing/api-0.32/src/test/groovy/OpenTracing32Test.groovy
@@ -2,8 +2,6 @@ import datadog.trace.agent.test.AgentTestRunner
 import datadog.trace.api.DDId
 import datadog.trace.api.DDTags
 import datadog.trace.api.interceptor.MutableSpan
-import static datadog.trace.api.sampling.PrioritySampling.*
-import static datadog.trace.api.sampling.SamplingMechanism.*
 import datadog.trace.context.TraceScope
 import datadog.trace.core.DDSpan
 import datadog.trace.core.propagation.ExtractedContext
@@ -18,6 +16,8 @@ import io.opentracing.util.GlobalTracer
 import spock.lang.Subject
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.api.sampling.PrioritySampling.*
+import static datadog.trace.api.sampling.SamplingMechanism.*
 
 class OpenTracing32Test extends AgentTestRunner {
 

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java
@@ -58,8 +58,20 @@ public final class TracerConfig {
   public static final String PROPAGATION_STYLE_INJECT = "propagation.style.inject";
 
   public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled";
+
+  /**
+   * Disables validation that prevents invalid combinations of sampling priority and sampling
+   * mechanism on the set sampling priority calls. This check is enabled by default.
+   */
   public static final String SAMPLING_MECHANISM_VALIDATION_DISABLED =
       "trace.sampling.mechanism.validation.disabled";
 
+  /**
+   * Enables sampling decision updates to upstream_services part of x-datadog-tags. Enabled by
+   * default.
+   */
+  public static final String UPSTREAM_SERVICES_TRACKING_ENABLED =
+      "trace.upstream.services.tracking.enabled";
+
   private TracerConfig() {}
 }
diff --git a/dd-trace-core/src/jmh/java/datadog/trace/core/PendingTraceWrite.java b/dd-trace-core/src/jmh/java/datadog/trace/core/PendingTraceWrite.java
@@ -58,7 +58,8 @@ public void init(TraceCounters counters, Blackhole blackhole) {
                 trace,
                 null,
                 false,
-                null));
+                null,
+                false));
     span =
         DDSpan.create(
             System.currentTimeMillis() * 1000,
@@ -80,7 +81,8 @@ public void init(TraceCounters counters, Blackhole blackhole) {
                 trace,
                 null,
                 false,
-                null));
+                null,
+                false));
   }
 
   @Threads(4)

diff --git a/dd-trace-core/src/jmh/java/datadog/trace/core/TracerMapperMap.java b/dd-trace-core/src/jmh/java/datadog/trace/core/TracerMapperMap.java
@@ -116,6 +116,7 @@ private DDSpan createSpanWithOrigin(int iter, final String origin) {
             trace,
             null,
             false,
-            null));
+            null,
+            false));
   }
 }
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
@@ -117,6 +117,7 @@ public static CoreTracerBuilder builder() {
   private final SamplingCheckpointer checkpointer;
   private final ExternalAgentLauncher externalAgentLauncher;
   private boolean disableSamplingMechanismValidation;
+  private boolean enableUpstreamServicesTracking;
 
   /**
    * JVM shutdown callback, keeping a reference to it to remove this if DDTracer gets destroyed
@@ -433,6 +434,7 @@ private CoreTracer(
     this.externalAgentLauncher = new ExternalAgentLauncher(config);
 
     this.disableSamplingMechanismValidation = config.isSamplingMechanismValidationDisabled();
+    this.enableUpstreamServicesTracking = config.isUpstreamServicesTrackingEnabled();
 
     if (sharedCommunicationObjects == null) {
       sharedCommunicationObjects = new SharedCommunicationObjects();
@@ -1104,7 +1106,8 @@ private DDSpanContext buildSpanContext() {
               parentTrace,
               requestContextData,
               disableSamplingMechanismValidation,
-              ddTags);
+              ddTags,
+              enableUpstreamServicesTracking);
 
       // By setting the tags on the context we apply decorators to any tags that have been set via
       // the builder. This is the order that the tags were added previously, but maybe the `tags`

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java b/dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java
@@ -105,7 +105,8 @@ public class DDSpanContext implements AgentSpan.Context, RequestContext<Object>,
 
   private final boolean disableSamplingMechanismValidation;
 
-  private DatadogTags ddTags;
+  private final boolean enableUpstreamServicesTracking;
+  private final DatadogTags ddTags;
 
   /** Aims to pack sampling priority and sampling mechanism into one value */
   protected static class SamplingDecision {
@@ -146,7 +147,8 @@ public DDSpanContext(
       final PendingTrace trace,
       final Object requestContextData,
       final boolean disableSamplingMechanismValidation,
-      final DatadogTags ddTags) {
+      final DatadogTags ddTags,
+      boolean enableUpstreamServicesTracking) {
 
     assert trace != null;
     this.trace = trace;
@@ -181,10 +183,9 @@ public DDSpanContext(
     this.spanType = spanType;
     this.origin = origin;
 
-    long samplingParams = SamplingDecision.create(samplingPriority, samplingMechanism);
-    if (samplingParams != SamplingDecision.UNSET_UNKNOWN) {
-      setSamplingPriority(samplingPriority, samplingMechanism);
-    }
+    // enableUpstreamServicesTracking has to be set prior to calling setSamplingPriority
+    this.enableUpstreamServicesTracking = enableUpstreamServicesTracking;
+    setSamplingPriority(samplingPriority, samplingMechanism);
 
     // Additional Metadata
     final Thread current = Thread.currentThread();
@@ -214,7 +215,6 @@ public String getServiceName() {
 
   public void setServiceName(final String serviceName) {
     String newServiceName = trace.getTracer().mapServiceName(serviceName);
-    ddTags.updateServiceName(serviceName, newServiceName);
     this.serviceName = newServiceName;
     this.topLevel = isTopLevel(parentServiceName, this.serviceName);
   }
@@ -356,7 +356,9 @@ public boolean setSamplingPriority(
       return false;
     }
 
-    ddTags.updateUpstreamServices(getServiceName(), newPriority, newMechanism, rate);
+    if (enableUpstreamServicesTracking) {
+      ddTags.updateUpstreamServices(getServiceName(), newPriority, newMechanism, rate);
+    }
     return true;
   }
 

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogTags.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogTags.java
@@ -80,13 +80,22 @@ public DatadogTags(String rawTags) {
    * @param rate - sampling rate, pass a negative value if not applicable
    */
   public void updateUpstreamServices(String serviceName, int priority, int mechanism, double rate) {
-    if (priority != PrioritySampling.UNSET) {
+    if (priority != PrioritySampling.UNSET
+        && (samplingDecision == null
+            || samplingDecision.priority != priority
+            || samplingDecision.mechanism != mechanism
+            || samplingDecision.rate != rate
+            || !samplingDecision.service.equals(serviceName))) {
       ServiceSamplingDecision newSamplingDecision =
           new ServiceSamplingDecision(serviceName, priority, mechanism, rate);
       SAMPLING_DECISION_UPDATER.compareAndSet(this, samplingDecision, newSamplingDecision);
     }
   }
 
+  public boolean isEmpty() {
+    return rawTags.isEmpty() && samplingDecision == null;
+  }
+
   /** @return encoded header value */
   public String encoded() {
     boolean isSamplingDecisionEmpty = samplingDecision == null || samplingDecision.isUnset();
@@ -147,13 +156,4 @@ private void encodeUpstreamServices(StringBuilder sb) {
     sb.append('=');
     samplingDecision.encoded(sb);
   }
-
-  public void updateServiceName(String serviceName, String newServiceName) {
-    // TODO update service name
-    // do we even need to update it?
-  }
-
-  public boolean isEmpty() {
-    return rawTags.isEmpty() && samplingDecision == null;
-  }
 }
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/writer/DDAgentApiTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/writer/DDAgentApiTest.groovy
@@ -441,7 +441,8 @@ class DDAgentApiTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.from(1)),
       null,
       false,
-      null)
+      null,
+      false)
 
     def span = DDSpan.create(timestamp, context)
     span.setTag(tag, value)

diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/writer/DDAgentWriterCombinedTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/writer/DDAgentWriterCombinedTest.groovy
@@ -285,7 +285,8 @@ class DDAgentWriterCombinedTest extends DDCoreSpecification {
       trace,
       null,
       false,
-      null)
+      null,
+      false)
   }
 
   def createMinimalTrace() {

diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/writer/DDAgentWriterTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/writer/DDAgentWriterTest.groovy
@@ -168,7 +168,8 @@ class DDAgentWriterTest extends DDCoreSpecification {
     tracer.mapServiceName(_) >> { String serviceName -> serviceName }
     PendingTrace trace = Mock(PendingTrace)
     trace.getTracer() >> tracer
-    return new DDSpan(0, new DDSpanContext(
+
+    def context = new DDSpanContext(
       DDId.from(1),
       DDId.from(1),
       DDId.ZERO,
@@ -186,6 +187,8 @@ class DDAgentWriterTest extends DDCoreSpecification {
       trace,
       null,
       false,
-      null))
+      null,
+      false)
+    return new DDSpan(0, context)
   }
 }
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/common/writer/PayloadDispatcherTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/common/writer/PayloadDispatcherTest.groovy
@@ -151,7 +151,8 @@ class PayloadDispatcherTest extends DDSpecification {
     tracer.mapServiceName(_) >> { String serviceName -> serviceName }
     PendingTrace trace = Mock(PendingTrace)
     trace.getTracer() >> tracer
-    return new DDSpan(0, new DDSpanContext(
+
+    def context = new DDSpanContext(
       DDId.from(1),
       DDId.from(1),
       DDId.ZERO,
@@ -169,6 +170,8 @@ class PayloadDispatcherTest extends DDSpecification {
       trace,
       null,
       false,
-      null))
+      null,
+      false)
+    return new DDSpan(0, context)
   }
 }
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanSerializationTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanSerializationTest.groovy
@@ -148,7 +148,8 @@ class DDSpanSerializationTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.ONE),
       null,
       false,
-      null)
+      null,
+      false)
     context.setAllTags(tags)
     def span = DDSpan.create(0, context)
     CaptureBuffer capture = new CaptureBuffer()
@@ -218,7 +219,8 @@ class DDSpanSerializationTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.ONE),
       null,
       false,
-      null)
+      null,
+      false)
     context.setAllTags(tags)
     def span = DDSpan.create(0, context)
     CaptureBuffer capture = new CaptureBuffer()
@@ -298,7 +300,8 @@ class DDSpanSerializationTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.ONE),
       null,
       false,
-      null)
+      null,
+      false)
     ctx.setAllTags(["k1": "v1"])
     return ctx
   }

diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanTest.groovy
@@ -333,7 +333,8 @@ class DDSpanTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.ONE),
       null,
       false,
-      null)
+      null,
+      false)
     then:
     context.isTopLevel() == expectTopLevel
 
@@ -370,7 +371,8 @@ class DDSpanTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.ONE),
       null,
       false,
-      null)
+      null,
+      false)
 
     def span = null
 

diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/PendingTraceBufferTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/PendingTraceBufferTest.groovy
@@ -320,7 +320,8 @@ class PendingTraceBufferTest extends DDSpecification {
       trace,
       null,
       false,
-      null)
+      null,
+      false)
     return DDSpan.create(0, context)
   }
 
@@ -344,7 +345,8 @@ class PendingTraceBufferTest extends DDSpecification {
       trace,
       null,
       false,
-      null)
+      null,
+      false)
     return DDSpan.create(0, context)
   }
 }
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/B3HttpInjectorTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/propagation/B3HttpInjectorTest.groovy
@@ -42,7 +42,8 @@ class B3HttpInjectorTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.ONE),
       null,
       false,
-      null)
+      null,
+      false)
 
     final Map<String, String> carrier = Mock()
 
@@ -103,7 +104,8 @@ class B3HttpInjectorTest extends DDCoreSpecification {
       tracer.pendingTraceFactory.create(DDId.ONE),
       null,
       false,
-      null)
+      null,
+      false)
     final Map<String, String> carrier = Mock()
 
     when: