Don't attempt to replace the current SpanContext

In general, it turns out that replacing the current SpanContext is not a
good idea. Indeed, the OTel documentation [says][1]:

> Please note, since `SpanContext` is immutable, it is not possible to
> update `SpanContext` with a new `TraceState`. Such changes then make
> sense only right before `SpanContext` propagation or telemetry data
> exporting. In both cases, `Propagator`s and `SpanExporter`s may create
> a modified `TraceState` copy before serializing it to the wire.

So this commit updates how we set `TraceOptions` *during* a span.
Namely, when we update `TraceOptions` using `WithTraceOptions`, we no
longer modify the `SpanContext`. Instead, the value of `TraceOptions` is
stored on the `Context`, from where it can be retrieved by
`TraceOptionsPropagator` at propagation time.

This means that at propagation we will respect any `TraceOptions` set in
the `TraceState`, *unless* TraceOptions have been set in the `Context`,
in which case those will be used instead.

[1]: https://opentelemetry.io/docs/specs/otel/trace/api/#tracestate

Author: nickstenning

telemetry/propagators.go

@@ -0,0 +1,41 @@
+package telemetry
+
+import (
+	"context"
+
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
+)
+
+// Check TraceOptionsPropagator implements TextMapPropagator
+var _ propagation.TextMapPropagator = new(TraceOptionsPropagator)
+
+type TraceOptionsPropagator struct {
+	Next propagation.TextMapPropagator
+}
+
+func (p *TraceOptionsPropagator) Inject(ctx context.Context, carrier propagation.TextMapCarrier) {
+	sc := trace.SpanContextFromContext(ctx)
+	if !sc.IsValid() {
+		return
+	}
+
+	// If TraceOptions has been set directly in the context, then replace the
+	// SpanContext with one that has the appropriate TraceState.
+	//
+	// Note: it is generally only safe to do this in a propagator or an exporter.
+	if to, ok := traceOptionsFromContextOnly(ctx); ok {
+		ts := setTraceOptions(sc.TraceState(), to)
+		ctx = trace.ContextWithSpanContext(ctx, sc.WithTraceState(ts))
+	}
+
+	p.Next.Inject(ctx, carrier)
+}
+
+func (p *TraceOptionsPropagator) Extract(ctx context.Context, carrier propagation.TextMapCarrier) context.Context {
+	return p.Next.Extract(ctx, carrier)
+}
+
+func (p *TraceOptionsPropagator) Fields() []string {
+	return p.Next.Fields()
+}

telemetry/propagators_test.go

@@ -0,0 +1,95 @@
+package telemetry
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
+)
+
+func makeValidSpanContextConfig() trace.SpanContextConfig {
+	traceID, _ := trace.TraceIDFromHex("0123456789abcdef0123456789abcdef")
+	spanID, _ := trace.SpanIDFromHex("0123456789abcdef")
+	return trace.SpanContextConfig{
+		TraceID: traceID,
+		SpanID:  spanID,
+	}
+}
+
+func makeValidSpanContext() trace.SpanContext {
+	return trace.NewSpanContext(makeValidSpanContextConfig())
+}
+
+// Check that we're correctly passing the work onto the Next propagator.
+func TestTraceOptionsPropagatorUsesNextPropagator(t *testing.T) {
+	ctx := context.Background()
+	ctx = trace.ContextWithSpanContext(ctx, makeValidSpanContext())
+	propagator := TraceOptionsPropagator{
+		Next: propagation.TraceContext{},
+	}
+	carrier := propagation.MapCarrier{}
+
+	propagator.Inject(ctx, carrier)
+
+	require.Contains(t, carrier, "traceparent")
+}
+
+// Check that TraceOptions are respected both in SpanContext and
+// (preferentially) from the Context itself.
+func TestTraceOptionsPropagatorInjectsTraceOptions(t *testing.T) {
+	ctx := context.Background()
+
+	ts := trace.TraceState{}
+	ts, _ = ts.Insert("r8/sm", "always")
+
+	scc := makeValidSpanContextConfig()
+	scc.TraceState = ts
+	ctx = trace.ContextWithSpanContext(ctx, trace.NewSpanContext(scc))
+	propagator := TraceOptionsPropagator{
+		Next: propagation.TraceContext{},
+	}
+
+	// First check that only the sample mode field is set
+	{
+		carrier := propagation.MapCarrier{}
+		propagator.Inject(ctx, carrier)
+		require.Contains(t, carrier, "tracestate")
+		assert.Equal(t, carrier["tracestate"], "r8/sm=always")
+	}
+
+	// Then update TraceOptions locally and ensure that the values override those
+	// set in the SpanContext.
+	{
+		ctx := WithTraceOptions(ctx, TraceOptions{
+			DetailLevel: DetailLevelFull,
+			SampleMode:  SampleModeNever,
+		})
+		carrier := propagation.MapCarrier{}
+		propagator.Inject(ctx, carrier)
+		require.Contains(t, carrier, "tracestate")
+		assert.Contains(t, carrier["tracestate"], "r8/sm=never")
+		assert.Contains(t, carrier["tracestate"], "r8/dl=full")
+	}
+}
+
+func TestTraceOptionsPropagatorPrefersTraceOptionsFromContext(t *testing.T) {
+	ctx := trace.ContextWithSpanContext(context.Background(), makeValidSpanContext())
+	ctx = WithTraceOptions(ctx, TraceOptions{
+		DetailLevel: DetailLevelFull,
+		SampleMode:  SampleModeAlways,
+	})
+
+	propagator := TraceOptionsPropagator{
+		Next: propagation.TraceContext{},
+	}
+	carrier := propagation.MapCarrier{}
+
+	propagator.Inject(ctx, carrier)
+
+	require.Contains(t, carrier, "tracestate")
+	assert.Contains(t, carrier["tracestate"], "r8/sm=always")
+	assert.Contains(t, carrier["tracestate"], "r8/dl=full")
+}

telemetry/telemetry.go

@@ -37,10 +37,12 @@ func Start(ctx context.Context) (*Telemetry, error) {
 	otel.SetErrorHandler(ErrorHandler{})
 	otel.SetTracerProvider(tp)
 	otel.SetTextMapPropagator(
-		propagation.NewCompositeTextMapPropagator(
-			propagation.TraceContext{},
-			propagation.Baggage{},
-		),
+		&TraceOptionsPropagator{
+			Next: propagation.NewCompositeTextMapPropagator(
+				propagation.TraceContext{},
+				propagation.Baggage{},
+			),
+		},
 	)
 
 	return &Telemetry{tp}, nil

telemetry/tracestate.go

@@ -7,6 +7,10 @@ import (
 	"go.uber.org/zap"
 )
 
+type traceOptionsContextKeyT string
+
+const traceOptionsContextKey = traceOptionsContextKeyT("traceOptions")
+
 const (
 	TraceStateKeyDetailLevel = "r8/dl"
 	TraceStateKeySampleMode  = "r8/sm"
@@ -57,16 +61,20 @@ type TraceOptions struct {
 // TraceOptionsFromContext extracts any custom trace options from the trace
 // state carried in the passed context.
 func TraceOptionsFromContext(ctx context.Context) TraceOptions {
-	ts := trace.SpanContextFromContext(ctx).TraceState()
-	return parseTraceOptions(ts)
+	// First we see if any TraceOptions are set directly in the context. If so,
+	// they override any in the SpanContext TraceState.
+	if to, ok := traceOptionsFromContextOnly(ctx); ok {
+		return to
+	}
+
+	// Otherwise we fall back to using any TraceOptions set in the SpanContext.
+	return parseTraceOptions(trace.SpanContextFromContext(ctx).TraceState())
 }
 
 // WithTraceOptions returns a copy of the provided context with the passed
 // TraceOptions set.
 func WithTraceOptions(ctx context.Context, to TraceOptions) context.Context {
-	sc := trace.SpanContextFromContext(ctx)
-	ts := setTraceOptions(sc.TraceState(), to)
-	return trace.ContextWithSpanContext(ctx, sc.WithTraceState(ts))
+	return context.WithValue(ctx, traceOptionsContextKey, to)
 }
 
 // WithFullTrace returns a new context with full tracing mode enabled. This
@@ -78,6 +86,15 @@ func WithFullTrace(ctx context.Context) context.Context {
 	return WithTraceOptions(ctx, to)
 }
 
+func traceOptionsFromContextOnly(ctx context.Context) (TraceOptions, bool) {
+	if v := ctx.Value(traceOptionsContextKey); v != nil {
+		if to, ok := v.(TraceOptions); ok {
+			return to, true
+		}
+	}
+	return TraceOptions{}, false
+}
+
 func parseTraceOptions(ts trace.TraceState) TraceOptions {
 	to := TraceOptions{}
 

telemetry/tracestate_test.go

@@ -2,7 +2,6 @@ package telemetry
 
 import (
 	"context"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -32,17 +31,19 @@ func TestWithTraceOptionsSetsTraceOptions(t *testing.T) {
 	assert.Equal(t, SampleModeAlways, to.SampleMode)
 }
 
-func TestWithTraceOptionsSerialization(t *testing.T) {
+func TestTraceOptionsUsesSpanContextTraceOptions(t *testing.T) {
 	ctx := context.Background()
 
-	ctx = WithTraceOptions(ctx, TraceOptions{
-		DetailLevel: DetailLevelFull,
-		SampleMode:  SampleModeAlways,
-	})
+	ts := trace.TraceState{}
+	ts, _ = ts.Insert("r8/dl", "full")
+	ts, _ = ts.Insert("r8/sm", "always")
 
-	tsString := trace.SpanContextFromContext(ctx).TraceState().String()
-	elements := strings.Split(tsString, ",")
+	scc := makeValidSpanContextConfig()
+	scc.TraceState = ts
+	sc := trace.NewSpanContext(scc)
+	ctx = trace.ContextWithSpanContext(ctx, sc)
 
-	assert.Contains(t, elements, "r8/sm=always")
-	assert.Contains(t, elements, "r8/dl=full")
+	to := TraceOptionsFromContext(ctx)
+	assert.Equal(t, DetailLevelFull, to.DetailLevel)
+	assert.Equal(t, SampleModeAlways, to.SampleMode)
 }