Fix Single Span Sampling for parent spans

dd-trace-core/src/jmh/java/datadog/trace/core/propagation/InjectorBenchmark.java

@@ -137,7 +137,7 @@ public void injectContext(Blackhole blackhole) {
     blackhole.consume(headers);
     if (modifyPropagationTags) {
       int sm = mechanism = (mechanism + 1) % 4;
-      propagationTags.updateTraceSamplingPriority(1, sm, "service");
+      propagationTags.updateTraceSamplingPriority(1, sm);
     }
   }
 

dd-trace-core/src/main/java/datadog/trace/core/DDSpanContext.java

@@ -113,6 +113,13 @@ public class DDSpanContext implements AgentSpan.Context, RequestContext, TraceSe
 
   private volatile int samplingPriority = PrioritySampling.UNSET;
 
+  // A flag that this span has been kept by the Single Span Sampling mechanism. This is needed for
+  // when it's a root span that is set to be dropped by the trace sampling, but kept by the single
+  // span sampling. In this case, we should NOT override the original samplingPriority that is used
+  // by the child spans that don't have an explicitly set trace sampling priority and gets it from
+  // the root span.
+  private volatile boolean isSelectedBySingleSpanSampling = false;
+
   /** The origin of the trace. (eg. Synthetics, CI App) */
   private volatile CharSequence origin;
 
@@ -300,8 +307,7 @@ private void forceKeepThisSpan(byte samplingMechanism) {
     // even if the old sampling priority and mechanism have already propagated
     if (SAMPLING_PRIORITY_UPDATER.getAndSet(this, PrioritySampling.USER_KEEP)
         == PrioritySampling.UNSET) {
-      propagationTags.updateTraceSamplingPriority(
-          PrioritySampling.USER_KEEP, samplingMechanism, serviceName);
+      propagationTags.updateTraceSamplingPriority(PrioritySampling.USER_KEEP, samplingMechanism);
     }
   }
 
@@ -342,7 +348,7 @@ private boolean setThisSpanSamplingPriority(final int newPriority, final int new
       return false;
     }
     // set trace level sampling priority tag propagationTags
-    propagationTags.updateTraceSamplingPriority(newPriority, newMechanism, serviceName);
+    propagationTags.updateTraceSamplingPriority(newPriority, newMechanism);
     return true;
   }
 
@@ -372,19 +378,23 @@ private boolean validateSamplingPriority(final int newPriority, final int newMec
     return true;
   }
 
-  /** @return the sampling priority of this span's trace, or null if no priority has been set */
+  /**
+   * @return the trace sampling priority of this span's trace, or null if no priority has been set
+   */
   public int getSamplingPriority() {
     return getRootSpanContextOrThis().samplingPriority;
   }
 
   public void setSpanSamplingPriority(double rate, int limit) {
     synchronized (unsafeTags) {
-      forceKeepThisSpan(SamplingMechanism.SPAN_SAMPLING_RATE);
       unsafeSetTag(SPAN_SAMPLING_MECHANISM_TAG, SamplingMechanism.SPAN_SAMPLING_RATE);
       unsafeSetTag(SPAN_SAMPLING_RULE_RATE_TAG, rate);
       if (limit != Integer.MAX_VALUE) {
         unsafeSetTag(SPAN_SAMPLING_MAX_PER_SECOND_TAG, limit);
       }
+      propagationTags.updateTraceSamplingPriority(
+          PrioritySampling.USER_KEEP, SamplingMechanism.SPAN_SAMPLING_RATE);
+      isSelectedBySingleSpanSampling = true;
     }
   }
 
@@ -594,7 +604,7 @@ public void processTagsAndBaggage(final MetadataConsumer consumer) {
               threadName,
               postProcessor.processTags(unsafeTags),
               baggageItemsWithPropagationTags,
-              samplingPriority != PrioritySampling.UNSET ? samplingPriority : getSamplingPriority(),
+              getEffectiveSamplingPriority(),
               measured,
               topLevel,
               httpStatusCode == 0 ? null : HTTP_STATUSES.get(httpStatusCode),
@@ -603,6 +613,19 @@ public void processTagsAndBaggage(final MetadataConsumer consumer) {
     }
   }
 
+  // used in tests
+  int getEffectiveSamplingPriority() {
+    if (isSelectedBySingleSpanSampling) {
+      // use span sampling priority if this span has been selected by the single span sampler
+      return PrioritySampling.USER_KEEP;
+    } else if (samplingPriority != PrioritySampling.UNSET) {
+      // use trace sampling priority, if it's been set
+      return samplingPriority;
+    }
+    // otherwise get sampling priority from the root span
+    return getSamplingPriority();
+  }
+
   @Override
   public String toString() {
     final StringBuilder s =

dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTags.java

@@ -47,8 +47,7 @@ public interface Factory {
    * Updates the trace-level sampling priority decision if it hasn't already been made and _dd.p.dm
    * tag doesn't exist. Called on the root span context.
    */
-  public abstract void updateTraceSamplingPriority(
-      int samplingPriority, int samplingMechanism, String serviceName);
+  public abstract void updateTraceSamplingPriority(int samplingPriority, int samplingMechanism);
 
   /**
    * Constructs a header value that includes valid propagated _dd.p.* tags and possibly a new

dd-trace-core/src/main/java/datadog/trace/core/propagation/PropagationTagsFactory.java

@@ -53,8 +53,7 @@ private ValidPropagationTags(List<String> tagPairs, int tagsSize, boolean hasDec
     }
 
     @Override
-    public void updateTraceSamplingPriority(
-        int samplingPriority, int samplingMechanism, String serviceName) {
+    public void updateTraceSamplingPriority(int samplingPriority, int samplingMechanism) {
 
       if (samplingPriority != PrioritySampling.UNSET && isDecisionMakerTagMissing) {
         if (samplingPriority > 0) {
@@ -110,8 +109,7 @@ private InvalidPropagationTags(String error) {
     }
 
     @Override
-    public void updateTraceSamplingPriority(
-        int samplingPriority, int samplingMechanism, String serviceName) {}
+    public void updateTraceSamplingPriority(int samplingPriority, int samplingMechanism) {}
 
     @Override
     public String headerValue(HeaderType headerType) {

dd-trace-core/src/test/groovy/datadog/trace/common/sampling/SingleSpanSamplerTest.groovy

@@ -7,6 +7,10 @@ import datadog.trace.core.test.DDCoreSpecification
 
 import static datadog.trace.api.config.TracerConfig.SPAN_SAMPLING_RULES
 import static datadog.trace.api.config.TracerConfig.SPAN_SAMPLING_RULES_FILE
+import static datadog.trace.api.config.TracerConfig.TRACE_SAMPLE_RATE
+import static datadog.trace.api.sampling.SamplingMechanism.DEFAULT
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP
+import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP
 import static datadog.trace.api.sampling.SamplingMechanism.SPAN_SAMPLING_RATE
 
 class SingleSpanSamplerTest extends DDCoreSpecification {
@@ -72,6 +76,50 @@ class SingleSpanSamplerTest extends DDCoreSpecification {
     """[ { "service": "*", "name": "operation-b", "sample_rate": 0.5 } ]"""                           | false          | null               | null         | null
   }
 
+  def "Parent/child scenarios when the trace is dropped but individual spans are kept by the single span sampler"() {
+    given:
+    Properties properties = new Properties()
+    if (rules != null) {
+      properties.setProperty(SPAN_SAMPLING_RULES, rules)
+      properties.setProperty(TRACE_SAMPLE_RATE, "0")
+    }
+    def tracer = tracerBuilder().writer(new ListWriter()).build()
+
+    when:
+    SingleSpanSampler sampler = SingleSpanSampler.Builder.forConfig(Config.get(properties))
+
+    DDSpan rootSpan = tracer.buildSpan("web.request")
+      .withServiceName("webserver")
+      .ignoreActiveSpan().start() as DDSpan
+
+    DDSpan childSpan = tracer.buildSpan("web.handler")
+      .withServiceName("webserver")
+      .asChildOf(rootSpan)
+      .ignoreActiveSpan().start() as DDSpan
+
+    then:
+    // set trace sampling priority to drop the trace
+    rootSpan.setSamplingPriority(SAMPLER_DROP, DEFAULT)
+
+    // set spans sampling priority
+    sampler.setSamplingPriority(rootSpan) == sampleRoot
+    sampler.setSamplingPriority(childSpan) == sampleChild
+
+    rootSpan.context().effectiveSamplingPriority == (sampleRoot ? USER_KEEP : SAMPLER_DROP)
+    childSpan.context().effectiveSamplingPriority == (sampleChild ? USER_KEEP : SAMPLER_DROP)
+
+    expect:
+    rootSpan.getTag("_dd.span_sampling.mechanism") == rootMechanism
+    childSpan.getTag("_dd.span_sampling.mechanism") == childMechanism
+
+    where:
+    rules                                                   | sampleRoot | sampleChild | rootMechanism      | childMechanism
+    """[{"service": "webserver", "name": "web.request"}]""" | true       | false       | SPAN_SAMPLING_RATE | null
+    """[{"service": "webserver", "name": "web.handler"}]""" | false      | true        | null               | SPAN_SAMPLING_RATE
+    """[{"service": "webserver", "name": "web.*"}]"""       | true       | true        | SPAN_SAMPLING_RATE | SPAN_SAMPLING_RATE
+    """[{"service": "other-server"}]"""                     | false      | false       | null               | null
+  }
+
   def "Single Span Sampler set sampling priority with the max-per-second limit"() {
     given:
     Properties properties = new Properties()

dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanContextTest.groovy

@@ -217,6 +217,7 @@ class DDSpanContextTest extends DDCoreSpecification {
 
     expect:
     context.getSamplingPriority() == UNSET
+    context.effectiveSamplingPriority == UNSET
 
     when:
     context.setSpanSamplingPriority(rate, limit)
@@ -225,7 +226,9 @@ class DDSpanContextTest extends DDCoreSpecification {
     context.getTag(SPAN_SAMPLING_MECHANISM_TAG) == SPAN_SAMPLING_RATE
     context.getTag(SPAN_SAMPLING_RULE_RATE_TAG) == rate
     context.getTag(SPAN_SAMPLING_MAX_PER_SECOND_TAG) == (limit == Integer.MAX_VALUE ? null : limit)
-    context.getSamplingPriority() == USER_KEEP
+    // single span sampling should not change the trace sampling priority
+    context.getSamplingPriority() == UNSET
+    context.effectiveSamplingPriority == USER_KEEP
     context.getPropagationTags().createTagMap() == ["_dd.p.dm":"-" + SPAN_SAMPLING_RATE]
 
     where:

dd-trace-core/src/test/groovy/datadog/trace/core/DDSpanTest.groovy

@@ -423,7 +423,9 @@ class DDSpanTest extends DDCoreSpecification {
     span.getTag(SPAN_SAMPLING_MECHANISM_TAG) == SPAN_SAMPLING_RATE
     span.getTag(SPAN_SAMPLING_RULE_RATE_TAG) == rate
     span.getTag(SPAN_SAMPLING_MAX_PER_SECOND_TAG) == (limit == Integer.MAX_VALUE ? null : limit)
-    span.samplingPriority() == USER_KEEP
+    // single span sampling should not change the trace sampling priority
+    span.samplingPriority() == UNSET
+    span.context.effectiveSamplingPriority == USER_KEEP
 
     where:
     rate | limit

dd-trace-core/src/test/groovy/datadog/trace/core/propagation/DatadogPropagationTagsTest.groovy

@@ -75,7 +75,7 @@ class DatadogPropagationTagsTest extends DDCoreSpecification {
     def propagationTags = propagationTagsFactory.fromHeaderValue(PropagationTags.HeaderType.DATADOG, originalTagSet)
 
     when:
-    propagationTags.updateTraceSamplingPriority(priority, mechanism, "service-1")
+    propagationTags.updateTraceSamplingPriority(priority, mechanism)
 
     then:
     propagationTags.headerValue(PropagationTags.HeaderType.DATADOG) == expectedHeaderValue
@@ -115,7 +115,7 @@ class DatadogPropagationTagsTest extends DDCoreSpecification {
     def propagationTags = PropagationTags.factory(limit).fromHeaderValue(PropagationTags.HeaderType.DATADOG, tags)
 
     when:
-    propagationTags.updateTraceSamplingPriority(USER_KEEP, MANUAL, "service-name")
+    propagationTags.updateTraceSamplingPriority(USER_KEEP, MANUAL)
 
     then:
     propagationTags.headerValue(PropagationTags.HeaderType.DATADOG) == null
@@ -129,7 +129,7 @@ class DatadogPropagationTagsTest extends DDCoreSpecification {
     def propagationTags = PropagationTags.factory(limit).fromHeaderValue(PropagationTags.HeaderType.DATADOG, tags)
 
     when:
-    propagationTags.updateTraceSamplingPriority(USER_KEEP, MANUAL, "service-name")
+    propagationTags.updateTraceSamplingPriority(USER_KEEP, MANUAL)
 
     then:
     propagationTags.headerValue(PropagationTags.HeaderType.DATADOG) == null
@@ -141,7 +141,7 @@ class DatadogPropagationTagsTest extends DDCoreSpecification {
     def propagationTags = PropagationTags.factory(0).fromHeaderValue(PropagationTags.HeaderType.DATADOG, "")
 
     when:
-    propagationTags.updateTraceSamplingPriority(USER_KEEP, MANUAL, "service-name")
+    propagationTags.updateTraceSamplingPriority(USER_KEEP, MANUAL)
 
     then:
     propagationTags.headerValue(PropagationTags.HeaderType.DATADOG) == null