add request-id to lambda start/end invocation calls

Author: ojproductions

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java

@@ -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;
@@ -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) {
@@ -120,8 +121,14 @@ static void exit(
         if (throwable != null) {
           span.addThrowable(throwable);
         }
+        String lambdaRequestId = awsContext.getAwsRequestId();
+
+        // QUESTION would it be better to get request id from span tag instead? So we dont have to
+        // change method signature
+        // E.G:
+        // String lambdaRequestIdFromSpan = (String) span.getTag("request_id");
         span.finish();
-        AgentTracer.get().notifyExtensionEnd(span, result, null != throwable);
+        AgentTracer.get().notifyExtensionEnd(span, result, null != throwable, lambdaRequestId);
       } finally {
         scope.close();
       }

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -1194,13 +1194,13 @@ public void closeActive() {
   }
 
   @Override
-  public AgentSpanContext notifyExtensionStart(Object event) {
-    return LambdaHandler.notifyStartInvocation(this, event);
+  public AgentSpanContext notifyExtensionStart(Object event, String requestId) {
+    return LambdaHandler.notifyStartInvocation(this, event, requestId);
   }
 
   @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 requestId) {
+    LambdaHandler.notifyEndInvocation(span, result, isError, requestId);
   }
 
   @Override

dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java

@@ -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";
@@ -72,14 +73,16 @@ public class LambdaHandler {
 
   private static String EXTENSION_BASE_URL = "http://127.0.0.1:8124";
 
-  public static AgentSpanContext notifyStartInvocation(CoreTracer tracer, Object event) {
+  public static AgentSpanContext notifyStartInvocation(
+      CoreTracer tracer, Object event, String requestId) {
     RequestBody body = RequestBody.create(jsonMediaType, writeValueAsString(event));
     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, requestId)
                     .post(body)
                     .build())
             .execute()) {
@@ -99,7 +102,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 requestId) {
     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");
@@ -113,6 +117,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, requestId)
             .post(body);
 
     Object errorMessage = span.getTag(DDTags.ERROR_MSG);

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentTracer.java

@@ -406,9 +406,9 @@ default SpanBuilder singleSpanBuilder(CharSequence spanName) {
 
     CallbackProvider getUniversalCallbackProvider();
 
-    AgentSpanContext notifyExtensionStart(Object event);
+    AgentSpanContext notifyExtensionStart(Object event, String requestId);
 
-    void notifyExtensionEnd(AgentSpan span, Object result, boolean isError);
+    void notifyExtensionEnd(AgentSpan span, Object result, boolean isError, String requestId);
 
     AgentDataStreamsMonitoring getDataStreamsMonitoring();
 
@@ -653,12 +653,13 @@ public EndpointTracker onRootSpanStarted(AgentSpan root) {
     }
 
     @Override
-    public AgentSpanContext notifyExtensionStart(Object event) {
+    public AgentSpanContext notifyExtensionStart(Object event, String requestId) {
       return null;
     }
 
     @Override
-    public void notifyExtensionEnd(AgentSpan span, Object result, boolean isError) {}
+    public void notifyExtensionEnd(
+        AgentSpan span, Object result, boolean isError, String requestId) {}
 
     @Override
     public AgentDataStreamsMonitoring getDataStreamsMonitoring() {