diff --git a/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index 98efeb6d9e3..b281639d7ff 100644 --- a/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -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; @@ -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, + @Advice.Argument(1) final Object out, + @Advice.Argument(2) final Context awsContext, @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; } diff --git a/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy b/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy index 5982e5d18ee..4efa77ff176 100644 --- a/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy +++ b/dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy @@ -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() { @@ -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) { @@ -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) @@ -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