Pass Lambda Request ID to Extension (#8814)

* update LambdaHandlerInstrumentation.java

* update unit test

Author: nhulston

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;
     }

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