Change when we check the sampler. (#657)

* Change when we check the sampler.

* Fix comments.

* Deprecate setIsSampled without an argument.

1 files changed, 26 insertions, 6 deletions

diff --git a/impl_core/src/main/java/io/opencensus/implcore/trace/SpanBuilderImpl.java b/impl_core/src/main/java/io/opencensus/implcore/trace/SpanBuilderImpl.java
index b66a354f..3cee1def 100644
--- a/impl_core/src/main/java/io/opencensus/implcore/trace/SpanBuilderImpl.java
+++ b/impl_core/src/main/java/io/opencensus/implcore/trace/SpanBuilderImpl.java
@@ -75,12 +75,8 @@ final class SpanBuilderImpl extends SpanBuilder {
       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 @@ final class SpanBuilderImpl extends SpanBuilder {
     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);