Fix probability sampler and sampling. (#661)

* Fix probability sampler and sampling.

* Add more tests and fix comments.

3 files changed, 182 insertions, 34 deletions

diff --git a/api/src/test/java/io/opencensus/trace/NoopSpan.java b/api/src/test/java/io/opencensus/trace/NoopSpan.java
new file mode 100644
index 00000000..0217529d
--- /dev/null
+++ 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
index 1e05e60f..327b9abe 100644
--- 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 static org.mockito.Mockito.verify;
 
 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 class SpanTest {
     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
index 0999e8f8..6cb41975 100644
--- 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 @@ package io.opencensus.trace.samplers;
 
 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 class SamplersTest {
 
   // 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 @@ public class SamplersTest {
           TraceId.generateRandomId(random),
           SpanId.generateRandomId(random),
           SPAN_NAME,
-          Collections.<Span>emptyList())) {
+          parentLinks)) {
         count++;
       }
     }
@@ -139,17 +145,133 @@ public class SamplersTest {
   }
 
   @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