Skip to content

add request-id to lambda start/end invocation calls #9991

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: master
Choose a base branch
from
Open

add request-id to lambda start/end invocation calls #9991

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c52606f
add request-id to lambda start/end invocation calls
ojproductions Nov 17, 2025
dcf196f
tests
ojproductions Nov 18, 2025
1886e72
Merge branch 'master' into onzia/request-id-header
ojproductions Nov 18, 2025
1f69e8e
commenting and naming
ojproductions Nov 18, 2025
8979821
Merge branch 'master' into onzia/request-id-header
ojproductions Nov 19, 2025
3241b0f
remove unused coreTracer param
ojproductions Nov 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions ...c/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,15 +88,15 @@ static AgentScope enter(
if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) {
return null;
}

AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in);
String lambdaRequestId = awsContext.getAwsRequestId();
AgentSpanContext lambdaContext = AgentTracer.get().notifyExtensionStart(in, lambdaRequestId);
final AgentSpan span;
if (null == lambdaContext) {
span = startSpan(INVOCATION_SPAN_NAME);
} else {
span = startSpan(INVOCATION_SPAN_NAME, lambdaContext);
}
span.setTag("request_id", awsContext.getAwsRequestId());
span.setTag("request_id", lambdaRequestId);

final AgentScope scope = activateSpan(span);
return scope;
Expand All @@ -107,6 +107,7 @@ static void exit(
@Origin String method,
@Enter final AgentScope scope,
@Advice.Argument(1) final Object result,
@Advice.Argument(2) final Context awsContext,
@Advice.Thrown final Throwable throwable) {

if (scope == null) {
Expand All @@ -120,8 +121,10 @@ static void exit(
if (throwable != null) {
span.addThrowable(throwable);
}
String lambdaRequestId = awsContext.getAwsRequestId();

span.finish();
AgentTracer.get().notifyExtensionEnd(span, result, null != throwable);
AgentTracer.get().notifyExtensionEnd(span, result, null != throwable, lambdaRequestId);
} finally {
scope.close();
}
Expand Down
9 changes: 5 additions & 4 deletions dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1194,13 +1194,14 @@ public void closeActive() {
}

@Override
public AgentSpanContext notifyExtensionStart(Object event) {
return LambdaHandler.notifyStartInvocation(this, event);
public AgentSpanContext notifyExtensionStart(Object event, String lambdaRequestId) {
return LambdaHandler.notifyStartInvocation(this, event, lambdaRequestId);
}

@Override
public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {
LambdaHandler.notifyEndInvocation(span, result, isError);
public void notifyExtensionEnd(
AgentSpan span, Object result, boolean isError, String lambdaRequestId) {
LambdaHandler.notifyEndInvocation(span, result, isError, lambdaRequestId);
}

@Override
Expand Down
14 changes: 11 additions & 3 deletions dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ public class LambdaHandler {
private static final String DATADOG_INVOCATION_ERROR_MSG = "x-datadog-invocation-error-msg";
private static final String DATADOG_INVOCATION_ERROR_TYPE = "x-datadog-invocation-error-type";
private static final String DATADOG_INVOCATION_ERROR_STACK = "x-datadog-invocation-error-stack";
private static final String LAMBDA_RUNTIME_AWS_REQUEST_ID = "lambda-runtime-aws-request-id";

private static final String START_INVOCATION = "/lambda/start-invocation";
private static final String END_INVOCATION = "/lambda/end-invocation";
Expand Down Expand Up @@ -72,14 +73,19 @@ public class LambdaHandler {

private static String EXTENSION_BASE_URL = "http://127.0.0.1:8124";

public static AgentSpanContext notifyStartInvocation(CoreTracer tracer, Object event) {
RequestBody body = RequestBody.create(jsonMediaType, writeValueAsString(event));
public static AgentSpanContext notifyStartInvocation(
CoreTracer tracer, Object event, String lambdaRequestId) {
RequestBody body =
RequestBody.create(
jsonMediaType,
writeValueAsString(event)); // MAYBE remove CoreTracer parameter since it is not used
Copy link
Contributor

@mcculls mcculls Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If there's no plan to use it in the near future then I would remove the CoreTracer parameter now

try (Response response =
HTTP_CLIENT
.newCall(
new Request.Builder()
.url(EXTENSION_BASE_URL + START_INVOCATION)
.addHeader(DATADOG_META_LANG, "java")
.addHeader(LAMBDA_RUNTIME_AWS_REQUEST_ID, lambdaRequestId)
.post(body)
.build())
.execute()) {
Expand All @@ -99,7 +105,8 @@ public static AgentSpanContext notifyStartInvocation(CoreTracer tracer, Object e
return null;
}

public static boolean notifyEndInvocation(AgentSpan span, Object result, boolean isError) {
public static boolean notifyEndInvocation(
AgentSpan span, Object result, boolean isError, String lambdaRequestId) {
if (null == span || null == span.getSamplingPriority()) {
log.error(
"could not notify the extension as the lambda span is null or no sampling priority has been found");
Expand All @@ -113,6 +120,7 @@ public static boolean notifyEndInvocation(AgentSpan span, Object result, boolean
.addHeader(DATADOG_SPAN_ID, DDSpanId.toString(span.getSpanId()))
.addHeader(DATADOG_SAMPLING_PRIORITY, span.getSamplingPriority().toString())
.addHeader(DATADOG_META_LANG, "java")
.addHeader(LAMBDA_RUNTIME_AWS_REQUEST_ID, lambdaRequestId)
.post(body);

Object errorMessage = span.getTag(DDTags.ERROR_MSG);
Expand Down
30 changes: 18 additions & 12 deletions dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,12 @@ class LambdaHandlerTest extends DDCoreSpecification {
LambdaHandler.setExtensionBaseUrl(server.address.toString())

when:
def objTest = LambdaHandler.notifyStartInvocation(ct, obj)
def objTest = LambdaHandler.notifyStartInvocation(ct, obj, "lambda-request-123")

then:
objTest.getTraceId().toString() == traceId
objTest.getSamplingPriority() == samplingPriority
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "lambda-request-123"

cleanup:
server.close()
Expand Down Expand Up @@ -84,11 +85,12 @@ class LambdaHandlerTest extends DDCoreSpecification {
LambdaHandler.setExtensionBaseUrl(server.address.toString())

when:
def objTest = LambdaHandler.notifyStartInvocation(ct, obj)
def objTest = LambdaHandler.notifyStartInvocation(ct, obj, "lambda-request-123")

then:
objTest.getTraceId().toHexString() == traceId
objTest.getSamplingPriority() == samplingPriority
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "lambda-request-123"

cleanup:
server.close()
Expand All @@ -115,10 +117,11 @@ class LambdaHandlerTest extends DDCoreSpecification {
LambdaHandler.setExtensionBaseUrl(server.address.toString())

when:
def objTest = LambdaHandler.notifyStartInvocation(ct, obj)
def objTest = LambdaHandler.notifyStartInvocation(ct, obj, "my-lambda-request")

then:
objTest == expected
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "my-lambda-request"

cleanup:
server.close()
Expand Down Expand Up @@ -148,22 +151,23 @@ class LambdaHandlerTest extends DDCoreSpecification {
}

when:
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue)
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue, lambdaReqIdHeaderValue)

then:
server.lastRequest.headers.get("x-datadog-invocation-error") == eHeaderValue
server.lastRequest.headers.get("x-datadog-trace-id") == tIdHeaderValue
server.lastRequest.headers.get("x-datadog-span-id") == sIdHeaderValue
server.lastRequest.headers.get("x-datadog-sampling-priority") == sPIdHeaderValue
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == lambdaReqIdHeaderValue
result == expected

cleanup:
server.close()

where:
expected | eHeaderValue | tIdHeaderValue | sIdHeaderValue | sPIdHeaderValue | lambdaResult | boolValue
true | "true" | "1234" | "5678" | "2" | {} | true
true | null | "1234" | "5678" | "2" | "12345" | false
expected | eHeaderValue | tIdHeaderValue | sIdHeaderValue | sPIdHeaderValue | lambdaResult | boolValue | lambdaReqIdHeaderValue
true | "true" | "1234" | "5678" | "2" | {} | true | "request123"
true | null | "1234" | "5678" | "2" | "12345" | false | "request456"
}

def "test end invocation failure"() {
Expand All @@ -185,19 +189,20 @@ class LambdaHandlerTest extends DDCoreSpecification {
}

when:
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue)
def result = LambdaHandler.notifyEndInvocation(span, lambdaResult, boolValue, lambdaReqIdHeaderValue)

then:
result == expected
server.lastRequest.headers.get("x-datadog-invocation-error") == headerValue
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == lambdaReqIdHeaderValue

cleanup:
server.close()

where:
expected | headerValue | lambdaResult | boolValue
false | "true" | {} | true
false | null | "12345" | false
expected | headerValue | lambdaResult | boolValue | lambdaReqIdHeaderValue
false | "true" | {} | true | "request123"
false | null | "12345" | false | "request456"
}

def "test end invocation success with error metadata"() {
Expand All @@ -222,13 +227,14 @@ class LambdaHandlerTest extends DDCoreSpecification {
}

when:
LambdaHandler.notifyEndInvocation(span, {}, true)
LambdaHandler.notifyEndInvocation(span, {}, true, "lambda-request-123")

then:
server.lastRequest.headers.get("x-datadog-invocation-error") == "true"
server.lastRequest.headers.get("x-datadog-invocation-error-msg") == "custom error message"
server.lastRequest.headers.get("x-datadog-invocation-error-type") == "java.lang.Throwable"
server.lastRequest.headers.get("x-datadog-invocation-error-stack") == "ZXJyb3JTdGFjawogCXRlc3Q="
server.lastRequest.headers.get("lambda-runtime-aws-request-id") == "lambda-request-123"

cleanup:
server.close()
Expand Down
9 changes: 5 additions & 4 deletions internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -406,9 +406,9 @@ default SpanBuilder singleSpanBuilder(CharSequence spanName) {

CallbackProvider getUniversalCallbackProvider();

AgentSpanContext notifyExtensionStart(Object event);
AgentSpanContext notifyExtensionStart(Object event, String lambdaRequestId);

void notifyExtensionEnd(AgentSpan span, Object result, boolean isError);
void notifyExtensionEnd(AgentSpan span, Object result, boolean isError, String lambdaRequestId);

AgentDataStreamsMonitoring getDataStreamsMonitoring();

Expand Down Expand Up @@ -653,12 +653,13 @@ public EndpointTracker onRootSpanStarted(AgentSpan root) {
}

@Override
public AgentSpanContext notifyExtensionStart(Object event) {
public AgentSpanContext notifyExtensionStart(Object event, String lambdaRequestId) {
return null;
}

@Override
public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {}
public void notifyExtensionEnd(
AgentSpan span, Object result, boolean isError, String lambdaRequestId) {}

@Override
public AgentDataStreamsMonitoring getDataStreamsMonitoring() {
Expand Down
Loading