From 7827c658d48b3ec17fb41e14d358875c694175a0 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Wed, 12 Nov 2025 14:18:55 -0800 Subject: [PATCH 1/7] Add request-id as header to start/end invocation. --- internal/extension/extension.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/internal/extension/extension.go b/internal/extension/extension.go index 68e2813..e8bf52c 100644 --- a/internal/extension/extension.go +++ b/internal/extension/extension.go @@ -24,6 +24,7 @@ import ( "time" "github.com/DataDog/datadog-lambda-go/internal/logger" + "github.com/aws/aws-lambda-go/lambdacontext" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" ) @@ -119,6 +120,8 @@ func (em *ExtensionManager) checkAgentRunning() { func (em *ExtensionManager) SendStartInvocationRequest(ctx context.Context, eventPayload json.RawMessage) context.Context { body := bytes.NewBuffer(eventPayload) req, _ := http.NewRequest(http.MethodPost, em.startInvocationUrl, body) + lc, _ := lambdacontext.FromContext(ctx) + req.Header.Set("x-datadog-aws-request-id", lc.AwsRequestID) response, err := em.httpClient.Do(req) if response != nil && response.Body != nil { defer func() { @@ -157,6 +160,8 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi } body := bytes.NewBuffer(content) req, _ := http.NewRequest(http.MethodPost, em.endInvocationUrl, body) + lc, _ := lambdacontext.FromContext(ctx) + req.Header.Set("x-datadog-aws-request-id", lc.AwsRequestID) // Mark the invocation as an error if any if cfg.Error != nil { From 9b1a4a5d991d3cbbf94c9c0f35d8362159464c5e Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Thu, 13 Nov 2025 09:41:58 -0800 Subject: [PATCH 2/7] Use consistent header name. --- internal/extension/extension.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/internal/extension/extension.go b/internal/extension/extension.go index e8bf52c..1006675 100644 --- a/internal/extension/extension.go +++ b/internal/extension/extension.go @@ -44,6 +44,8 @@ const ( DdSeverlessSpan ddTraceContext = "dd-tracer-serverless-span" DdLambdaResponse ddTraceContext = "dd-response" + + lambdaRuntimeAwsRequestIdHeader = "lambda-runtime-aws-request-id" ) const ( @@ -121,7 +123,7 @@ func (em *ExtensionManager) SendStartInvocationRequest(ctx context.Context, even body := bytes.NewBuffer(eventPayload) req, _ := http.NewRequest(http.MethodPost, em.startInvocationUrl, body) lc, _ := lambdacontext.FromContext(ctx) - req.Header.Set("x-datadog-aws-request-id", lc.AwsRequestID) + req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) response, err := em.httpClient.Do(req) if response != nil && response.Body != nil { defer func() { @@ -161,7 +163,7 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi body := bytes.NewBuffer(content) req, _ := http.NewRequest(http.MethodPost, em.endInvocationUrl, body) lc, _ := lambdacontext.FromContext(ctx) - req.Header.Set("x-datadog-aws-request-id", lc.AwsRequestID) + req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) // Mark the invocation as an error if any if cfg.Error != nil { From 4b2cbd4f0b186951dae75c5cec8e369b574f7a98 Mon Sep 17 00:00:00 2001 From: Olivier John Ndjike Nzia Date: Mon, 17 Nov 2025 11:16:07 -0500 Subject: [PATCH 3/7] add tests to check for lambda request id in headers --- internal/extension/extension_test.go | 45 ++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/internal/extension/extension_test.go b/internal/extension/extension_test.go index 5b391f4..81694f8 100644 --- a/internal/extension/extension_test.go +++ b/internal/extension/extension_test.go @@ -19,6 +19,7 @@ import ( "testing" "github.com/DataDog/datadog-lambda-go/internal/logger" + "github.com/aws/aws-lambda-go/lambdacontext" "github.com/stretchr/testify/assert" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" @@ -153,6 +154,27 @@ func TestExtensionStartInvoke(t *testing.T) { assert.Nil(t, samplingPriority) } +func TestExtensionStartInvokeLambdaRequestId(t *testing.T) { + headers := http.Header{} + capturingClient := capturingClient{hdr: headers} + + em := &ExtensionManager{ + startInvocationUrl: startInvocationUrl, + httpClient: capturingClient, + } + + lc := &lambdacontext.LambdaContext{ + AwsRequestID: "test-request-id-12345", + } + ctx := lambdacontext.NewContext(context.TODO(), lc) + em.SendStartInvocationRequest(ctx, []byte{}) + + err := em.Flush() + + assert.Nil(t, err) + assert.Equal(t, "test-request-id-12345", headers.Get("lambda-runtime-aws-request-id")) +} + func TestExtensionStartInvokeWithTraceContext(t *testing.T) { headers := http.Header{} headers.Set(string(DdTraceId), mockTraceId) @@ -215,6 +237,29 @@ func TestExtensionEndInvocation(t *testing.T) { assert.Equal(t, 1, len(lines)) } +func TestExtensionEndInvokeLambdaRequestId(t *testing.T) { + headers := http.Header{} + capturingClient := capturingClient{hdr: headers} + + em := &ExtensionManager{ + startInvocationUrl: startInvocationUrl, + httpClient: capturingClient, + } + + lc := &lambdacontext.LambdaContext{ + AwsRequestID: "test-request-id-12345", + } + + ctx := lambdacontext.NewContext(context.TODO(), lc) + span := tracer.StartSpan("aws.lambda") + span.Finish() + cfg := ddtrace.FinishConfig{} + em.SendEndInvocationRequest(ctx, span, cfg) + err := em.Flush() + assert.Nil(t, err) + assert.Equal(t, "test-request-id-12345", headers.Get("lambda-runtime-aws-request-id")) +} + func TestExtensionEndInvocationError(t *testing.T) { em := &ExtensionManager{ endInvocationUrl: endInvocationUrl, From 07f6ea8f2ca01aa0e3af56e1f1edda17dd50b7eb Mon Sep 17 00:00:00 2001 From: Olivier John Ndjike Nzia Date: Mon, 17 Nov 2025 11:51:45 -0500 Subject: [PATCH 4/7] add error tests for lambda requestIDs --- internal/extension/extension.go | 16 ++++++++--- internal/extension/extension_test.go | 41 +++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/internal/extension/extension.go b/internal/extension/extension.go index 1006675..07ab6b4 100644 --- a/internal/extension/extension.go +++ b/internal/extension/extension.go @@ -122,8 +122,13 @@ func (em *ExtensionManager) checkAgentRunning() { func (em *ExtensionManager) SendStartInvocationRequest(ctx context.Context, eventPayload json.RawMessage) context.Context { body := bytes.NewBuffer(eventPayload) req, _ := http.NewRequest(http.MethodPost, em.startInvocationUrl, body) - lc, _ := lambdacontext.FromContext(ctx) - req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) + + if lc, ok := lambdacontext.FromContext(ctx); ok { + req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) + } else { + logger.Error(fmt.Errorf("missing lambda Context. Unable to set lambda-runtime-aws-request-id header")) + } + response, err := em.httpClient.Do(req) if response != nil && response.Body != nil { defer func() { @@ -162,8 +167,11 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi } body := bytes.NewBuffer(content) req, _ := http.NewRequest(http.MethodPost, em.endInvocationUrl, body) - lc, _ := lambdacontext.FromContext(ctx) - req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) + if lc, ok := lambdacontext.FromContext(ctx); ok { + req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) + } else { + logger.Error(fmt.Errorf("missing lambda Context. Unable to set lambda-runtime-aws-request-id header")) + } // Mark the invocation as an error if any if cfg.Error != nil { diff --git a/internal/extension/extension_test.go b/internal/extension/extension_test.go index 81694f8..100ccf7 100644 --- a/internal/extension/extension_test.go +++ b/internal/extension/extension_test.go @@ -175,6 +175,23 @@ func TestExtensionStartInvokeLambdaRequestId(t *testing.T) { assert.Equal(t, "test-request-id-12345", headers.Get("lambda-runtime-aws-request-id")) } +func TestExtensionStartInvokeLambdaRequestIdError(t *testing.T) { + headers := http.Header{} + capturingClient := capturingClient{hdr: headers} + + em := &ExtensionManager{ + startInvocationUrl: startInvocationUrl, + httpClient: capturingClient, + } + + logOutput := captureLog(func() { em.SendStartInvocationRequest(context.TODO(), []byte{}) }) + err := em.Flush() + assert.Nil(t, err) + assert.Contains(t, logOutput, "missing lambda Context. Unable to set lambda-runtime-aws-request-id header") + lines := strings.Split(strings.TrimSpace(logOutput), "\n") + assert.Equal(t, 1, len(lines)) +} + func TestExtensionStartInvokeWithTraceContext(t *testing.T) { headers := http.Header{} headers.Set(string(DdTraceId), mockTraceId) @@ -227,8 +244,9 @@ func TestExtensionEndInvocation(t *testing.T) { endInvocationUrl: endInvocationUrl, httpClient: &ClientSuccessEndInvoke{}, } + ctx := lambdacontext.NewContext(context.TODO(), &lambdacontext.LambdaContext{}) span := tracer.StartSpan("aws.lambda") - logOutput := captureLog(func() { em.SendEndInvocationRequest(context.TODO(), span, ddtrace.FinishConfig{}) }) + logOutput := captureLog(func() { em.SendEndInvocationRequest(ctx, span, ddtrace.FinishConfig{}) }) span.Finish() // Expected because the noopSpanContext doesn't have the SamplingPriority() and we cannot use the mock for the agent assert.Contains(t, logOutput, "could not get sampling priority from getSamplingPriority()") @@ -260,6 +278,27 @@ func TestExtensionEndInvokeLambdaRequestId(t *testing.T) { assert.Equal(t, "test-request-id-12345", headers.Get("lambda-runtime-aws-request-id")) } +func TestExtensionEndInvokeLambdaRequestIdError(t *testing.T) { + headers := http.Header{} + capturingClient := capturingClient{hdr: headers} + ctx := context.WithValue(context.TODO(), DdSamplingPriority, mockSamplingPriority) + ctx = context.WithValue(ctx, DdTraceId, mockTraceId) + em := &ExtensionManager{ + startInvocationUrl: startInvocationUrl, + httpClient: capturingClient, + } + + span := tracer.StartSpan("aws.lambda") + logOutput := captureLog(func() { em.SendEndInvocationRequest(ctx, span, ddtrace.FinishConfig{}) }) + span.Finish() + + err := em.Flush() + assert.Nil(t, err) + assert.Contains(t, logOutput, "missing lambda Context. Unable to set lambda-runtime-aws-request-id header") + lines := strings.Split(strings.TrimSpace(logOutput), "\n") + assert.Equal(t, 1, len(lines)) +} + func TestExtensionEndInvocationError(t *testing.T) { em := &ExtensionManager{ endInvocationUrl: endInvocationUrl, From 2b3cb98131e7500903f430ba4acd2c50945dd400 Mon Sep 17 00:00:00 2001 From: Olivier John Ndjike Nzia Date: Mon, 17 Nov 2025 14:45:20 -0500 Subject: [PATCH 5/7] fix tests --- internal/extension/extension_test.go | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/internal/extension/extension_test.go b/internal/extension/extension_test.go index 100ccf7..6535c03 100644 --- a/internal/extension/extension_test.go +++ b/internal/extension/extension_test.go @@ -176,12 +176,9 @@ func TestExtensionStartInvokeLambdaRequestId(t *testing.T) { } func TestExtensionStartInvokeLambdaRequestIdError(t *testing.T) { - headers := http.Header{} - capturingClient := capturingClient{hdr: headers} - em := &ExtensionManager{ startInvocationUrl: startInvocationUrl, - httpClient: capturingClient, + httpClient: &ClientSuccessStartInvoke{}, } logOutput := captureLog(func() { em.SendStartInvocationRequest(context.TODO(), []byte{}) }) @@ -260,8 +257,8 @@ func TestExtensionEndInvokeLambdaRequestId(t *testing.T) { capturingClient := capturingClient{hdr: headers} em := &ExtensionManager{ - startInvocationUrl: startInvocationUrl, - httpClient: capturingClient, + endInvocationUrl: endInvocationUrl, + httpClient: capturingClient, } lc := &lambdacontext.LambdaContext{ @@ -284,10 +281,10 @@ func TestExtensionEndInvokeLambdaRequestIdError(t *testing.T) { ctx := context.WithValue(context.TODO(), DdSamplingPriority, mockSamplingPriority) ctx = context.WithValue(ctx, DdTraceId, mockTraceId) em := &ExtensionManager{ - startInvocationUrl: startInvocationUrl, - httpClient: capturingClient, + endInvocationUrl: endInvocationUrl, + httpClient: capturingClient, } - + span := tracer.StartSpan("aws.lambda") logOutput := captureLog(func() { em.SendEndInvocationRequest(ctx, span, ddtrace.FinishConfig{}) }) span.Finish() From 08a7cfc2dd55a7b2653fc2f422a27ea987b00e7a Mon Sep 17 00:00:00 2001 From: Olivier John Ndjike Nzia Date: Mon, 17 Nov 2025 16:49:12 -0500 Subject: [PATCH 6/7] better logs --- internal/extension/extension_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/extension/extension_test.go b/internal/extension/extension_test.go index 6535c03..07ae0e9 100644 --- a/internal/extension/extension_test.go +++ b/internal/extension/extension_test.go @@ -184,7 +184,7 @@ func TestExtensionStartInvokeLambdaRequestIdError(t *testing.T) { logOutput := captureLog(func() { em.SendStartInvocationRequest(context.TODO(), []byte{}) }) err := em.Flush() assert.Nil(t, err) - assert.Contains(t, logOutput, "missing lambda Context. Unable to set lambda-runtime-aws-request-id header") + assert.Contains(t, logOutput, "missing AWS Lambda context. Unable to set lambda-runtime-aws-request-id header") lines := strings.Split(strings.TrimSpace(logOutput), "\n") assert.Equal(t, 1, len(lines)) } @@ -291,7 +291,7 @@ func TestExtensionEndInvokeLambdaRequestIdError(t *testing.T) { err := em.Flush() assert.Nil(t, err) - assert.Contains(t, logOutput, "missing lambda Context. Unable to set lambda-runtime-aws-request-id header") + assert.Contains(t, logOutput, "missing AWS Lambda context. Unable to set lambda-runtime-aws-request-id header") lines := strings.Split(strings.TrimSpace(logOutput), "\n") assert.Equal(t, 1, len(lines)) } From 7bca976727d8adaf74b41107ac69875dd5a5651c Mon Sep 17 00:00:00 2001 From: Olivier John Ndjike Nzia Date: Mon, 17 Nov 2025 16:49:26 -0500 Subject: [PATCH 7/7] better logs p2 --- internal/extension/extension.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/extension/extension.go b/internal/extension/extension.go index 07ab6b4..7b32714 100644 --- a/internal/extension/extension.go +++ b/internal/extension/extension.go @@ -126,7 +126,7 @@ func (em *ExtensionManager) SendStartInvocationRequest(ctx context.Context, even if lc, ok := lambdacontext.FromContext(ctx); ok { req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) } else { - logger.Error(fmt.Errorf("missing lambda Context. Unable to set lambda-runtime-aws-request-id header")) + logger.Error(fmt.Errorf("missing AWS Lambda context. Unable to set lambda-runtime-aws-request-id header")) } response, err := em.httpClient.Do(req) @@ -170,7 +170,7 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi if lc, ok := lambdacontext.FromContext(ctx); ok { req.Header.Set(lambdaRuntimeAwsRequestIdHeader, lc.AwsRequestID) } else { - logger.Error(fmt.Errorf("missing lambda Context. Unable to set lambda-runtime-aws-request-id header")) + logger.Error(fmt.Errorf("missing AWS Lambda context. Unable to set lambda-runtime-aws-request-id header")) } // Mark the invocation as an error if any