sdk/trace: Test the Sampling Probability (#247)

* sdk/trace: Test the Sampling Probability

Closes #163

* Add tests for when the parent is sampled.

.gitignore

@@ -1,5 +1,6 @@
 .tools/
 .idea/
+.vscode/
 *.iml
 *.so
 coverage.*

sdk/trace/provider.go

@@ -85,6 +85,8 @@ func NewProvider(opts ...ProviderOption) (*Provider, error) {
 	return tp, nil
 }
 
+// GetTracer with the given name. If a tracer for the given name does not exist,
+// it is created first. If the name is empty, DefaultTracerName is used.
 func (p *Provider) GetTracer(name string) apitrace.Tracer {
 	p.mu.Lock()
 	defer p.mu.Unlock()

sdk/trace/sampling.go

@@ -39,16 +39,18 @@ type SamplingDecision struct {
 	Sample bool
 }
 
-// ProbabilitySampler returns a Sampler that samples a given fraction of traces.
-//
-// It also samples spans whose parents are sampled.
+// ProbabilitySampler samples a given fraction of traces. Fractions >= 1 will
+// always sample. If the parent span is sampled, then it's child spans will
+// automatically be sampled. Fractions <0 are treated as zero, but spans may
+// still be sampled if their parent is.
 func ProbabilitySampler(fraction float64) Sampler {
-	if !(fraction >= 0) {
-		fraction = 0
-	} else if fraction >= 1 {
+	if fraction >= 1 {
 		return AlwaysSample()
 	}
 
+	if fraction <= 0 {
+		fraction = 0
+	}
 	traceIDUpperBound := uint64(fraction * (1 << 63))
 	return Sampler(func(p SamplingParameters) SamplingDecision {
 		if p.ParentContext.IsSampled() {

sdk/trace/trace_test.go

@@ -17,6 +17,7 @@ package trace
 import (
 	"context"
 	"fmt"
+	"math"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -134,7 +135,71 @@ func TestRecordingIsOff(t *testing.T) {
 	}
 }
 
-// TODO: [rghetia] enable sampling test when Sampling is working.
+func TestSampling(t *testing.T) {
+	idg := defIDGenerator()
+	total := 10000
+	for name, tc := range map[string]struct {
+		sampler       Sampler
+		expect        float64
+		tolerance     float64
+		parent        bool
+		sampledParent bool
+	}{
+		// Span w/o a parent
+		"NeverSample":            {sampler: NeverSample(), expect: 0, tolerance: 0},
+		"AlwaysSample":           {sampler: AlwaysSample(), expect: 1.0, tolerance: 0},
+		"ProbabilitySampler_-1":  {sampler: ProbabilitySampler(-1.0), expect: 0, tolerance: 0},
+		"ProbabilitySampler_.25": {sampler: ProbabilitySampler(0.25), expect: .25, tolerance: 0.015},
+		"ProbabilitySampler_.50": {sampler: ProbabilitySampler(0.50), expect: .5, tolerance: 0.015},
+		"ProbabilitySampler_.75": {sampler: ProbabilitySampler(0.75), expect: .75, tolerance: 0.015},
+		"ProbabilitySampler_2.0": {sampler: ProbabilitySampler(2.0), expect: 1, tolerance: 0},
+		// Spans with a parent that is *not* sampled act like spans w/o a parent
+		"UnsampledParentSpanWithProbabilitySampler_-1":  {sampler: ProbabilitySampler(-1.0), expect: 0, tolerance: 0, parent: true},
+		"UnsampledParentSpanWithProbabilitySampler_.25": {sampler: ProbabilitySampler(.25), expect: .25, tolerance: 0.015, parent: true},
+		"UnsampledParentSpanWithProbabilitySampler_.50": {sampler: ProbabilitySampler(0.50), expect: .5, tolerance: 0.015, parent: true},
+		"UnsampledParentSpanWithProbabilitySampler_.75": {sampler: ProbabilitySampler(0.75), expect: .75, tolerance: 0.015, parent: true},
+		"UnsampledParentSpanWithProbabilitySampler_2.0": {sampler: ProbabilitySampler(2.0), expect: 1, tolerance: 0, parent: true},
+		// Spans with a parent that is sampled, will always sample, regardless of the probability
+		"SampledParentSpanWithProbabilitySampler_-1":  {sampler: ProbabilitySampler(-1.0), expect: 1, tolerance: 0, parent: true, sampledParent: true},
+		"SampledParentSpanWithProbabilitySampler_.25": {sampler: ProbabilitySampler(.25), expect: 1, tolerance: 0, parent: true, sampledParent: true},
+		"SampledParentSpanWithProbabilitySampler_2.0": {sampler: ProbabilitySampler(2.0), expect: 1, tolerance: 0, parent: true, sampledParent: true},
+		// Spans with a sampled parent, but when using the NeverSample Sampler, aren't sampled
+		"SampledParentSpanWithNeverSample": {sampler: NeverSample(), expect: 0, tolerance: 0, parent: true, sampledParent: true},
+	} {
+		tc := tc
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			p, err := NewProvider(WithConfig(Config{DefaultSampler: tc.sampler}))
+			if err != nil {
+				t.Fatal("unexpected error:", err)
+			}
+			tr := p.GetTracer("test")
+			var sampled int
+			for i := 0; i < total; i++ {
+				var opts []apitrace.SpanOption
+				if tc.parent {
+					psc := core.SpanContext{
+						TraceID: idg.NewTraceID(),
+						SpanID:  idg.NewSpanID(),
+					}
+					if tc.sampledParent {
+						psc.TraceFlags = core.TraceFlagsSampled
+					}
+					opts = append(opts, apitrace.ChildOf(psc))
+				}
+				_, span := tr.Start(context.Background(), "test", opts...)
+				if span.SpanContext().IsSampled() {
+					sampled++
+				}
+			}
+			got := float64(sampled) / float64(total)
+			diff := math.Abs(got - tc.expect)
+			if diff > tc.tolerance {
+				t.Errorf("got %f (diff: %f), expected %f (w/tolerance: %f)", got, diff, tc.expect, tc.tolerance)
+			}
+		})
+	}
+}
 
 func TestStartSpanWithChildOf(t *testing.T) {
 	tp, _ := NewProvider()