Fix probability sampler and sampling. (#661)

Bogdan Drutu · web-flow · commit 23d498078cf3 · 2017-09-22T14:08:46.000-07:00
* Fix probability sampler and sampling.

* Add more tests and fix comments.
diff --git a/api/src/main/java/io/opencensus/trace/samplers/ProbabilitySampler.java b/api/src/main/java/io/opencensus/trace/samplers/ProbabilitySampler.java
@@ -56,7 +56,7 @@ abstract class ProbabilitySampler extends Sampler {
   static ProbabilitySampler create(double probability) {
     checkArgument(
         probability >= 0.0 && probability <= 1.0, "probability must be in range [0.0, 1.0]");
-    long idUpperBound = 0;
+    long idUpperBound;
     // Special case the limits, to avoid any possible issues with lack of precision across
     // double/long boundaries. For probability == 0.0, we use Long.MIN_VALUE as this guarantees
     // that we will never sample a trace, even in the case where the id == Long.MIN_VALUE, since
@@ -79,6 +79,16 @@ public final boolean shouldSample(
       SpanId spanId,
       String name,
       @Nullable List<Span> parentLinks) {
+    // If the parent is sampled keep the sampling decision.
+    if (parentContext != null && parentContext.getTraceOptions().isSampled()) {
+      return true;
+    }
+    // If any parent link is sampled keep the sampling decision.
+    for (Span parentLink : parentLinks) {
+      if (parentLink.getContext().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.
diff --git a/api/src/test/java/io/opencensus/trace/NoopSpan.java b/api/src/test/java/io/opencensus/trace/NoopSpan.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright 2017, OpenCensus Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opencensus.trace;
+
+import java.util.EnumSet;
+import java.util.Map;
+import javax.annotation.Nullable;
+
+/**
+ * Class to be used in tests where an implementation for the Span is needed.
+ *
+ * <p>Not final to allow Mockito to "spy" this class.
+ */
+public class NoopSpan extends Span {
+
+  /** Creates a new {@code NoopSpan}. */
+  public NoopSpan(SpanContext context, @Nullable EnumSet<Options> options) {
+    super(context, options);
+  }
+
+  @Override
+  public void putAttributes(Map<String, AttributeValue> attributes) {}
+
+  @Override
+  public void addAnnotation(String description, Map<String, AttributeValue> attributes) {}
+
+  @Override
+  public void addAnnotation(Annotation annotation) {}
+
+  @Override
+  public void addNetworkEvent(NetworkEvent networkEvent) {}
+
+  @Override
+  public void addLink(Link link) {}
+
+  @Override
+  public void end(EndSpanOptions options) {}
+}
diff --git a/api/src/test/java/io/opencensus/trace/SpanTest.java b/api/src/test/java/io/opencensus/trace/SpanTest.java
@@ -23,7 +23,6 @@
 
 import java.util.Collections;
 import java.util.EnumSet;
-import java.util.Map;
 import java.util.Random;
 import org.junit.Before;
 import org.junit.Test;
@@ -99,29 +98,4 @@ public void endCallsEndWithDefaultOptions() {
     span.end();
     verify(span).end(same(EndSpanOptions.DEFAULT));
   }
-
-  // No-op implementation of the Span for testing only.
-  private static class NoopSpan extends Span {
-    private NoopSpan(SpanContext context, EnumSet<Span.Options> options) {
-      super(context, options);
-    }
-
-    @Override
-    public void putAttributes(Map<String, AttributeValue> attributes) {}
-
-    @Override
-    public void addAnnotation(String description, Map<String, AttributeValue> attributes) {}
-
-    @Override
-    public void addAnnotation(Annotation annotation) {}
-
-    @Override
-    public void addNetworkEvent(NetworkEvent networkEvent) {}
-
-    @Override
-    public void addLink(Link link) {}
-
-    @Override
-    public void end(EndSpanOptions options) {}
-  }
 }
diff --git a/api/src/test/java/io/opencensus/trace/samplers/SamplersTest.java b/api/src/test/java/io/opencensus/trace/samplers/SamplersTest.java
@@ -18,13 +18,17 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import io.opencensus.trace.NoopSpan;
 import io.opencensus.trace.Sampler;
 import io.opencensus.trace.Span;
 import io.opencensus.trace.SpanContext;
 import io.opencensus.trace.SpanId;
 import io.opencensus.trace.TraceId;
 import io.opencensus.trace.TraceOptions;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.EnumSet;
+import java.util.List;
 import java.util.Random;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -43,6 +47,8 @@ public class SamplersTest {
       SpanContext.create(traceId, parentSpanId, TraceOptions.builder().setIsSampled(true).build());
   private final SpanContext notSampledSpanContext =
       SpanContext.create(traceId, parentSpanId, TraceOptions.DEFAULT);
+  private final Span sampledSpan =
+      new NoopSpan(sampledSpanContext, EnumSet.of(Span.Options.RECORD_EVENTS));
 
   @Test
   public void alwaysSampleSampler_AlwaysReturnTrue() {
@@ -118,7 +124,7 @@ public void probabilitySampler_outOfRangeLowProbability() {
 
   // Applies the given sampler to NUM_SAMPLE_TRIES random traceId/spanId pairs.
   private static void assertSamplerSamplesWithProbability(
-      Sampler sampler, SpanContext parent, double probability) {
+      Sampler sampler, SpanContext parent, List<Span> parentLinks, double probability) {
     Random random = new Random(1234);
     int count = 0; // Count of spans with sampling enabled
     for (int i = 0; i < NUM_SAMPLE_TRIES; i++) {
@@ -128,7 +134,7 @@ private static void assertSamplerSamplesWithProbability(
           TraceId.generateRandomId(random),
           SpanId.generateRandomId(random),
           SPAN_NAME,
-          Collections.<Span>emptyList())) {
+          parentLinks)) {
         count++;
       }
     }
@@ -139,17 +145,133 @@ private static void assertSamplerSamplesWithProbability(
   }
 
   @Test
-  public void probabilitySampler_differentProbabilities() {
+  public void probabilitySampler_DifferentProbabilities_NotSampledParent() {
     final Sampler neverSample = Samplers.probabilitySampler(0.0);
-    assertSamplerSamplesWithProbability(neverSample, sampledSpanContext, 0.0);
+    assertSamplerSamplesWithProbability(
+        neverSample, notSampledSpanContext, Collections.<Span>emptyList(), 0.0);
     final Sampler alwaysSample = Samplers.probabilitySampler(1.0);
-    assertSamplerSamplesWithProbability(alwaysSample, sampledSpanContext, 1.0);
+    assertSamplerSamplesWithProbability(
+        alwaysSample, notSampledSpanContext, Collections.<Span>emptyList(), 1.0);
     final Sampler fiftyPercentSample = Samplers.probabilitySampler(0.5);
-    assertSamplerSamplesWithProbability(fiftyPercentSample, sampledSpanContext, 0.5);
+    assertSamplerSamplesWithProbability(
+        fiftyPercentSample, notSampledSpanContext, Collections.<Span>emptyList(), 0.5);
     final Sampler twentyPercentSample = Samplers.probabilitySampler(0.2);
-    assertSamplerSamplesWithProbability(twentyPercentSample, sampledSpanContext, 0.2);
+    assertSamplerSamplesWithProbability(
+        twentyPercentSample, notSampledSpanContext, Collections.<Span>emptyList(), 0.2);
     final Sampler twoThirdsSample = Samplers.probabilitySampler(2.0 / 3.0);
-    assertSamplerSamplesWithProbability(twoThirdsSample, sampledSpanContext, 2.0 / 3.0);
+    assertSamplerSamplesWithProbability(
+        twoThirdsSample, notSampledSpanContext, Collections.<Span>emptyList(), 2.0 / 3.0);
+  }
+
+  @Test
+  public void probabilitySampler_DifferentProbabilities_SampledParent() {
+    final Sampler neverSample = Samplers.probabilitySampler(0.0);
+    assertSamplerSamplesWithProbability(
+        neverSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
+    final Sampler alwaysSample = Samplers.probabilitySampler(1.0);
+    assertSamplerSamplesWithProbability(
+        alwaysSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
+    final Sampler fiftyPercentSample = Samplers.probabilitySampler(0.5);
+    assertSamplerSamplesWithProbability(
+        fiftyPercentSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
+    final Sampler twentyPercentSample = Samplers.probabilitySampler(0.2);
+    assertSamplerSamplesWithProbability(
+        twentyPercentSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
+    final Sampler twoThirdsSample = Samplers.probabilitySampler(2.0 / 3.0);
+    assertSamplerSamplesWithProbability(
+        twoThirdsSample, sampledSpanContext, Collections.<Span>emptyList(), 1.0);
+  }
+
+  @Test
+  public void probabilitySampler_DifferentProbabilities_SampledParentLink() {
+    final Sampler neverSample = Samplers.probabilitySampler(0.0);
+    assertSamplerSamplesWithProbability(
+        neverSample, notSampledSpanContext, Arrays.asList(sampledSpan), 1.0);
+    final Sampler alwaysSample = Samplers.probabilitySampler(1.0);
+    assertSamplerSamplesWithProbability(
+        alwaysSample, notSampledSpanContext, Arrays.asList(sampledSpan), 1.0);
+    final Sampler fiftyPercentSample = Samplers.probabilitySampler(0.5);
+    assertSamplerSamplesWithProbability(
+        fiftyPercentSample, notSampledSpanContext, Arrays.asList(sampledSpan), 1.0);
+    final Sampler twentyPercentSample = Samplers.probabilitySampler(0.2);
+    assertSamplerSamplesWithProbability(
+        twentyPercentSample, notSampledSpanContext, Arrays.asList(sampledSpan), 1.0);
+    final Sampler twoThirdsSample = Samplers.probabilitySampler(2.0 / 3.0);
+    assertSamplerSamplesWithProbability(
+        twoThirdsSample, notSampledSpanContext, Arrays.asList(sampledSpan), 1.0);
+  }
+
+  @Test
+  public void probabilitySampler_SampleBasedOnTraceId() {
+    final Sampler defaultProbability = Samplers.probabilitySampler(0.0001);
+    // This traceId will not be sampled by the ProbabilitySampler because the first 8 bytes as long
+    // is not less than probability * Long.MAX_VALUE;
+    TraceId notSampledtraceId =
+        TraceId.fromBytes(
+            new byte[] {
+              (byte) 0x8F,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              0,
+              0,
+              0,
+              0,
+              0,
+              0,
+              0,
+              0
+            });
+    assertThat(
+            defaultProbability.shouldSample(
+                null,
+                false,
+                notSampledtraceId,
+                SpanId.generateRandomId(random),
+                SPAN_NAME,
+                Collections.<Span>emptyList()))
+        .isFalse();
+    // This traceId will be sampled by the ProbabilitySampler because the first 8 bytes as long
+    // is less than probability * Long.MAX_VALUE;
+    TraceId sampledtraceId =
+        TraceId.fromBytes(
+            new byte[] {
+              (byte) 0x00,
+              (byte) 0x00,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              (byte) 0xFF,
+              0,
+              0,
+              0,
+              0,
+              0,
+              0,
+              0,
+              0
+            });
+    assertThat(
+            defaultProbability.shouldSample(
+                null,
+                false,
+                sampledtraceId,
+                SpanId.generateRandomId(random),
+                SPAN_NAME,
+                Collections.<Span>emptyList()))
+        .isTrue();
+  }
+
+  @Test
+  public void probabilitySampler_getDescription() {
+    assertThat((Samplers.probabilitySampler(0.5)).getDescription())
+        .isEqualTo("ProbabilitySampler{0.500000}");
   }
 
   @Test
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
@@ -118,7 +118,16 @@ private static boolean makeSamplingDecision(
           .shouldSample(parent, hasRemoteParent, traceId, spanId, name, parentLinks);
     }
     // Parent is always different than null because otherwise we use the default sampler.
-    return parent.getTraceOptions().isSampled();
+    return parent.getTraceOptions().isSampled() || isAnyParentLinkSampled(parentLinks);
+  }
+
+  private static boolean isAnyParentLinkSampled(List<Span> parentLinks) {
+    for (Span parentLink : parentLinks) {
+      if (parentLink.getContext().getTraceOptions().isSampled()) {
+        return true;
+      }
+    }
+    return false;
   }
 
   private static void linkSpans(Span span, List<Span> parentLinks) {
diff --git a/impl_core/src/test/java/io/opencensus/implcore/trace/SpanBuilderImplTest.java b/impl_core/src/test/java/io/opencensus/implcore/trace/SpanBuilderImplTest.java
