tracer: always propagate the tracestate header (#2339)

Always forwards the tracestate from other vendors, regardless of whether the w3c trace context is during extraction.

ddtrace/tracer/textmap.go

@@ -15,6 +15,7 @@ import (
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/samplernames"
 )
@@ -155,11 +156,12 @@ func NewPropagator(cfg *PropagatorConfig, propagators ...Propagator) Propagator
 	if cfg.PriorityHeader == "" {
 		cfg.PriorityHeader = DefaultPriorityHeader
 	}
+	cp := new(chainedPropagator)
+	cp.onlyExtractFirst = internal.BoolEnv("DD_TRACE_PROPAGATION_EXTRACT_FIRST", false)
 	if len(propagators) > 0 {
-		return &chainedPropagator{
-			injectors:  propagators,
-			extractors: propagators,
-		}
+		cp.injectors = propagators
+		cp.extractors = propagators
+		return cp
 	}
 	injectorsPs := os.Getenv(headerPropagationStyleInject)
 	if injectorsPs == "" {
@@ -173,24 +175,20 @@ func NewPropagator(cfg *PropagatorConfig, propagators ...Propagator) Propagator
 			log.Warn("%v is deprecated. Please use %v or %v instead.\n", headerPropagationStyleExtractDeprecated, headerPropagationStyleExtract, headerPropagationStyle)
 		}
 	}
-	injectors, injectorNames := getPropagators(cfg, injectorsPs)
-	extractors, extractorsNames := getPropagators(cfg, extractorsPs)
-	return &chainedPropagator{
-		injectors,
-		extractors,
-		injectorNames,
-		extractorsNames,
-	}
+	cp.injectors, cp.injectorNames = getPropagators(cfg, injectorsPs)
+	cp.extractors, cp.extractorsNames = getPropagators(cfg, extractorsPs)
+	return cp
 }
 
 // chainedPropagator implements Propagator and applies a list of injectors and extractors.
 // When injecting, all injectors are called to propagate the span context.
 // When extracting, it tries each extractor, selecting the first successful one.
 type chainedPropagator struct {
-	injectors       []Propagator
-	extractors      []Propagator
-	injectorNames   string
-	extractorsNames string
+	injectors        []Propagator
+	extractors       []Propagator
+	injectorNames    string
+	extractorsNames  string
+	onlyExtractFirst bool // value of DD_TRACE_PROPAGATION_EXTRACT_FIRST
 }
 
 // getPropagators returns a list of propagators based on ps, which is a comma seperated
@@ -265,21 +263,70 @@ func (p *chainedPropagator) Inject(spanCtx ddtrace.SpanContext, carrier interfac
 	return nil
 }
 
-// Extract implements Propagator.
+// Extract implements Propagator. This method will attempt to extract the context
+// based on the precedence order of the propagators. Generally, the first valid
+// trace context that could be extracted will be returned, and other extractors will
+// be ignored. However, the W3C tracestate header value will always be extracted and
+// stored in the local trace context even if a previous propagator has already succeeded
+// so long as the trace-ids match.
 func (p *chainedPropagator) Extract(carrier interface{}) (ddtrace.SpanContext, error) {
+	var ctx ddtrace.SpanContext
 	for _, v := range p.extractors {
-		ctx, err := v.Extract(carrier)
 		if ctx != nil {
-			// first extractor returns
-			log.Debug("Extracted span context: %#v", ctx)
-			return ctx, nil
+			// A local trace context has already been extracted.
+			p, isW3C := v.(*propagatorW3c)
+			if !isW3C {
+				continue // Ignore other propagators.
+			}
+			p.propagateTracestate(ctx.(*spanContext), carrier)
+			break
 		}
-		if err == ErrSpanContextNotFound {
-			continue
+		var err error
+		ctx, err = v.Extract(carrier)
+		if ctx != nil {
+			if p.onlyExtractFirst {
+				// Return early if the customer configured that only the first successful
+				// extraction should occur.
+				return ctx, nil
+			}
+		} else if err != ErrSpanContextNotFound {
+			return nil, err
 		}
-		return nil, err
 	}
-	return nil, ErrSpanContextNotFound
+	if ctx == nil {
+		return nil, ErrSpanContextNotFound
+	}
+	log.Debug("Extracted span context: %#v", ctx)
+	return ctx, nil
+}
+
+// propagateTracestate will add the tracestate propagating tag to the given
+// *spanContext. The W3C trace context will be extracted from the provided
+// carrier. The trace id of this W3C trace context must match the trace id
+// provided by the given *spanContext. If it matches, then the tracestate
+// will be re-composed based on the composition of the given *spanContext,
+// but will include the non-DD vendors in the W3C trace context's tracestate.
+func (p *propagatorW3c) propagateTracestate(ctx *spanContext, carrier interface{}) {
+	w3cCtx, _ := p.Extract(carrier)
+	if w3cCtx == nil {
+		return // It's not valid, so ignore it.
+	}
+	if ctx.TraceID() != w3cCtx.TraceID() {
+		return // The trace-ids must match.
+	}
+	if w3cCtx.(*spanContext).trace == nil {
+		return // this shouldn't happen, since it should have a propagating tag already
+	}
+	if ctx.trace == nil {
+		ctx.trace = newTrace()
+	}
+	// Get the tracestate header from extracted w3C context, and propagate
+	// it to the span context that will be returned.
+	// Note: Other trace context fields like sampling priority, propagated tags,
+	// and origin will remain unchanged.
+	ts := w3cCtx.(*spanContext).trace.propagatingTag(tracestateHeader)
+	priority, _ := ctx.SamplingPriority()
+	setPropagatingTag(ctx, tracestateHeader, composeTracestate(ctx, priority, ts))
 }
 
 // propagator implements Propagator and injects/extracts span contexts

ddtrace/tracer/textmap_test.go

@@ -107,6 +107,117 @@ func TestTextMapCarrierForeachKeyError(t *testing.T) {
 	assert.Equal(t, got, want)
 }
 
+func TestTextMapExtractTracestatePropagation(t *testing.T) {
+	tests := []struct {
+		name, propagationStyle, traceparent string
+		onlyExtractFirst                    bool // value of DD_TRACE_PROPAGATION_EXTRACT_FIRST
+		wantTracestatePropagation           bool
+	}{
+		{
+			/*
+				With only Datadog propagation set, the tracestate header should
+				be ignored, and not propagated to the returned trace context.
+			*/
+			name:             "datadog-only",
+			propagationStyle: "datadog",
+			traceparent:      "00-00000000000000000000000000000004-2222222222222222-01",
+		},
+		{
+			/*
+				With Datadog, B3, AND w3c propagation set, the tracestate header should
+				be propagated to the returned trace context. This test also verifies that
+				b3 extraction doesn't override the local context value.
+			*/
+			name:                      "datadog-b3-w3c",
+			propagationStyle:          "datadog,b3,tracecontext",
+			traceparent:               "00-00000000000000000000000000000004-2222222222222222-01",
+			wantTracestatePropagation: true,
+		},
+		{
+			/*
+				With Datadog AND w3c propagation set, the tracestate header should
+				be propagated to the returned trace context.
+			*/
+			name:                      "datadog-and-w3c",
+			propagationStyle:          "datadog,tracecontext",
+			traceparent:               "00-00000000000000000000000000000004-2222222222222222-01",
+			wantTracestatePropagation: true,
+		},
+		{
+			/*
+				With Datadog AND w3c propagation set, but mismatching trace-ids,
+				the tracestate header should be ignored and not propagated to
+				the returned trace context.
+			*/
+			name:             "datadog-and-w3c-mismatching-ids",
+			propagationStyle: "datadog,tracecontext",
+			traceparent:      "00-00000000000000000000000000000088-2222222222222222-01",
+		},
+		{
+			/*
+				With Datadog AND w3c propagation set, but the traceparent is malformed,
+				the tracestate header should be ignored and not propagated to
+				the returned trace context.
+			*/
+			name:             "datadog-and-w3c-malformed",
+			propagationStyle: "datadog,tracecontext",
+			traceparent:      "00-00000000000000000000000000000004-22asdf!2-01",
+		},
+		{
+			/*
+				With Datadog AND w3c propagation set, but there is no traceparent,
+				the tracestate header should be ignored and not propagated to
+				the returned trace context.
+			*/
+			name:             "datadog-and-w3c-no-traceparent",
+			propagationStyle: "datadog,tracecontext",
+		},
+		{
+			/*
+				With Datadog AND w3c propagation set, but DD_TRACE_PROPAGATION_EXTRACT_FIRST
+				is true, the tracestate header should be ignored and not propagated to
+				the returned trace context.
+			*/
+			name:             "datadog-and-w3c-only-extract-first",
+			propagationStyle: "datadog,tracecontext",
+			traceparent:      "00-00000000000000000000000000000004-2222222222222222-01",
+			onlyExtractFirst: true,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(fmt.Sprintf("TestTextMapExtractTracestatePropagation-%s", tc.name), func(t *testing.T) {
+			t.Setenv(headerPropagationStyle, tc.propagationStyle)
+			if tc.onlyExtractFirst {
+				os.Setenv("DD_TRACE_PROPAGATION_EXTRACT_FIRST", "true")
+				defer os.Unsetenv("DD_TRACE_PROPAGATION_EXTRACT_FIRST")
+			}
+			tracer := newTracer()
+			assert := assert.New(t)
+			headers := TextMapCarrier(map[string]string{
+				DefaultTraceIDHeader:  "4",
+				DefaultParentIDHeader: "1",
+				originHeader:          "synthetics",
+				b3TraceIDHeader:       "0021dc1807524785",
+				traceparentHeader:     tc.traceparent,
+				tracestateHeader:      "dd=s:2;o:rum;t.tid:1230000000000000~~,othervendor=t61rcWkgMzE",
+			})
+
+			ctx, err := tracer.Extract(headers)
+			assert.Nil(err)
+			sctx, ok := ctx.(*spanContext)
+			assert.True(ok)
+			assert.Equal("00000000000000000000000000000004", sctx.traceID.HexEncoded())
+			assert.Equal(uint64(1), sctx.spanID)    // should use x-datadog-parent-id, not the id in the tracestate
+			assert.Equal("synthetics", sctx.origin) // should use x-datadog-origin, not the origin in the tracestate
+			if tc.wantTracestatePropagation {
+				assert.Equal("dd=s:0;o:synthetics,othervendor=t61rcWkgMzE", sctx.trace.propagatingTag(tracestateHeader))
+			} else if sctx.trace != nil {
+				assert.False(sctx.trace.hasPropagatingTag(tracestateHeader))
+			}
+		})
+	}
+}
+
 func TestTextMapPropagatorErrors(t *testing.T) {
 	t.Setenv(headerPropagationStyleExtract, "datadog")
 	propagator := NewPropagator(nil)