Update Samplers to conform to Spec (#531)

* Refactor SDK Sampler API to conform to Spec

* Sampler is now an interface rather than a function type
* SamplingParameters include the span Kind, Attributes, and Links
* SamplingResult includes a SamplingDecision with three possible values, as well as Attributes

* Add attributes retruned from a Sampler to the span

* Add SpanKind, Attributes, and Links to API Sampler.ShouldSample() parameters

* Drop "Get" from sdk Sampler.GetDescription to match api Sampler

* Make spanID parameter in API Sampler interface a core.SpanID

* Fix types and printf format per PR feedback from krnowak

* Ensure unit test error messages reflect new reality

Co-authored-by: Joshua MacDonald <[email protected]>

api/trace/always_sampler.go

@@ -33,8 +33,11 @@ func (as alwaysSampleSampler) ShouldSample(
 	_ core.SpanContext,
 	_ bool,
 	_ core.TraceID,
-	_ uint64,
+	_ core.SpanID,
 	_ string,
+	_ SpanKind,
+	_ []core.KeyValue,
+	_ []Link,
 ) Decision {
 	return alwaysSampleDecision
 }

api/trace/always_sampler_test.go

@@ -24,7 +24,7 @@ import (
 
 func TestShouldSample(t *testing.T) {
 	gotD := AlwaysSampleSampler().ShouldSample(
-		core.SpanContext{}, false, core.TraceID{}, 0, "span")
+		core.SpanContext{}, false, core.TraceID{}, core.SpanID{}, "span", SpanKindClient, []core.KeyValue{}, []Link{})
 	wantD := Decision{Sampled: true}
 	if diff := cmp.Diff(wantD, gotD); diff != "" {
 		t.Errorf("Decision: +got, -want%v", diff)

api/trace/never_sampler.go

@@ -33,8 +33,11 @@ func (ns neverSampleSampler) ShouldSample(
 	_ core.SpanContext,
 	_ bool,
 	_ core.TraceID,
-	_ uint64,
+	_ core.SpanID,
 	_ string,
+	_ SpanKind,
+	_ []core.KeyValue,
+	_ []Link,
 ) Decision {
 	return neverSampledecision
 }

api/trace/never_sampler_test.go

@@ -24,7 +24,7 @@ import (
 
 func TestNeverSamperShouldSample(t *testing.T) {
 	gotD := NeverSampleSampler().ShouldSample(
-		core.SpanContext{}, false, core.TraceID{}, 0, "span")
+		core.SpanContext{}, false, core.TraceID{}, core.SpanID{}, "span", SpanKindClient, []core.KeyValue{}, []Link{})
 	wantD := Decision{Sampled: false}
 	if diff := cmp.Diff(wantD, gotD); diff != "" {
 		t.Errorf("Decision: +got, -want%v", diff)

api/trace/sampler.go

@@ -24,8 +24,11 @@ type Sampler interface {
 		sc core.SpanContext,
 		remote bool,
 		traceID core.TraceID,
-		spanID uint64,
+		spanID core.SpanID,
 		spanName string,
+		spanKind SpanKind,
+		attributes []core.KeyValue,
+		links []Link,
 	) Decision
 
 	// Description returns of the sampler. It contains its name or short description

sdk/trace/sampling.go

@@ -16,12 +16,17 @@ package trace
 
 import (
 	"encoding/binary"
+	"fmt"
 
 	"go.opentelemetry.io/otel/api/core"
+	api "go.opentelemetry.io/otel/api/trace"
 )
 
 // Sampler decides whether a trace should be sampled and exported.
-type Sampler func(SamplingParameters) SamplingDecision
+type Sampler interface {
+	ShouldSample(SamplingParameters) SamplingResult
+	Description() string
+}
 
 // SamplingParameters contains the values passed to a Sampler.
 type SamplingParameters struct {
@@ -30,11 +35,46 @@ type SamplingParameters struct {
 	SpanID          core.SpanID
 	Name            string
 	HasRemoteParent bool
+	Kind            api.SpanKind
+	Attributes      []core.KeyValue
+	Links           []api.Link
+}
+
+// SamplingDecision indicates whether a span is recorded and sampled.
+type SamplingDecision uint8
+
+// Valid sampling decisions
+const (
+	NotRecord SamplingDecision = iota
+	Record
+	RecordAndSampled
+)
+
+// SamplingResult conveys a SamplingDecision and a set of Attributes.
+type SamplingResult struct {
+	Decision   SamplingDecision
+	Attributes []core.KeyValue
+}
+
+type probabilitySampler struct {
+	traceIDUpperBound uint64
+	description       string
 }
 
-// SamplingDecision is the value returned by a Sampler.
-type SamplingDecision struct {
-	Sample bool
+func (ps probabilitySampler) ShouldSample(p SamplingParameters) SamplingResult {
+	if p.ParentContext.IsSampled() {
+		return SamplingResult{Decision: RecordAndSampled}
+	}
+
+	x := binary.BigEndian.Uint64(p.TraceID[0:8]) >> 1
+	if x < ps.traceIDUpperBound {
+		return SamplingResult{Decision: RecordAndSampled}
+	}
+	return SamplingResult{Decision: NotRecord}
+}
+
+func (ps probabilitySampler) Description() string {
+	return ps.description
 }
 
 // ProbabilitySampler samples a given fraction of traces. Fractions >= 1 will
@@ -49,37 +89,53 @@ func ProbabilitySampler(fraction float64) Sampler {
 	if fraction <= 0 {
 		fraction = 0
 	}
-	traceIDUpperBound := uint64(fraction * (1 << 63))
-	return func(p SamplingParameters) SamplingDecision {
-		if p.ParentContext.IsSampled() {
-			return SamplingDecision{Sample: true}
-		}
-		x := binary.BigEndian.Uint64(p.TraceID[0:8]) >> 1
-		return SamplingDecision{Sample: x < traceIDUpperBound}
+
+	return &probabilitySampler{
+		traceIDUpperBound: uint64(fraction * (1 << 63)),
+		description:       fmt.Sprintf("ProbabilitySampler{%g}", fraction),
 	}
 }
 
+type alwaysOnSampler struct{}
+
+func (as alwaysOnSampler) ShouldSample(p SamplingParameters) SamplingResult {
+	return SamplingResult{Decision: RecordAndSampled}
+}
+
+func (as alwaysOnSampler) Description() string {
+	return "AlwaysOnSampler"
+}
+
 // AlwaysSample returns a Sampler that samples every trace.
 // Be careful about using this sampler in a production application with
 // significant traffic: a new trace will be started and exported for every
 // request.
 func AlwaysSample() Sampler {
-	return func(p SamplingParameters) SamplingDecision {
-		return SamplingDecision{Sample: true}
-	}
+	return alwaysOnSampler{}
+}
+
+type alwaysOffSampler struct{}
+
+func (as alwaysOffSampler) ShouldSample(p SamplingParameters) SamplingResult {
+	return SamplingResult{Decision: NotRecord}
+}
+
+func (as alwaysOffSampler) Description() string {
+	return "AlwaysOffSampler"
 }
 
 // NeverSample returns a Sampler that samples no traces.
 func NeverSample() Sampler {
-	return func(p SamplingParameters) SamplingDecision {
-		return SamplingDecision{Sample: false}
-	}
+	return alwaysOffSampler{}
 }
 
 // AlwaysParentSample returns a Sampler that samples a trace only
 // if the parent span is sampled.
 // This Sampler is a passthrough to the ProbabilitySampler with
 // a fraction of value 0.
 func AlwaysParentSample() Sampler {
-	return ProbabilitySampler(0)
+	return &probabilitySampler{
+		traceIDUpperBound: 0,
+		description:       "AlwaysParentSampler",
+	}
 }

sdk/trace/sampling_test.go

@@ -30,8 +30,8 @@ func TestAlwaysParentSampleWithParentSampled(t *testing.T) {
 		SpanID:     spanID,
 		TraceFlags: core.TraceFlagsSampled,
 	}
-	if !sampler(sdktrace.SamplingParameters{ParentContext: parentCtx}).Sample {
-		t.Error("Sampling decision should be true")
+	if sampler.ShouldSample(sdktrace.SamplingParameters{ParentContext: parentCtx}).Decision != sdktrace.RecordAndSampled {
+		t.Error("Sampling decision should be RecordAndSampled")
 	}
 }
 
@@ -43,7 +43,7 @@ func TestAlwaysParentSampleWithParentNotSampled(t *testing.T) {
 		TraceID: traceID,
 		SpanID:  spanID,
 	}
-	if sampler(sdktrace.SamplingParameters{ParentContext: parentCtx}).Sample {
-		t.Error("Sampling decision should be false")
+	if sampler.ShouldSample(sdktrace.SamplingParameters{ParentContext: parentCtx}).Decision != sdktrace.NotRecord {
+		t.Error("Sampling decision should be NotRecord")
 	}
 }

sdk/trace/span.go

@@ -222,8 +222,17 @@ func (s *span) SetName(name string) {
 		name:         name,
 		cfg:          s.tracer.provider.config.Load().(*Config),
 		span:         s,
+		attributes:   s.data.Attributes,
+		links:        s.data.Links,
+		kind:         s.data.SpanKind,
+	}
+	sampled := makeSamplingDecision(data)
+
+	// Adding attributes directly rather than using s.SetAttributes()
+	// as s.mu is already locked and attempting to do so would deadlock.
+	for _, a := range sampled.Attributes {
+		s.attributes.add(a)
 	}
-	makeSamplingDecision(data)
 }
 
 func (s *span) addLink(link apitrace.Link) {
@@ -308,8 +317,11 @@ func startSpanInternal(tr *tracer, name string, parent core.SpanContext, remoteP
 		name:         name,
 		cfg:          cfg,
 		span:         span,
+		attributes:   o.Attributes,
+		links:        o.Links,
+		kind:         o.SpanKind,
 	}
-	makeSamplingDecision(data)
+	sampled := makeSamplingDecision(data)
 
 	// TODO: [rghetia] restore when spanstore is added.
 	// if !internal.LocalSpanStoreEnabled && !span.spanContext.IsSampled() && !o.Record {
@@ -332,6 +344,8 @@ func startSpanInternal(tr *tracer, name string, parent core.SpanContext, remoteP
 	span.messageEvents = newEvictedQueue(cfg.MaxEventsPerSpan)
 	span.links = newEvictedQueue(cfg.MaxLinksPerSpan)
 
+	span.SetAttributes(sampled.Attributes...)
+
 	if !noParent {
 		span.data.ParentSpanID = parent.SpanID
 	}
@@ -354,9 +368,12 @@ type samplingData struct {
 	name         string
 	cfg          *Config
 	span         *span
+	attributes   []core.KeyValue
+	links        []apitrace.Link
+	kind         apitrace.SpanKind
 }
 
-func makeSamplingDecision(data samplingData) {
+func makeSamplingDecision(data samplingData) SamplingResult {
 	if data.noParent || data.remoteParent {
 		// If this span is the child of a local span and no
 		// Sampler is set in the options, keep the parent's
@@ -369,16 +386,25 @@ func makeSamplingDecision(data samplingData) {
 		//	sampler = o.Sampler
 		//}
 		spanContext := &data.span.spanContext
-		sampled := sampler(SamplingParameters{
+		sampled := sampler.ShouldSample(SamplingParameters{
 			ParentContext:   data.parent,
 			TraceID:         spanContext.TraceID,
 			SpanID:          spanContext.SpanID,
 			Name:            data.name,
-			HasRemoteParent: data.remoteParent}).Sample
-		if sampled {
+			HasRemoteParent: data.remoteParent,
+			Kind:            data.kind,
+			Attributes:      data.attributes,
+			Links:           data.links,
+		})
+		if sampled.Decision == RecordAndSampled {
 			spanContext.TraceFlags |= core.TraceFlagsSampled
 		} else {
 			spanContext.TraceFlags &^= core.TraceFlagsSampled
 		}
+		return sampled
+	}
+	if data.parent.TraceFlags&core.TraceFlagsSampled != 0 {
+		return SamplingResult{Decision: RecordAndSampled}
 	}
+	return SamplingResult{Decision: NotRecord}
 }

sdk/trace/trace_test.go

@@ -67,13 +67,28 @@ func (t *testExporter) ExportSpan(ctx context.Context, d *export.SpanData) {
 	t.spans = append(t.spans, d)
 }
 
+type testSampler struct {
+	callCount int
+	prefix    string
+	t         *testing.T
+}
+
+func (ts *testSampler) ShouldSample(p SamplingParameters) SamplingResult {
+	ts.callCount++
+	ts.t.Logf("called sampler for name %q", p.Name)
+	decision := NotRecord
+	if strings.HasPrefix(p.Name, ts.prefix) {
+		decision = RecordAndSampled
+	}
+	return SamplingResult{Decision: decision, Attributes: []core.KeyValue{core.Key("callCount").Int(ts.callCount)}}
+}
+
+func (ts testSampler) Description() string {
+	return "testSampler"
+}
+
 func TestSetName(t *testing.T) {
-	samplerIsCalled := false
-	fooSampler := Sampler(func(p SamplingParameters) SamplingDecision {
-		samplerIsCalled = true
-		t.Logf("called sampler for name %q", p.Name)
-		return SamplingDecision{Sample: strings.HasPrefix(p.Name, "foo")}
-	})
+	fooSampler := &testSampler{prefix: "foo", t: t}
 	tp, _ := NewProvider(WithConfig(Config{DefaultSampler: fooSampler}))
 
 	type testCase struct {
@@ -109,18 +124,18 @@ func TestSetName(t *testing.T) {
 		},
 	} {
 		span := startNamedSpan(tp, "SetName", tt.name)
-		if !samplerIsCalled {
+		if fooSampler.callCount == 0 {
 			t.Errorf("%d: the sampler was not even called during span creation", idx)
 		}
-		samplerIsCalled = false
+		fooSampler.callCount = 0
 		if gotSampledBefore := span.SpanContext().IsSampled(); tt.sampledBefore != gotSampledBefore {
 			t.Errorf("%d: invalid sampling decision before rename, expected %v, got %v", idx, tt.sampledBefore, gotSampledBefore)
 		}
 		span.SetName(tt.newName)
-		if !samplerIsCalled {
+		if fooSampler.callCount == 0 {
 			t.Errorf("%d: the sampler was not even called during span rename", idx)
 		}
-		samplerIsCalled = false
+		fooSampler.callCount = 0
 		if gotSampledAfter := span.SpanContext().IsSampled(); tt.sampledAfter != gotSampledAfter {
 			t.Errorf("%d: invalid sampling decision after rename, expected %v, got %v", idx, tt.sampledAfter, gotSampledAfter)
 		}