Merge remote-tracking branch 'origin/dougqh/sampling-provenance' into mcculls/handle-remote-sampling-rules_MERGE

Author: mcculls

dd-trace-core/src/main/java/datadog/trace/common/sampling/RateSamplingRule.java

@@ -1,5 +1,6 @@
 package datadog.trace.common.sampling;
 
+import datadog.trace.api.sampling.SamplingMechanism;
 import datadog.trace.core.CoreSpan;
 import datadog.trace.core.util.Matcher;
 import datadog.trace.core.util.Matchers;
@@ -10,9 +11,16 @@
 
 public abstract class RateSamplingRule {
   private final RateSampler sampler;
+  private final byte mechanism;
 
-  public RateSamplingRule(final RateSampler sampler) {
+  public RateSamplingRule(final RateSampler sampler, byte mechanism) {
     this.sampler = sampler;
+    this.mechanism = mechanism;
+  }
+
+  @Deprecated
+  public RateSamplingRule(final RateSampler sampler) {
+    this(sampler, SamplingMechanism.LOCAL_USER_RULE);
   }
 
   public abstract <T extends CoreSpan<T>> boolean matches(T span);
@@ -25,10 +33,13 @@ public RateSampler getSampler() {
     return sampler;
   }
 
-  public static class AlwaysMatchesSamplingRule extends RateSamplingRule {
+  public final byte getMechanism() {
+    return mechanism;
+  }
 
-    public AlwaysMatchesSamplingRule(final RateSampler sampler) {
-      super(sampler);
+  public static class AlwaysMatchesSamplingRule extends RateSamplingRule {
+    public AlwaysMatchesSamplingRule(final RateSampler sampler, byte samplingMechanism) {
+      super(sampler, samplingMechanism);
     }
 
     @Override
@@ -87,8 +98,10 @@ public TraceSamplingRule(
         final String operationGlob,
         final String resourceGlob,
         final Map<String, String> tags,
-        final RateSampler sampler) {
-      super(sampler);
+        final RateSampler sampler,
+        final byte samplingMechanism) {
+      super(sampler, samplingMechanism);
+
       serviceMatcher = Matchers.compileGlob(serviceGlob);
       operationMatcher = Matchers.compileGlob(operationGlob);
       resourceMatcher = Matchers.compileGlob(resourceGlob);

dd-trace-core/src/main/java/datadog/trace/common/sampling/RuleBasedTraceSampler.java

@@ -1,7 +1,6 @@
 package datadog.trace.common.sampling;
 
 import datadog.trace.api.config.TracerConfig;
-import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.api.sampling.SamplingMechanism;
 import datadog.trace.api.sampling.SamplingRule;
 import datadog.trace.core.CoreSpan;
@@ -68,7 +67,8 @@ public static RuleBasedTraceSampler build(
                 rule.getName(),
                 rule.getResource(),
                 rule.getTags(),
-                new DeterministicSampler.TraceSampler(rule.getSampleRate()));
+                new DeterministicSampler.TraceSampler(rule.getSampleRate()),
+                samplingMechanism(rule.getProvenance()));
         samplingRules.add(samplingRule);
       }
     } else {
@@ -102,16 +102,31 @@ public static RuleBasedTraceSampler build(
       }
     }
 
+    // Per spec, defaultRate is treated as "rule".  Arguably a defaultRate set via RC should be
+    // remote rule,
+    // but that's not currenlty part of the spec.
     if (defaultRate != null) {
       final RateSamplingRule samplingRule =
           new RateSamplingRule.AlwaysMatchesSamplingRule(
-              new DeterministicSampler.TraceSampler(defaultRate));
+              new DeterministicSampler.TraceSampler(defaultRate),
+              SamplingMechanism.LOCAL_USER_RULE);
       samplingRules.add(samplingRule);
     }
 
     return new RuleBasedTraceSampler(samplingRules, rateLimit, new RateByServiceTraceSampler());
   }
 
+  private static byte samplingMechanism(SamplingRule.Provenance provenance) {
+    switch (provenance) {
+      case DYNAMIC:
+        return SamplingMechanism.REMOTE_ADAPTIVE_RULE;
+      case CUSTOMER:
+        return SamplingMechanism.REMOTE_USER_RULE;
+      default:
+        return SamplingMechanism.LOCAL_USER_RULE;
+    }
+  }
+
   @Override
   public <T extends CoreSpan<T>> boolean sample(final T span) {
     return true;
@@ -137,21 +152,21 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
               PrioritySampling.USER_KEEP,
               SAMPLING_RULE_RATE,
               matchedRule.getSampler().getSampleRate(),
-              SamplingMechanism.RULE);
+              matchedRule.getMechanism());
         } else {
           span.setSamplingPriority(
               PrioritySampling.USER_DROP,
               SAMPLING_RULE_RATE,
               matchedRule.getSampler().getSampleRate(),
-              SamplingMechanism.RULE);
+              matchedRule.getMechanism());
         }
         span.setMetric(SAMPLING_LIMIT_RATE, rateLimit);
       } else {
         span.setSamplingPriority(
             PrioritySampling.USER_DROP,
             SAMPLING_RULE_RATE,
             matchedRule.getSampler().getSampleRate(),
-            SamplingMechanism.RULE);
+            matchedRule.getMechanism());
       }
     }
   }