Add tests for propagation of Sampler Tracestate changes (#1655)

* Add tests for propagation of Sampler Tracestate changes

Sampler specification indicates that SamplingResult.Tracestate
should be associated with the SpanContext of the newly created span.
See
https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/sdk.md#sampler

* Fix SamplingResult TraceState propagation

* Add Changelog entry

Co-authored-by: Tyler Yahn <MrAlias@users.noreply.github.com>

CHANGELOG.md

@@ -16,6 +16,11 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 - Removed the exported `SimpleSpanProcessor` and `BatchSpanProcessor` structs.
    These are now returned as a SpanProcessor interface from their respective constructors. (#1638)
 
+### Fixed
+
+- `SamplingResult.TraceState` is correctly propagated to a newly created
+  span's `SpanContext`. (#1655)
+
 ## [0.18.0] - 2020-03-03
 
 ### Added
@@ -41,7 +46,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 ### Changed
 
 - Replaced interface `oteltest.SpanRecorder` with its existing implementation
-  `StandardSpanRecorder` (#1542).
+  `StandardSpanRecorder`. (#1542)
 - Default span limit values to 128. (#1535)
 - Rename `MaxEventsPerSpan`, `MaxAttributesPerSpan` and `MaxLinksPerSpan` to `EventCountLimit`, `AttributeCountLimit` and `LinkCountLimit`, and move these fields into `SpanLimits`. (#1535)
 - Renamed the `otel/label` package to `otel/attribute`. (#1541)

sdk/trace/span.go

@@ -543,6 +543,7 @@ func startSpanInternal(ctx context.Context, tr *tracer, name string, parent trac
 		kind:         o.SpanKind,
 	}
 	sampled := makeSamplingDecision(data)
+	span.spanContext.TraceState = sampled.Tracestate
 
 	if !span.spanContext.IsSampled() && !o.Record {
 		return span

sdk/trace/trace_test.go

@@ -58,12 +58,22 @@ var (
 	sid trace.SpanID
 
 	handler *storingHandler = &storingHandler{}
+
+	k1, k2, k3    attribute.Key
+	kv1, kv2, kv3 attribute.KeyValue
 )
 
 func init() {
 	tid, _ = trace.TraceIDFromHex("01020304050607080102040810203040")
 	sid, _ = trace.SpanIDFromHex("0102040810203040")
 
+	k1 = attribute.Key("k1")
+	kv1 = k1.String("v1")
+	k2 = attribute.Key("k2")
+	kv2 = k2.String("v2")
+	k3 = attribute.Key("k3")
+	kv3 = k3.String("v3")
+
 	otel.SetErrorHandler(handler)
 }
 
@@ -1526,3 +1536,145 @@ func TestAddLinksWithMoreAttributesThanLimit(t *testing.T) {
 		t.Errorf("Link: -got +want %s", diff)
 	}
 }
+
+type stateSampler struct {
+	prefix string
+	f      func(trace.TraceState) trace.TraceState
+}
+
+func (s *stateSampler) ShouldSample(p SamplingParameters) SamplingResult {
+	decision := Drop
+	if strings.HasPrefix(p.Name, s.prefix) {
+		decision = RecordAndSample
+	}
+	return SamplingResult{Decision: decision, Tracestate: s.f(p.ParentContext.TraceState)}
+}
+
+func (s stateSampler) Description() string {
+	return "stateSampler"
+}
+
+// Check that a new span propagates the SamplerResult.TraceState
+func TestSamplerTraceState(t *testing.T) {
+	mustTS := func(t trace.TraceState, err error) trace.TraceState { return t }
+	makeInserter := func(k attribute.KeyValue, prefix string) Sampler {
+		return &stateSampler{
+			prefix: prefix,
+			f:      func(t trace.TraceState) trace.TraceState { return mustTS(t.Insert(k)) },
+		}
+	}
+	makeDeleter := func(k attribute.Key, prefix string) Sampler {
+		return &stateSampler{
+			prefix: prefix,
+			f:      func(t trace.TraceState) trace.TraceState { return mustTS(t.Delete(k)) },
+		}
+	}
+	clearer := func(prefix string) Sampler {
+		return &stateSampler{
+			prefix: prefix,
+			f:      func(t trace.TraceState) trace.TraceState { return trace.TraceState{} },
+		}
+	}
+
+	tests := []struct {
+		name       string
+		sampler    Sampler
+		spanName   string
+		input      trace.TraceState
+		want       trace.TraceState
+		exportSpan bool
+	}{
+		{
+			name:       "alwaysOn",
+			sampler:    AlwaysSample(),
+			input:      mustTS(trace.TraceStateFromKeyValues(kv1)),
+			want:       mustTS(trace.TraceStateFromKeyValues(kv1)),
+			exportSpan: true,
+		},
+		{
+			name:       "alwaysOff",
+			sampler:    NeverSample(),
+			input:      mustTS(trace.TraceStateFromKeyValues(kv1)),
+			want:       mustTS(trace.TraceStateFromKeyValues(kv1)),
+			exportSpan: false,
+		},
+		{
+			name:       "insertKeySampled",
+			sampler:    makeInserter(kv2, "span"),
+			spanName:   "span0",
+			input:      mustTS(trace.TraceStateFromKeyValues(kv1)),
+			want:       mustTS(trace.TraceStateFromKeyValues(kv2, kv1)),
+			exportSpan: true,
+		},
+		{
+			name:       "insertKeyDropped",
+			sampler:    makeInserter(kv2, "span"),
+			spanName:   "nospan0",
+			input:      mustTS(trace.TraceStateFromKeyValues(kv1)),
+			want:       mustTS(trace.TraceStateFromKeyValues(kv2, kv1)),
+			exportSpan: false,
+		},
+		{
+			name:       "deleteKeySampled",
+			sampler:    makeDeleter(k1, "span"),
+			spanName:   "span0",
+			input:      mustTS(trace.TraceStateFromKeyValues(kv1, kv2)),
+			want:       mustTS(trace.TraceStateFromKeyValues(kv2)),
+			exportSpan: true,
+		},
+		{
+			name:       "deleteKeyDropped",
+			sampler:    makeDeleter(k1, "span"),
+			spanName:   "nospan0",
+			input:      mustTS(trace.TraceStateFromKeyValues(kv1, kv2, kv3)),
+			want:       mustTS(trace.TraceStateFromKeyValues(kv2, kv3)),
+			exportSpan: false,
+		},
+		{
+			name:       "clearer",
+			sampler:    clearer("span"),
+			spanName:   "span0",
+			input:      mustTS(trace.TraceStateFromKeyValues(kv1, kv3)),
+			want:       mustTS(trace.TraceStateFromKeyValues()),
+			exportSpan: true,
+		},
+	}
+
+	for _, ts := range tests {
+		ts := ts
+		t.Run(ts.name, func(t *testing.T) {
+			te := NewTestExporter()
+			tp := NewTracerProvider(WithConfig(Config{DefaultSampler: ts.sampler}), WithSyncer(te), WithResource(resource.Empty()))
+			tr := tp.Tracer("TraceState")
+
+			sc1 := trace.SpanContext{
+				TraceID:    tid,
+				SpanID:     sid,
+				TraceFlags: trace.FlagsSampled,
+				TraceState: ts.input,
+			}
+			ctx := trace.ContextWithRemoteSpanContext(context.Background(), sc1)
+			_, span := tr.Start(ctx, ts.spanName)
+
+			// span's TraceState should be set regardless of Sampled/NonSampled state.
+			require.Equal(t, ts.want, span.SpanContext().TraceState)
+
+			span.End()
+
+			got := te.Spans()
+			if len(got) > 0 != ts.exportSpan {
+				t.Errorf("unexpected number of exported spans %d", len(got))
+			}
+			if len(got) == 0 {
+				return
+			}
+
+			receivedState := got[0].SpanContext.TraceState
+
+			if diff := cmpDiff(receivedState, ts.want); diff != "" {
+				t.Errorf("TraceState not propagated: -got +want %s", diff)
+			}
+		})
+	}
+
+}