Skip to content

Pass Lambda Request ID to Extension #8814

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

Merged
merged 2 commits into from
May 13, 2025
Merged

Pass Lambda Request ID to Extension #8814

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fa65a68
update LambdaHandlerInstrumentation.java
nhulston May 13, 2025
bdeda1d
update unit test
nhulston May 13, 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
9 changes: 7 additions & 2 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 @@ -13,6 +13,7 @@
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;

import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.RequestHandler;
import com.google.auto.service.AutoService;
import datadog.trace.agent.tooling.Instrumenter;
Expand Down Expand Up @@ -79,20 +80,24 @@ public static class ExtensionCommunicationAdvice {
@OnMethodEnter
static AgentScope enter(
@This final Object that,
@Advice.Argument(0) final Object event,
@Advice.Argument(0) final Object in,
Copy link

@duncanista duncanista May 13, 2025

Choose a reason for hiding this comment

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

Should we keep calling it event?

Copy link
Contributor Author

@nhulston nhulston May 13, 2025

Choose a reason for hiding this comment

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

The 2-arg handler takes args (InputType input, Context context) and the 3-arg handler tags args (InputStream input, OutputStream output, Context context) so in/out is better than event here

Copy link

@duncanista duncanista May 13, 2025

Choose a reason for hiding this comment

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

Ah, got it, thanks!

Copy link

@duncanista duncanista May 13, 2025

Choose a reason for hiding this comment

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

Maybe add a comment about the different types of handlers? Is it worth adding here?

Copy link
Contributor Author

@nhulston nhulston May 13, 2025

Choose a reason for hiding this comment

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

Oops I already merged, but the PR from https://github.com/DataDog/dd-trace-java/pull/8264/files added some comments explaining this

@Advice.Argument(1) final Object out,
@Advice.Argument(2) final Context awsContext,
Comment on lines -82 to +85
Copy link
Contributor Author

@nhulston nhulston May 13, 2025

Choose a reason for hiding this comment

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

There's a Lambda handler method header with 2 arguments and 3 arguments. The 2-arg header calls the 3-arg header under-the-hood, so as long as we instrument the 3-arg header, we support all cases.

We were already instrumenting the 3-arg header thanks to #8264;
this change just specifies all 3 arguments but doesn't change what we're actually instrumenting.

So TLDR this is non-breaking

@Origin("#m") final String methodName) {

if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) {
return null;
}

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

final AgentScope scope = activateSpan(span);
return scope;
}
Expand Down
13 changes: 11 additions & 2 deletions ...nstrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import datadog.trace.agent.test.naming.VersionedNamingTestBase
import java.nio.charset.StandardCharsets
import com.amazonaws.services.lambda.runtime.Context

abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase {
def requestId = "test-request-id"

@Override
String service() {
Expand All @@ -12,7 +14,10 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
when:
def input = new ByteArrayInputStream(StandardCharsets.UTF_8.encode("Hello").array())
def output = new ByteArrayOutputStream()
new HandlerStreaming().handleRequest(input, output, null)
def ctx = Stub(Context) {
getAwsRequestId() >> requestId
}
new HandlerStreaming().handleRequest(input, output, ctx)

then:
assertTraces(1) {
Expand All @@ -29,7 +34,10 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
when:
def input = new ByteArrayInputStream(StandardCharsets.UTF_8.encode("Hello").array())
def output = new ByteArrayOutputStream()
new HandlerStreamingWithError().handleRequest(input, output, null)
def ctx = Stub(Context) {
getAwsRequestId() >> requestId
}
new HandlerStreamingWithError().handleRequest(input, output, ctx)

then:
thrown(Error)
Expand All @@ -39,6 +47,7 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
operationName operation()
errored true
tags {
tag "request_id", requestId
tag "error.type", "java.lang.Error"
tag "error.message", "Some error"
tag "error.stack", String
Expand Down