From e251d3c95f3227a6999d77026ef924b637f6d1a8 Mon Sep 17 00:00:00 2001 From: Nick Stenning Date: Fri, 17 Nov 2023 13:20:11 +0100 Subject: [PATCH] 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 --- telemetry/propagators.go | 41 +++++++++++++++ telemetry/propagators_test.go | 95 +++++++++++++++++++++++++++++++++++ telemetry/telemetry.go | 10 ++-- telemetry/tracestate.go | 27 ++++++++-- telemetry/tracestate_test.go | 21 ++++---- 5 files changed, 175 insertions(+), 19 deletions(-) create mode 100644 telemetry/propagators.go create mode 100644 telemetry/propagators_test.go diff --git a/telemetry/propagators.go b/telemetry/propagators.go new file mode 100644 index 0000000..325f9f7 --- /dev/null +++ b/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() +} diff --git a/telemetry/propagators_test.go b/telemetry/propagators_test.go new file mode 100644 index 0000000..e5b7f44 --- /dev/null +++ b/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") +} diff --git a/telemetry/telemetry.go b/telemetry/telemetry.go index 90e3e98..e267ffa 100644 --- a/telemetry/telemetry.go +++ b/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 diff --git a/telemetry/tracestate.go b/telemetry/tracestate.go index 3fb1d83..5d5116a 100644 --- a/telemetry/tracestate.go +++ b/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{} diff --git a/telemetry/tracestate_test.go b/telemetry/tracestate_test.go index b0d5ef9..0ab9e1a 100644 --- a/telemetry/tracestate_test.go +++ b/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) }