Change when we check the sampler. (#657)

* Change when we check the sampler.

* Fix comments.

* Deprecate setIsSampled without an argument.

Author: Bogdan Drutu

api/src/main/java/io/opencensus/trace/TraceOptions.java

@@ -182,12 +182,26 @@ private Builder(byte options) {
     }
 
     /**
-     * Marks this trace as sampled.
-     *
+     * @deprecated Use {@code Builder.setIsSampled(true)}.
      * @return this.
      */
+    @Deprecated
     public Builder setIsSampled() {
-      options |= IS_SAMPLED;
+      return setIsSampled(true);
+    }
+
+    /**
+     * Sets the sampling bit in the options.
+     *
+     * @param isSampled the sampling bit.
+     * @return this.
+     */
+    public Builder setIsSampled(boolean isSampled) {
+      if (isSampled) {
+        options = (byte) (options | IS_SAMPLED);
+      } else {
+        options = (byte) (options & ~IS_SAMPLED);;
+      }
       return this;
     }
 

api/src/main/java/io/opencensus/trace/samplers/ProbabilitySampler.java

@@ -79,10 +79,6 @@ public final boolean shouldSample(
       SpanId spanId,
       String name,
       @Nullable List<Span> parentLinks) {
-    // Always enable sampling if parent was sampled.
-    if (parentContext != null && parentContext.getTraceOptions().isSampled()) {
-      return true;
-    }
     // Always sample if we are within probability range. This is true even for child spans (that
     // may have had a different sampling decision made) to allow for different sampling policies,
     // and dynamic increases to sampling probabilities for debugging purposes.

api/src/test/java/io/opencensus/trace/SpanContextTest.java

@@ -41,7 +41,7 @@ public class SpanContextTest {
       SpanContext.create(
           TraceId.fromBytes(secondTraceIdBytes),
           SpanId.fromBytes(secondSpanIdBytes),
-          TraceOptions.builder().setIsSampled().build());
+          TraceOptions.builder().setIsSampled(true).build());
 
   @Test
   public void invalidSpanContext() {
@@ -82,7 +82,8 @@ public void getSpanId() {
   @Test
   public void getTraceOptions() {
     assertThat(first.getTraceOptions()).isEqualTo(TraceOptions.DEFAULT);
-    assertThat(second.getTraceOptions()).isEqualTo(TraceOptions.builder().setIsSampled().build());
+    assertThat(second.getTraceOptions())
+        .isEqualTo(TraceOptions.builder().setIsSampled(true).build());
   }
 
   @Test
@@ -93,13 +94,17 @@ public void spanContext_EqualsAndHashCode() {
         SpanContext.create(
             TraceId.fromBytes(firstTraceIdBytes),
             SpanId.fromBytes(firstSpanIdBytes),
-            TraceOptions.DEFAULT));
+            TraceOptions.DEFAULT),
+        SpanContext.create(
+            TraceId.fromBytes(firstTraceIdBytes),
+            SpanId.fromBytes(firstSpanIdBytes),
+            TraceOptions.builder().setIsSampled(false).build()));
     tester.addEqualityGroup(
         second,
         SpanContext.create(
             TraceId.fromBytes(secondTraceIdBytes),
             SpanId.fromBytes(secondSpanIdBytes),
-            TraceOptions.builder().setIsSampled().build()));
+            TraceOptions.builder().setIsSampled(true).build()));
     tester.testEquals();
   }
 
@@ -111,6 +116,6 @@ public void spanContext_ToString() {
     assertThat(second.toString()).contains(TraceId.fromBytes(secondTraceIdBytes).toString());
     assertThat(second.toString()).contains(SpanId.fromBytes(secondSpanIdBytes).toString());
     assertThat(second.toString())
-        .contains(TraceOptions.builder().setIsSampled().build().toString());
+        .contains(TraceOptions.builder().setIsSampled(true).build().toString());
   }
 }

api/src/test/java/io/opencensus/trace/SpanTest.java

@@ -46,7 +46,7 @@ public void setUp() {
         SpanContext.create(
             TraceId.generateRandomId(random),
             SpanId.generateRandomId(random),
-            TraceOptions.builder().setIsSampled().build());
+            TraceOptions.builder().setIsSampled(true).build());
     notSampledSpanContext =
         SpanContext.create(
             TraceId.generateRandomId(random),

api/src/test/java/io/opencensus/trace/TraceOptionsTest.java

@@ -33,7 +33,10 @@ public class TraceOptionsTest {
   @Test
   public void getOptions() {
     assertThat(TraceOptions.DEFAULT.getOptions()).isEqualTo(0);
-    assertThat(TraceOptions.builder().setIsSampled().build().getOptions()).isEqualTo(1);
+    assertThat(TraceOptions.builder().setIsSampled(false).build().getOptions()).isEqualTo(0);
+    assertThat(TraceOptions.builder().setIsSampled(true).build().getOptions()).isEqualTo(1);
+    assertThat(TraceOptions.builder().setIsSampled(true).setIsSampled(false).build().getOptions())
+        .isEqualTo(0);
     assertThat(TraceOptions.fromBytes(firstBytes).getOptions()).isEqualTo(-1);
     assertThat(TraceOptions.fromBytes(secondBytes).getOptions()).isEqualTo(1);
     assertThat(TraceOptions.fromBytes(thirdBytes).getOptions()).isEqualTo(6);
@@ -42,7 +45,7 @@ public void getOptions() {
   @Test
   public void isSampled() {
     assertThat(TraceOptions.DEFAULT.isSampled()).isFalse();
-    assertThat(TraceOptions.builder().setIsSampled().build().isSampled()).isTrue();
+    assertThat(TraceOptions.builder().setIsSampled(true).build().isSampled()).isTrue();
   }
 
   @Test
@@ -56,7 +59,7 @@ public void toFromBytes() {
   public void builder_FromOptions() {
     assertThat(
             TraceOptions.builder(TraceOptions.fromBytes(thirdBytes))
-                .setIsSampled()
+                .setIsSampled(true)
                 .build()
                 .getOptions())
         .isEqualTo(6 | 1);
@@ -67,14 +70,15 @@ public void traceOptions_EqualsAndHashCode() {
     EqualsTester tester = new EqualsTester();
     tester.addEqualityGroup(TraceOptions.DEFAULT);
     tester.addEqualityGroup(
-        TraceOptions.fromBytes(secondBytes), TraceOptions.builder().setIsSampled().build());
+        TraceOptions.fromBytes(secondBytes), TraceOptions.builder().setIsSampled(true).build());
     tester.addEqualityGroup(TraceOptions.fromBytes(firstBytes));
     tester.testEquals();
   }
 
   @Test
   public void traceOptions_ToString() {
     assertThat(TraceOptions.DEFAULT.toString()).contains("sampled=false");
-    assertThat(TraceOptions.builder().setIsSampled().build().toString()).contains("sampled=true");
+    assertThat(TraceOptions.builder().setIsSampled(true).build().toString())
+        .contains("sampled=true");
   }
 }

api/src/test/java/io/opencensus/trace/samplers/SamplersTest.java

@@ -33,12 +33,14 @@
 /** Unit tests for {@link Samplers}. */
 @RunWith(JUnit4.class)
 public class SamplersTest {
+  private static final String SPAN_NAME = "MySpanName";
+  private static final int NUM_SAMPLE_TRIES = 1000;
   private final Random random = new Random(1234);
   private final TraceId traceId = TraceId.generateRandomId(random);
   private final SpanId parentSpanId = SpanId.generateRandomId(random);
   private final SpanId spanId = SpanId.generateRandomId(random);
   private final SpanContext sampledSpanContext =
-      SpanContext.create(traceId, parentSpanId, TraceOptions.builder().setIsSampled().build());
+      SpanContext.create(traceId, parentSpanId, TraceOptions.builder().setIsSampled(true).build());
   private final SpanContext notSampledSpanContext =
       SpanContext.create(traceId, parentSpanId, TraceOptions.DEFAULT);
 
@@ -114,59 +116,40 @@ public void probabilitySampler_outOfRangeLowProbability() {
     Samplers.probabilitySampler(-0.00001);
   }
 
-  private final void probabilitySampler_AlwaysReturnTrueForSampled(Sampler sampler) {
-    final int numSamples = 100; // Number of traces for which to generate sampling decisions.
-    for (int i = 0; i < numSamples; i++) {
-      assertThat(
-              sampler.shouldSample(
-                  sampledSpanContext,
-                  false,
-                  TraceId.generateRandomId(random),
-                  spanId,
-                  "bar",
-                  Collections.<Span>emptyList()))
-          .isTrue();
-    }
-  }
-
-  private final void probabilitySampler_SamplesWithProbabilityForUnsampled(
-      Sampler sampler, double probability) {
-    final int numSamples = 1000; // Number of traces for which to generate sampling decisions.
+  // Applies the given sampler to NUM_SAMPLE_TRIES random traceId/spanId pairs.
+  private static void assertSamplerSamplesWithProbability(
+      Sampler sampler, SpanContext parent, double probability) {
+    Random random = new Random(1234);
     int count = 0; // Count of spans with sampling enabled
-    for (int i = 0; i < numSamples; i++) {
+    for (int i = 0; i < NUM_SAMPLE_TRIES; i++) {
       if (sampler.shouldSample(
-          notSampledSpanContext,
+          parent,
           false,
           TraceId.generateRandomId(random),
-          spanId,
-          "bar",
+          SpanId.generateRandomId(random),
+          SPAN_NAME,
           Collections.<Span>emptyList())) {
         count++;
       }
     }
-    double proportionSampled = (double) count / numSamples;
+    double proportionSampled = (double) count / NUM_SAMPLE_TRIES;
     // Allow for a large amount of slop (+/- 10%) in number of sampled traces, to avoid flakiness.
     assertThat(proportionSampled < probability + 0.1 && proportionSampled > probability - 0.1)
         .isTrue();
   }
 
   @Test
-  public void probabilitySamper_SamplesWithProbability() {
+  public void probabilitySampler_differentProbabilities() {
     final Sampler neverSample = Samplers.probabilitySampler(0.0);
-    probabilitySampler_AlwaysReturnTrueForSampled(neverSample);
-    probabilitySampler_SamplesWithProbabilityForUnsampled(neverSample, 0.0);
+    assertSamplerSamplesWithProbability(neverSample, sampledSpanContext, 0.0);
     final Sampler alwaysSample = Samplers.probabilitySampler(1.0);
-    probabilitySampler_AlwaysReturnTrueForSampled(alwaysSample);
-    probabilitySampler_SamplesWithProbabilityForUnsampled(alwaysSample, 1.0);
+    assertSamplerSamplesWithProbability(alwaysSample, sampledSpanContext, 1.0);
     final Sampler fiftyPercentSample = Samplers.probabilitySampler(0.5);
-    probabilitySampler_AlwaysReturnTrueForSampled(fiftyPercentSample);
-    probabilitySampler_SamplesWithProbabilityForUnsampled(fiftyPercentSample, 0.5);
+    assertSamplerSamplesWithProbability(fiftyPercentSample, sampledSpanContext, 0.5);
     final Sampler twentyPercentSample = Samplers.probabilitySampler(0.2);
-    probabilitySampler_AlwaysReturnTrueForSampled(twentyPercentSample);
-    probabilitySampler_SamplesWithProbabilityForUnsampled(twentyPercentSample, 0.2);
+    assertSamplerSamplesWithProbability(twentyPercentSample, sampledSpanContext, 0.2);
     final Sampler twoThirdsSample = Samplers.probabilitySampler(2.0 / 3.0);
-    probabilitySampler_AlwaysReturnTrueForSampled(twoThirdsSample);
-    probabilitySampler_SamplesWithProbabilityForUnsampled(twoThirdsSample, 2.0 / 3.0);
+    assertSamplerSamplesWithProbability(twoThirdsSample, sampledSpanContext, 2.0 / 3.0);
   }
 
   @Test

impl_core/src/main/java/io/opencensus/implcore/trace/SpanBuilderImpl.java

@@ -75,12 +75,8 @@ private SpanImpl startSpanInternal(
       parentSpanId = parent.getSpanId();
       traceOptionsBuilder = TraceOptions.builder(parent.getTraceOptions());
     }
-    if (sampler == null) {
-      sampler = activeTraceParams.getSampler();
-    }
-    if (sampler.shouldSample(parent, hasRemoteParent, traceId, spanId, name, parentLinks)) {
-      traceOptionsBuilder.setIsSampled();
-    }
+    traceOptionsBuilder.setIsSampled(makeSamplingDecision(
+        parent, hasRemoteParent, name, sampler, parentLinks, traceId, spanId, activeTraceParams));
     TraceOptions traceOptions = traceOptionsBuilder.build();
     EnumSet<Span.Options> spanOptions = EnumSet.noneOf(Span.Options.class);
     if (traceOptions.isSampled() || Boolean.TRUE.equals(recordEvents)) {
@@ -101,6 +97,30 @@ private SpanImpl startSpanInternal(
     return span;
   }
 
+  private static boolean makeSamplingDecision(
+      @Nullable SpanContext parent,
+      @Nullable Boolean hasRemoteParent,
+      String name,
+      @Nullable Sampler sampler,
+      List<Span> parentLinks,
+      TraceId traceId,
+      SpanId spanId,
+      TraceParams activeTraceParams) {
+    // If users set a specific sampler in the SpanBuilder, use it.
+    if (sampler != null) {
+      return sampler.shouldSample(parent, hasRemoteParent, traceId, spanId, name, parentLinks);
+    }
+    // Use the default sampler if this is a root Span or this is an entry point Span (has remote
+    // parent).
+    if (Boolean.TRUE.equals(hasRemoteParent) || parent == null || !parent.isValid()) {
+      return activeTraceParams
+          .getSampler()
+          .shouldSample(parent, hasRemoteParent, traceId, spanId, name, parentLinks);
+    }
+    // Parent is always different than null because otherwise we use the default sampler.
+    return parent.getTraceOptions().isSampled();
+  }
+
   private static void linkSpans(Span span, List<Span> parentLinks) {
     if (!parentLinks.isEmpty()) {
       Link childLink = Link.fromSpanContext(span.getContext(), Type.CHILD_LINKED_SPAN);

impl_core/src/test/java/io/opencensus/implcore/trace/SpanBuilderImplTest.java

@@ -107,16 +107,14 @@ public void startChildSpan() {
     assertThat(rootSpan.getContext().isValid()).isTrue();
     assertThat(rootSpan.getOptions().contains(Options.RECORD_EVENTS)).isTrue();
     assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isTrue();
-    assertThat(((SpanImpl) rootSpan).toSpanData().getHasRemoteParent())
-        .isNull();
+    assertThat(((SpanImpl) rootSpan).toSpanData().getHasRemoteParent()).isNull();
     Span childSpan =
         SpanBuilderImpl.createWithParent(SPAN_NAME, rootSpan, spanBuilderOptions).startSpan();
     assertThat(childSpan.getContext().isValid()).isTrue();
     assertThat(childSpan.getContext().getTraceId()).isEqualTo(rootSpan.getContext().getTraceId());
     assertThat(((SpanImpl) childSpan).toSpanData().getParentSpanId())
         .isEqualTo(rootSpan.getContext().getSpanId());
-    assertThat(((SpanImpl) childSpan).toSpanData().getHasRemoteParent())
-        .isFalse();
+    assertThat(((SpanImpl) childSpan).toSpanData().getHasRemoteParent()).isFalse();
     assertThat(((SpanImpl) childSpan).getTimestampConverter())
         .isEqualTo(((SpanImpl) rootSpan).getTimestampConverter());
   }
@@ -164,6 +162,95 @@ public void startRemoteSpan() {
     assertThat(spanData.getHasRemoteParent()).isTrue();
   }
 
+  @Test
+  public void startRootSpan_WithSpecifiedSampler() {
+    // Apply given sampler before default sampler for root spans.
+    Span rootSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
+            .setSampler(Samplers.neverSample())
+            .startSpan();
+    assertThat(rootSpan.getContext().isValid()).isTrue();
+    assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isFalse();
+  }
+
+  @Test
+  public void startRootSpan_WithoutSpecifiedSampler() {
+    // Apply default sampler (always true in the tests) for root spans.
+    Span rootSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions).startSpan();
+    assertThat(rootSpan.getContext().isValid()).isTrue();
+    assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isTrue();
+  }
+
+  @Test
+  public void startRemoteChildSpan_WithSpecifiedSampler() {
+    Span rootSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
+            .setSampler(Samplers.alwaysSample())
+            .startSpan();
+    assertThat(rootSpan.getContext().isValid()).isTrue();
+    assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isTrue();
+    // Apply given sampler before default sampler for spans with remote parent.
+    Span childSpan =
+        SpanBuilderImpl.createWithRemoteParent(SPAN_NAME, rootSpan.getContext(), spanBuilderOptions)
+            .setSampler(Samplers.neverSample())
+            .startSpan();
+    assertThat(childSpan.getContext().isValid()).isTrue();
+    assertThat(childSpan.getContext().getTraceId()).isEqualTo(rootSpan.getContext().getTraceId());
+    assertThat(childSpan.getContext().getTraceOptions().isSampled()).isFalse();
+  }
+
+  @Test
+  public void startRemoteChildSpan_WithoutSpecifiedSampler() {
+    Span rootSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
+            .setSampler(Samplers.neverSample())
+            .startSpan();
+    assertThat(rootSpan.getContext().isValid()).isTrue();
+    assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isFalse();
+    // Apply default sampler (always true in the tests) for spans with remote parent.
+    Span childSpan =
+        SpanBuilderImpl.createWithRemoteParent(SPAN_NAME, rootSpan.getContext(), spanBuilderOptions)
+            .startSpan();
+    assertThat(childSpan.getContext().isValid()).isTrue();
+    assertThat(childSpan.getContext().getTraceId()).isEqualTo(rootSpan.getContext().getTraceId());
+    assertThat(childSpan.getContext().getTraceOptions().isSampled()).isTrue();
+  }
+
+  @Test
+  public void startChildSpan_WithSpecifiedSampler() {
+    Span rootSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
+            .setSampler(Samplers.alwaysSample())
+            .startSpan();
+    assertThat(rootSpan.getContext().isValid()).isTrue();
+    assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isTrue();
+    // Apply the given sampler for child spans.
+    Span childSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, rootSpan, spanBuilderOptions)
+            .setSampler(Samplers.neverSample())
+            .startSpan();
+    assertThat(childSpan.getContext().isValid()).isTrue();
+    assertThat(childSpan.getContext().getTraceId()).isEqualTo(rootSpan.getContext().getTraceId());
+    assertThat(childSpan.getContext().getTraceOptions().isSampled()).isFalse();
+  }
+
+  @Test
+  public void startChildSpan_WithoutSpecifiedSampler() {
+    Span rootSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, null, spanBuilderOptions)
+            .setSampler(Samplers.neverSample())
+            .startSpan();
+    assertThat(rootSpan.getContext().isValid()).isTrue();
+    assertThat(rootSpan.getContext().getTraceOptions().isSampled()).isFalse();
+    // Don't apply the default sampler (always true) for child spans.
+    Span childSpan =
+        SpanBuilderImpl.createWithParent(SPAN_NAME, rootSpan, spanBuilderOptions).startSpan();
+    assertThat(childSpan.getContext().isValid()).isTrue();
+    assertThat(childSpan.getContext().getTraceId()).isEqualTo(rootSpan.getContext().getTraceId());
+    assertThat(childSpan.getContext().getTraceOptions().isSampled()).isFalse();
+  }
+
   private static final class FakeRandomHandler extends RandomHandler {
     private final Random random;