Read lambda trace context header instead of using current segment

Gopkg.toml

@@ -27,7 +27,7 @@
 
 [[constraint]]
   name = "github.com/aws/aws-lambda-go"
-  version = "1.10.0"
+  version = "^1.10.0"
 
 [prune]
   go-tests = true

internal/trace/trace_context.go

@@ -9,6 +9,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/aws/aws-xray-sdk-go/header"
 	"github.com/aws/aws-xray-sdk-go/xray"
 )
 
@@ -53,9 +54,11 @@ func GetTraceHeaders(ctx context.Context, useCurrentSegmentAsParent bool) map[st
 
 		if useCurrentSegmentAsParent {
 			segment := xray.GetSegment(ctx)
-			newParentID, err := convertXRayEntityIDToAPMParentID(segment.ID)
-			if err == nil {
-				parentID = newParentID
+			if segment != nil {
+				newParentID, err := convertXRayEntityIDToAPMParentID(segment.ID)
+				if err == nil {
+					parentID = newParentID
+				}
 			}
 		}
 
@@ -113,27 +116,38 @@ func unmarshalEventForTraceContext(ev json.RawMessage) (map[string]string, bool)
 func convertTraceContextFromXRay(ctx context.Context) (map[string]string, error) {
 	traceContext := map[string]string{}
 
-	segment := xray.GetSegment(ctx)
-	if segment == nil {
+	header := getLambdaTraceHeaderFromContext(ctx)
+	if header == nil {
 		return traceContext, fmt.Errorf("xray segment doesn't exist, couldn't read trace context")
 	}
 
-	traceID, err := convertXRayTraceIDToAPMTraceID(segment.TraceID)
+	traceID, err := convertXRayTraceIDToAPMTraceID(header.TraceID)
 	if err != nil {
 		return traceContext, fmt.Errorf("couldn't read trace id from xray: %v", err)
 	}
-	parentID, err := convertXRayEntityIDToAPMParentID(segment.ID)
+	parentID, err := convertXRayEntityIDToAPMParentID(header.ParentID)
 	if err != nil {
 		return traceContext, fmt.Errorf("couldn't read parent id from xray: %v", err)
 	}
-	samplingPriority := convertXRaySampling(segment.Sampled)
+	samplingPriority := convertXRaySamplingDecision(header.SamplingDecision)
 
 	traceContext[traceIDHeader] = traceID
 	traceContext[parentIDHeader] = parentID
 	traceContext[samplingPriorityHeader] = samplingPriority
 	return traceContext, nil
 }
 
+// getTraceHeaderFromContext is used to extract xray segment metadata from the lambda context object
+func getLambdaTraceHeaderFromContext(ctx context.Context) *header.Header {
+	var traceHeader string
+
+	if traceHeaderValue := ctx.Value(xray.LambdaTraceHeaderKey); traceHeaderValue != nil {
+		traceHeader = traceHeaderValue.(string)
+		return header.FromString(traceHeader)
+	}
+	return nil
+}
+
 // Converts the last 63 bits of an X-Ray trace ID (hex) to a Datadog trace id (uint64).
 func convertXRayTraceIDToAPMTraceID(traceID string) (string, error) {
 	parts := strings.Split(traceID, "-")
@@ -177,8 +191,8 @@ func convertXRayEntityIDToAPMParentID(entityID string) (string, error) {
 }
 
 // Converts an X-Ray sampled flag into it's Datadog counterpart.
-func convertXRaySampling(sampled bool) string {
-	if sampled {
+func convertXRaySamplingDecision(decision header.SamplingDecision) string {
+	if decision == header.Sampled {
 		return userKeep
 	}
 	return userReject

internal/trace/trace_context_test.go

@@ -4,11 +4,29 @@ import (
 	"context"
 	"testing"
 
+	"github.com/aws/aws-xray-sdk-go/header"
+
 	"github.com/aws/aws-xray-sdk-go/xray"
 
 	"github.com/stretchr/testify/assert"
 )
 
+func mockLambdaTraceContext(ctx context.Context, traceID, parentID string, sampled bool) context.Context {
+	decision := header.NotSampled
+	if sampled {
+		decision = header.Sampled
+	}
+
+	traceHeader := header.Header{
+		TraceID:          traceID,
+		ParentID:         parentID,
+		SamplingDecision: decision,
+		AdditionalData:   make(map[string]string),
+	}
+	headerString := traceHeader.String()
+	return context.WithValue(ctx, xray.LambdaTraceHeaderKey, headerString)
+}
+
 func TestUnmarshalEventForTraceMetadataNonProxyEvent(t *testing.T) {
 	ev := loadRawJSON(t, "testdata/apig-event-metadata.json")
 
@@ -79,7 +97,8 @@ func TestXrayTraceContextNoSegment(t *testing.T) {
 	assert.Error(t, err)
 }
 func TestXrayTraceContextWithSegment(t *testing.T) {
-	ctx, _ := xray.BeginSegment(context.Background(), "Test-Segment")
+
+	ctx := mockLambdaTraceContext(context.Background(), "1-5ce31dc2-2c779014b90ce44db5e03875", "779014b90ce44db5e03875", true)
 
 	headers, err := convertTraceContextFromXRay(ctx)
 	assert.NoError(t, err)
@@ -90,7 +109,7 @@ func TestXrayTraceContextWithSegment(t *testing.T) {
 
 func TestExtractTraceContextFromContext(t *testing.T) {
 	ev := loadRawJSON(t, "testdata/apig-event-no-metadata.json")
-	ctx, _ := xray.BeginSegment(context.Background(), "Test-Segment")
+	ctx := mockLambdaTraceContext(context.Background(), "1-5ce31dc2-2c779014b90ce44db5e03875", "779014b90ce44db5e03875", true)
 
 	newCTX, err := ExtractTraceContext(ctx, *ev)
 	headers := GetTraceHeaders(newCTX, false)
@@ -102,7 +121,7 @@ func TestExtractTraceContextFromContext(t *testing.T) {
 }
 func TestExtractTraceContextFromEvent(t *testing.T) {
 	ev := loadRawJSON(t, "testdata/apig-event-metadata.json")
-	ctx, _ := xray.BeginSegment(context.Background(), "Test-Segment")
+	ctx := mockLambdaTraceContext(context.Background(), "1-5ce31dc2-2c779014b90ce44db5e03875", "779014b90ce44db5e03875", true)
 
 	newCTX, err := ExtractTraceContext(ctx, *ev)
 	headers := GetTraceHeaders(newCTX, false)
@@ -126,9 +145,12 @@ func TestExtractTraceContextFail(t *testing.T) {
 
 func TestGetTraceHeadersWithUpdatedParent(t *testing.T) {
 	ev := loadRawJSON(t, "testdata/apig-event-metadata.json")
-	ctx, _ := xray.BeginSegment(context.Background(), "Test-Segment")
+	ctx := mockLambdaTraceContext(context.Background(), "1-5ce31dc2-2c779014b90ce44db5e03875", "779014b90ce44db5e03874", true)
 
 	ctx, _ = ExtractTraceContext(ctx, *ev)
+
+	ctx, _ = xray.BeginSubsegment(ctx, "The Subsegment")
+
 	headers := GetTraceHeaders(ctx, true)
 	assert.Equal(t, "2", headers[samplingPriorityHeader])
 	assert.Equal(t, "1231452342", headers[traceIDHeader])