Added Lineage Support to XRay Propagator using W3CBaggage

propagators/aws/xray/propagator.go

@@ -2,12 +2,13 @@
 // SPDX-License-Identifier: Apache-2.0
 
 package xray // import "go.opentelemetry.io/contrib/propagators/aws/xray"
-
 import (
 	"context"
 	"errors"
+	"strconv"
 	"strings"
 
+	"go.opentelemetry.io/otel/baggage"
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -16,9 +17,11 @@ const (
 	traceHeaderKey       = "X-Amzn-Trace-Id"
 	traceHeaderDelimiter = ";"
 	kvDelimiter          = "="
+	lineageDelimiter     = ":"
 	traceIDKey           = "Root"
 	sampleFlagKey        = "Sampled"
 	parentIDKey          = "Parent"
+	lineageKey           = "Lineage"
 	traceIDVersion       = "1"
 	traceIDDelimiter     = "-"
 	isSampled            = "1"
@@ -31,6 +34,14 @@ const (
 	traceIDDelimitterIndex2 = 10
 	traceIDFirstPartLength  = 8
 	sampledFlagLength       = 1
+
+	lineageMaxLength   = 18
+	lineageMinLength   = 12
+	lineageHashLength  = 8
+	lineageMaxCounter1 = 32767
+	lineageMaxCounter2 = 255
+	lineageMinCounter  = 0
+	invalidLineage     = "-1:11111111:0"
 )
 
 var (
@@ -66,34 +77,45 @@ func (xray Propagator) Inject(ctx context.Context, carrier propagation.TextMapCa
 	if sc.TraceFlags() == traceFlagSampled {
 		samplingFlag = isSampled
 	}
+
 	headers := []string{
 		traceIDKey, kvDelimiter, xrayTraceID, traceHeaderDelimiter, parentIDKey,
 		kvDelimiter, parentID.String(), traceHeaderDelimiter, sampleFlagKey, kvDelimiter, samplingFlag,
 	}
 
-	carrier.Set(traceHeaderKey, strings.Join(headers, ""))
+	contextBaggage := baggage.FromContext(ctx)
+	lineage := contextBaggage.Member(lineageKey)
+
+	if lineage.Key() != "" {
+		headers = append(headers, traceHeaderDelimiter, lineageKey, kvDelimiter, lineage.Value())
+	}
+
+	carrier.Set(traceHeaderKey, strings.Join(headers, "")[:256])
 }
 
 // Extract gets a context from the carrier if it contains AWS X-Ray headers.
 func (xray Propagator) Extract(ctx context.Context, carrier propagation.TextMapCarrier) context.Context {
 	// extract tracing information
 	if header := carrier.Get(traceHeaderKey); header != "" {
-		sc, err := extract(header)
+		newContext, sc, err := extract(ctx, header)
 		if err == nil && sc.IsValid() {
-			return trace.ContextWithRemoteSpanContext(ctx, sc)
+			return trace.ContextWithRemoteSpanContext(newContext, sc)
 		}
 	}
 	return ctx
 }
 
 // extract extracts Span Context from context.
-func extract(headerVal string) (trace.SpanContext, error) {
+func extract(ctx context.Context, headerVal string) (context.Context, trace.SpanContext, error) {
 	var (
 		scc            = trace.SpanContextConfig{}
 		err            error
 		delimiterIndex int
 		part           string
 	)
+
+	currBaggage := baggage.FromContext(ctx)
+
 	pos := 0
 	for pos < len(headerVal) {
 		delimiterIndex = indexOf(headerVal, traceHeaderDelimiter, pos)
@@ -107,26 +129,38 @@ func extract(headerVal string) (trace.SpanContext, error) {
 		}
 		equalsIndex := strings.Index(part, kvDelimiter)
 		if equalsIndex < 0 {
-			return empty, errInvalidTraceHeader
+			return ctx, empty, errInvalidTraceHeader
 		}
 		value := part[equalsIndex+1:]
 		if strings.HasPrefix(part, traceIDKey) {
 			scc.TraceID, err = parseTraceID(value)
 			if err != nil {
-				return empty, err
+				return ctx, empty, err
 			}
 		} else if strings.HasPrefix(part, parentIDKey) {
 			// extract parentId
 			scc.SpanID, err = trace.SpanIDFromHex(value)
 			if err != nil {
-				return empty, errInvalidSpanIDLength
+				return ctx, empty, errInvalidSpanIDLength
 			}
 		} else if strings.HasPrefix(part, sampleFlagKey) {
 			// extract traceflag
 			scc.TraceFlags = parseTraceFlag(value)
+		} else if strings.HasPrefix(part, lineageKey) {
+			// extract lineage
+			lineageHeader := parseLineageHeader(value)
+			if isValidLineage(lineageHeader) {
+				lineageBaggage, _ := baggage.NewMember(lineageKey, lineageHeader)
+				currBaggage, _ = currBaggage.SetMember(lineageBaggage)
+			}
 		}
 	}
-	return trace.NewSpanContext(scc), nil
+
+	if currBaggage.Len() > 0 {
+		ctx = baggage.ContextWithBaggage(ctx, currBaggage)
+	}
+
+	return ctx, trace.NewSpanContext(scc), nil
 }
 
 // indexOf returns position of the first occurrence of a substr in str starting at pos index.
@@ -167,6 +201,40 @@ func parseTraceFlag(xraySampledFlag string) trace.TraceFlags {
 	return trace.FlagsSampled
 }
 
+func parseLineageHeader(xrayLineageHeader string) string {
+	numOfDelimiters := strings.Count(xrayLineageHeader, lineageDelimiter)
+	if len(xrayLineageHeader) < lineageMinLength ||
+		len(xrayLineageHeader) > lineageMaxLength ||
+		numOfDelimiters != 2 {
+		return invalidLineage
+	}
+
+	return xrayLineageHeader
+}
+
+func isValidLineage(lineage string) bool {
+	split := strings.Split(lineage, lineageDelimiter)
+	hash := split[1]
+	counter1 := parseIntWithBase(split[0], 10)
+	counter2 := parseIntWithBase(split[2], 10)
+
+	isHashValid := len(hash) == lineageHashLength && parseIntWithBase(hash, 16) != -1
+	isValidCounter2 := counter2 <= lineageMaxCounter2 && counter2 >= lineageMinCounter
+	isValidCounter1 := counter1 <= lineageMaxCounter1 && counter1 >= lineageMinCounter
+
+	return isHashValid && isValidCounter1 && isValidCounter2
+}
+
+func parseIntWithBase(s string, base int) int64 {
+	val, err := strconv.ParseInt(s, base, 64)
+
+	if err != nil {
+		return -1
+	}
+
+	return val
+}
+
 // Fields returns list of fields used by HTTPTextFormat.
 func (xray Propagator) Fields() []string {
 	return []string{traceHeaderKey}

propagators/aws/xray/propagator_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 
 	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/baggage"
 	"go.opentelemetry.io/otel/propagation"
 
 	"github.com/stretchr/testify/assert"
@@ -76,7 +77,7 @@ func TestAwsXrayExtract(t *testing.T) {
 			test.parentSpanID, traceHeaderDelimiter, sampleFlagKey, kvDelimiter, test.samplingFlag,
 		}, "")
 
-		sc, err := extract(headerVal)
+		_, sc, err := extract(context.Background(), headerVal)
 
 		info := []interface{}{
 			"trace ID: %q, parent span ID: %q, sampling flag: %q",
@@ -93,6 +94,149 @@ func TestAwsXrayExtract(t *testing.T) {
 	}
 }
 
+func TestAwsXrayExtractWithLineage(t *testing.T) {
+	testData := []struct {
+		lineage string
+
+		expectedBaggage map[string]string
+	}{
+		{
+			lineage: "32767:e65a2c4d:255",
+			expectedBaggage: map[string]string{
+				"Lineage": "32767:e65a2c4d:255",
+			},
+		},
+		{
+			lineage: "32767:e65a2c4d:255",
+			expectedBaggage: map[string]string{
+				"Lineage": "32767:e65a2c4d:255",
+				"cat":     "bark",
+				"dog":     "meow",
+			},
+		},
+		{
+			lineage:         "1::",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "1",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         ":",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "::",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "1:badc0de:13",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         ":fbadc0de:13",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "1:fbadc0de:",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "1::1",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "65535:fbadc0de:255",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "-213:e65a2c4d:255",
+			expectedBaggage: map[string]string{},
+		},
+		{
+			lineage:         "213:e65a2c4d:-22",
+			expectedBaggage: map[string]string{},
+		},
+	}
+
+	p := Propagator{}
+
+	for _, test := range testData {
+		carrier := make(map[string]string)
+		members := []baggage.Member{}
+		expectedSpanContext := trace.NewSpanContext(trace.SpanContextConfig{
+			TraceID:    traceID,
+			SpanID:     parentSpanID,
+			TraceFlags: trace.TraceFlags(0x00),
+			Remote:     true,
+		})
+
+		for key, value := range test.expectedBaggage {
+			member, _ := baggage.NewMember(key, value)
+			members = append(members, member)
+		}
+
+		expectedBaggage, _ := baggage.New(members...)
+		carrier[traceHeaderKey] = "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=0;Lineage=" + test.lineage
+
+		ctx := baggage.ContextWithBaggage(context.Background(), expectedBaggage)
+
+		if len(test.expectedBaggage) == 0 {
+			ctx = context.Background()
+		}
+
+		actualContext := p.Extract(ctx, propagation.MapCarrier(carrier))
+		spanContext := trace.SpanContextFromContext(actualContext)
+
+		assert.Equal(t, baggage.FromContext(actualContext), expectedBaggage)
+		assert.Equal(t, spanContext, expectedSpanContext)
+	}
+}
+
+func TestAwsXrayInjectWithLineage(t *testing.T) {
+	testData := []struct {
+		expectedBaggage map[string]string
+	}{
+		{
+			expectedBaggage: map[string]string{
+				"Lineage": "32767:e65a2c4d:255",
+				"cat":     "bark",
+				"dog":     "meow",
+			},
+		},
+		{
+			expectedBaggage: map[string]string{
+				"Lineage": "32767:e65a2c4d:255",
+			},
+		},
+	}
+
+	p := Propagator{}
+
+	for _, test := range testData {
+		carrier := make(map[string]string)
+		members := []baggage.Member{}
+
+		for key, value := range test.expectedBaggage {
+			member, _ := baggage.NewMember(key, value)
+			members = append(members, member)
+		}
+
+		expectedBaggage, _ := baggage.New(members...)
+
+		carrier[traceHeaderKey] = "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=0;Lineage=32767:e65a2c4d:255"
+
+		p.Inject(baggage.ContextWithBaggage(context.Background(), expectedBaggage), propagation.MapCarrier(carrier))
+
+		assert.Equal(t, carrier[traceHeaderKey], "Root=1-8a3c60f7-d188f8fa79d48a391a778fa6;Parent=53995c3f42cd8ad8;Sampled=0;Lineage=32767:e65a2c4d:255")
+	}
+}
+
 func BenchmarkPropagatorExtract(b *testing.B) {
 	propagator := Propagator{}