Upgraded DataDog tracing library to 1.14.0

vendor/gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer/textmap.go

@@ -2,10 +2,12 @@ package tracer
 
 import (
 	"net/http"
+	"os"
 	"strconv"
 	"strings"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
 )
 
 // HTTPHeadersCarrier wraps an http.Header as a TextMapWriter and TextMapReader, allowing
@@ -54,6 +56,11 @@ func (c TextMapCarrier) ForeachKey(handler func(key, val string) error) error {
 	return nil
 }
 
+const (
+	headerPropagationStyleInject  = "DD_PROPAGATION_STYLE_INJECT"
+	headerPropagationStyleExtract = "DD_PROPAGATION_STYLE_EXTRACT"
+)
+
 const (
 	// DefaultBaggageHeaderPrefix specifies the prefix that will be used in
 	// HTTP headers or text maps to prefix baggage keys.
@@ -72,6 +79,10 @@ const (
 	DefaultPriorityHeader = "x-datadog-sampling-priority"
 )
 
+// originHeader specifies the name of the header indicating the origin of the trace.
+// It is used with the Synthetics product and usually has the value "synthetics".
+const originHeader = "x-datadog-origin"
+
 // PropagatorConfig defines the configuration for initializing a propagator.
 type PropagatorConfig struct {
 	// BaggagePrefix specifies the prefix that will be used to store baggage
@@ -110,21 +121,86 @@ func NewPropagator(cfg *PropagatorConfig) Propagator {
 	if cfg.PriorityHeader == "" {
 		cfg.PriorityHeader = DefaultPriorityHeader
 	}
-	return &propagator{cfg}
+	return &chainedPropagator{
+		injectors:  getPropagators(cfg, headerPropagationStyleInject),
+		extractors: getPropagators(cfg, headerPropagationStyleExtract),
+	}
 }
 
-// propagator implements a propagator which uses TextMap internally.
-// It propagates the trace and span IDs, as well as the baggage from the
-// context.
-type propagator struct{ cfg *PropagatorConfig }
+// 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
+}
+
+// getPropagators returns a list of propagators based on the list found in the
+// given environment variable. If the list doesn't contain a value or has invalid
+// values, the default propagator will be returned.
+func getPropagators(cfg *PropagatorConfig, env string) []Propagator {
+	dd := &propagator{cfg}
+	ps := os.Getenv(env)
+	if ps == "" {
+		return []Propagator{dd}
+	}
+	var list []Propagator
+	for _, v := range strings.Split(ps, ",") {
+		switch strings.ToLower(v) {
+		case "datadog":
+			list = append(list, dd)
+		case "b3":
+			list = append(list, &propagatorB3{})
+		default:
+			// TODO(cgilmour): consider logging something for invalid/unknown styles.
+		}
+	}
+	if len(list) == 0 {
+		// return the default
+		return []Propagator{dd}
+	}
+	return list
+}
 
 // Inject defines the Propagator to propagate SpanContext data
 // out of the current process. The implementation propagates the
 // TraceID and the current active SpanID, as well as the Span baggage.
+func (p *chainedPropagator) Inject(spanCtx ddtrace.SpanContext, carrier interface{}) error {
+	for _, v := range p.injectors {
+		err := v.Inject(spanCtx, carrier)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// Extract implements Propagator.
+func (p *chainedPropagator) Extract(carrier interface{}) (ddtrace.SpanContext, error) {
+	for _, v := range p.extractors {
+		ctx, err := v.Extract(carrier)
+		if ctx != nil {
+			// first extractor returns
+			return ctx, nil
+		}
+		if err == ErrSpanContextNotFound {
+			continue
+		}
+		return nil, err
+	}
+	return nil, ErrSpanContextNotFound
+}
+
+// propagator implements Propagator and injects/extracts span contexts
+// using datadog headers. Only TextMap carriers are supported.
+type propagator struct {
+	cfg *PropagatorConfig
+}
+
 func (p *propagator) Inject(spanCtx ddtrace.SpanContext, carrier interface{}) error {
-	switch v := carrier.(type) {
+	switch c := carrier.(type) {
 	case TextMapWriter:
-		return p.injectTextMap(spanCtx, v)
+		return p.injectTextMap(spanCtx, c)
 	default:
 		return ErrInvalidCarrier
 	}
@@ -141,6 +217,9 @@ func (p *propagator) injectTextMap(spanCtx ddtrace.SpanContext, writer TextMapWr
 	if ctx.hasSamplingPriority() {
 		writer.Set(p.cfg.PriorityHeader, strconv.Itoa(ctx.samplingPriority()))
 	}
+	if ctx.origin != "" {
+		writer.Set(originHeader, ctx.origin)
+	}
 	// propagate OpenTracing baggage
 	for k, v := range ctx.baggage {
 		writer.Set(p.cfg.BaggagePrefix+k, v)
@@ -148,11 +227,10 @@ func (p *propagator) injectTextMap(spanCtx ddtrace.SpanContext, writer TextMapWr
 	return nil
 }
 
-// Extract implements Propagator.
 func (p *propagator) Extract(carrier interface{}) (ddtrace.SpanContext, error) {
-	switch v := carrier.(type) {
+	switch c := carrier.(type) {
 	case TextMapReader:
-		return p.extractTextMap(v)
+		return p.extractTextMap(c)
 	default:
 		return nil, ErrInvalidCarrier
 	}
@@ -180,6 +258,8 @@ func (p *propagator) extractTextMap(reader TextMapReader) (ddtrace.SpanContext,
 				return ErrSpanContextCorrupted
 			}
 			ctx.setSamplingPriority(priority)
+		case originHeader:
+			ctx.origin = v
 		default:
 			if strings.HasPrefix(key, p.cfg.BaggagePrefix) {
 				ctx.setBaggageItem(strings.TrimPrefix(key, p.cfg.BaggagePrefix), v)
@@ -195,3 +275,83 @@ func (p *propagator) extractTextMap(reader TextMapReader) (ddtrace.SpanContext,
 	}
 	return &ctx, nil
 }
+
+const (
+	b3TraceIDHeader = "x-b3-traceid"
+	b3SpanIDHeader  = "x-b3-spanid"
+	b3SampledHeader = "x-b3-sampled"
+)
+
+// propagatorB3 implements Propagator and injects/extracts span contexts
+// using B3 headers. Only TextMap carriers are supported.
+type propagatorB3 struct{}
+
+func (p *propagatorB3) Inject(spanCtx ddtrace.SpanContext, carrier interface{}) error {
+	switch c := carrier.(type) {
+	case TextMapWriter:
+		return p.injectTextMap(spanCtx, c)
+	default:
+		return ErrInvalidCarrier
+	}
+}
+
+func (*propagatorB3) injectTextMap(spanCtx ddtrace.SpanContext, writer TextMapWriter) error {
+	ctx, ok := spanCtx.(*spanContext)
+	if !ok || ctx.traceID == 0 || ctx.spanID == 0 {
+		return ErrInvalidSpanContext
+	}
+	writer.Set(b3TraceIDHeader, strconv.FormatUint(ctx.traceID, 16))
+	writer.Set(b3SpanIDHeader, strconv.FormatUint(ctx.spanID, 16))
+	if ctx.hasSamplingPriority() {
+		if ctx.samplingPriority() >= ext.PriorityAutoKeep {
+			writer.Set(b3SampledHeader, "1")
+		} else {
+			writer.Set(b3SampledHeader, "0")
+		}
+	}
+	return nil
+}
+
+func (p *propagatorB3) Extract(carrier interface{}) (ddtrace.SpanContext, error) {
+	switch c := carrier.(type) {
+	case TextMapReader:
+		return p.extractTextMap(c)
+	default:
+		return nil, ErrInvalidCarrier
+	}
+}
+
+func (*propagatorB3) extractTextMap(reader TextMapReader) (ddtrace.SpanContext, error) {
+	var ctx spanContext
+	err := reader.ForeachKey(func(k, v string) error {
+		var err error
+		key := strings.ToLower(k)
+		switch key {
+		case b3TraceIDHeader:
+			ctx.traceID, err = strconv.ParseUint(v, 16, 64)
+			if err != nil {
+				return ErrSpanContextCorrupted
+			}
+		case b3SpanIDHeader:
+			ctx.spanID, err = strconv.ParseUint(v, 16, 64)
+			if err != nil {
+				return ErrSpanContextCorrupted
+			}
+		case b3SampledHeader:
+			priority, err := strconv.Atoi(v)
+			if err != nil {
+				return ErrSpanContextCorrupted
+			}
+			ctx.setSamplingPriority(priority)
+		default:
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	if ctx.traceID == 0 || ctx.spanID == 0 {
+		return nil, ErrSpanContextNotFound
+	}
+	return &ctx, nil
+}