Read 128 bit trace ID in End Invocation (#30398)

Co-authored-by: Bryce Eadie <[email protected]>

pkg/serverless/invocationlifecycle/trace.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"regexp"
 	"strconv"
+	"strings"
 	"time"
 
 	json "github.com/json-iterator/go"
@@ -110,6 +111,9 @@ func (lp *LifecycleProcessor) endExecutionSpan(endDetails *InvocationEndDetails)
 		Meta:     lp.requestHandler.triggerTags,
 		Metrics:  lp.requestHandler.triggerMetrics,
 	}
+	if executionContext.TraceIDUpper64Hex != "" {
+		executionSpan.Meta[Upper64BitsTag] = executionContext.TraceIDUpper64Hex
+	}
 	executionSpan.Meta["request_id"] = endDetails.RequestID
 	executionSpan.Meta["cold_start"] = fmt.Sprintf("%t", endDetails.ColdStart)
 	if endDetails.ProactiveInit {
@@ -178,6 +182,9 @@ func (lp *LifecycleProcessor) completeInferredSpan(inferredSpan *inferredspan.In
 	}
 
 	inferredSpan.Span.TraceID = lp.GetExecutionInfo().TraceID
+	if lp.GetExecutionInfo().TraceIDUpper64Hex != "" {
+		inferredSpan.Span.Meta[Upper64BitsTag] = lp.GetExecutionInfo().TraceIDUpper64Hex
+	}
 
 	return inferredSpan.Span
 }
@@ -233,6 +240,30 @@ func InjectContext(executionContext *ExecutionStartInfo, headers http.Header) {
 		log.Debugf("injecting samplingPriority = %v", value)
 		executionContext.SamplingPriority = sampler.SamplingPriority(value)
 	}
+
+	upper64hex := getUpper64Hex(headers.Get(TraceTagsHeader))
+	if upper64hex != "" {
+		executionContext.TraceIDUpper64Hex = upper64hex
+	}
+}
+
+// searches traceTags for "_dd.p.tid=[upper 64 bits hex]" and returns that value if found
+func getUpper64Hex(traceTags string) string {
+	if !strings.Contains(traceTags, Upper64BitsTag) {
+		return ""
+	}
+	kvpairs := strings.Split(traceTags, ",")
+	for _, pair := range kvpairs {
+		if !strings.Contains(pair, Upper64BitsTag) {
+			continue
+		}
+		kv := strings.Split(pair, "=")
+		if len(kv) != 2 {
+			return ""
+		}
+		return kv[1]
+	}
+	return ""
 }
 
 // InjectSpanID injects the spanId

pkg/serverless/invocationlifecycle/trace_test.go

@@ -60,6 +60,20 @@ func TestInjectContextWithContext(t *testing.T) {
 	assert.Equal(t, sampler.PriorityUserKeep, currentExecutionInfo.SamplingPriority)
 }
 
+func TestInjectContextWith128BitTraceID(t *testing.T) {
+	currentExecutionInfo := newExecutionContextWithTime()
+	httpHeaders := http.Header{}
+	httpHeaders.Set("x-datadog-trace-id", "1234")
+	httpHeaders.Set("x-datadog-parent-id", "5678")
+	httpHeaders.Set("x-datadog-sampling-priority", "2")
+	httpHeaders.Set(TraceTagsHeader, Upper64BitsTag+"=1234567890abcdef")
+	InjectContext(currentExecutionInfo, httpHeaders)
+	assert.Equal(t, uint64(1234), currentExecutionInfo.TraceID)
+	assert.Equal(t, uint64(5678), currentExecutionInfo.parentID)
+	assert.Equal(t, sampler.PriorityUserKeep, currentExecutionInfo.SamplingPriority)
+	assert.Equal(t, "1234567890abcdef", currentExecutionInfo.TraceIDUpper64Hex)
+}
+
 func TestInjectSpanIDNoContext(t *testing.T) {
 	currentExecutionInfo := newExecutionContextWithTime()
 	InjectSpanID(currentExecutionInfo, nil)
@@ -510,6 +524,8 @@ func TestEndExecutionSpanWithNoError(t *testing.T) {
 		InvokeEventHeaders: http.Header{},
 	}
 	lp.startExecutionSpan(nil, []byte(testString), startDetails)
+	execInfo := lp.GetExecutionInfo()
+	execInfo.TraceIDUpper64Hex = "1234567890abcdef"
 
 	duration := 1 * time.Second
 	endTime := startTime.Add(duration)
@@ -530,16 +546,17 @@ func TestEndExecutionSpanWithNoError(t *testing.T) {
 		"cold_start":                "true",
 		"function.request.resource": "/users/create",
 		"function.request.path":     "/users/create",
-		"function.request.headers.x-datadog-parent-id":         "1480558859903409531",
-		"function.request.headers.x-datadog-trace-id":          "5736943178450432258",
+		"function.request.headers.x-datadog-parent-id": "1480558859903409531",
+		"function.request.headers.x-datadog-trace-id":  "5736943178450432258",
+		"_dd.p.tid": "1234567890abcdef",
 		"function.request.headers.x-datadog-sampling-priority": "1",
 		"function.request.httpMethod":                          "GET",
 		"function.request.headers.Accept":                      "*/*",
 		"function.request.headers.Accept-Encoding":             "gzip",
 		"function.response.response":                           "test response payload",
 		"language":                                             "dotnet",
 	}
-	assert.Equal(t, executionSpan.Meta, expectingResultMetaMap)
+	assert.Equal(t, expectingResultMetaMap, executionSpan.Meta)
 	assert.Equal(t, "aws.lambda", executionSpan.Name)
 	assert.Equal(t, "aws.lambda", executionSpan.Service)
 	assert.Equal(t, "TestFunction", executionSpan.Resource)
@@ -1059,3 +1076,43 @@ func TestCompleteInferredSpanWithAsync(t *testing.T) {
 	assert.Equal(t, duration.Nanoseconds(), span.Duration)
 	assert.Equal(t, int32(0), inferredSpan.Span.Error)
 }
+
+func Test_getUpper64Hex(t *testing.T) {
+	tests := []struct {
+		name string
+		tags string
+		want string
+	}{
+		{
+			name: "just a trace tag",
+			tags: "_dd.p.tid=1234567890abcdef",
+			want: "1234567890abcdef",
+		},
+		{
+			name: "nothing",
+			tags: "",
+			want: "",
+		},
+		{
+			name: "multiple tags 1",
+			tags: "some=tag,_dd.p.tid=1234567890abcdef",
+			want: "1234567890abcdef",
+		},
+		{
+			name: "multiple tags 2",
+			tags: "_dd.p.tid=1234567890abcdef,some=tag",
+			want: "1234567890abcdef",
+		},
+		{
+			name: "multiple tags 3",
+			tags: "some=tag,_dd.p.tid=1234567890abcdef,another=tag",
+			want: "1234567890abcdef",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert.Equalf(t, tt.want, getUpper64Hex(tt.tags), "getUpper64Hex(%v)", tt.tags)
+
+		})
+	}
+}

releasenotes/notes/read-128-bit-trace-id-in-end-invocation-9d44a9373f4f712d.yaml

@@ -0,0 +1,6 @@
+---
+enhancements:
+  - |
+    The AWS Lambda Extension is now  able to read the full 128-bit trace ID 
+    from the headers of the end-invocation HTTP request made by dd-trace or the
+    datadog-lambda-go library.