Add span error metadata to Datadog Lambda Extension request headers (#…

…6061) attach span error metadata to request headers
DataDog · Nov 20, 2023 · d3b600a · d3b600a
1 parent d18bc66
commit d3b600a
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 1 deletion.
diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java
@@ -6,12 +6,15 @@
 import com.squareup.moshi.JsonAdapter;
 import com.squareup.moshi.Moshi;
 import datadog.trace.api.DDSpanId;
+import datadog.trace.api.DDTags;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.core.propagation.ExtractedContext;
 import datadog.trace.core.propagation.PropagationTags;
 import java.io.ByteArrayInputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
 import okhttp3.MediaType;
 import okhttp3.OkHttpClient;
 import okhttp3.Request;
@@ -37,6 +40,9 @@ public class LambdaHandler {
   private static final String DATADOG_SPAN_ID = "x-datadog-span-id";
   private static final String DATADOG_SAMPLING_PRIORITY = "x-datadog-sampling-priority";
   private static final String DATADOG_INVOCATION_ERROR = "x-datadog-invocation-error";
+  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 DATADOG_TAGS_KEY = "x-datadog-tags";
 
   private static final String START_INVOCATION = "/lambda/start-invocation";
@@ -121,12 +127,30 @@ public static boolean notifyEndInvocation(AgentSpan span, Object result, boolean
     Request.Builder builder =
         new Request.Builder()
             .url(EXTENSION_BASE_URL + END_INVOCATION)
-            .addHeader(DATADOG_META_LANG, "java")
             .addHeader(DATADOG_TRACE_ID, span.getTraceId().toString())
             .addHeader(DATADOG_SPAN_ID, DDSpanId.toString(span.getSpanId()))
             .addHeader(DATADOG_SAMPLING_PRIORITY, span.getSamplingPriority().toString())
             .addHeader(DATADOG_META_LANG, "java")
             .post(body);
+
+    Object errorMessage = span.getTag(DDTags.ERROR_MSG);
+    if (errorMessage != null) {
+      builder.addHeader(DATADOG_INVOCATION_ERROR_MSG, errorMessage.toString());
+    }
+
+    Object errorType = span.getTag(DDTags.ERROR_TYPE);
+    if (errorType != null) {
+      builder.addHeader(DATADOG_INVOCATION_ERROR_TYPE, errorType.toString());
+    }
+
+    Object errorStack = span.getTag(DDTags.ERROR_STACK);
+    if (errorStack != null) {
+      String encodedErrStack =
+          Base64.getEncoder()
+              .encodeToString(errorStack.toString().getBytes(StandardCharsets.UTF_8));
+      builder.addHeader(DATADOG_INVOCATION_ERROR_STACK, encodedErrStack);
+    }
+
     if (isError) {
       builder.addHeader(DATADOG_INVOCATION_ERROR, "true");
     }

diff --git a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy
@@ -2,6 +2,7 @@ package datadog.trace.lambda
 
 import datadog.trace.api.Config
 import datadog.trace.api.DDSpanId
+import datadog.trace.api.DDTags
 import datadog.trace.api.DDTraceId
 import datadog.trace.core.propagation.PropagationTags
 import datadog.trace.core.test.DDCoreSpecification
@@ -161,6 +162,40 @@ class LambdaHandlerTest extends DDCoreSpecification {
     false    | null        | "12345"      | false
   }
 
+  def "test end invocation success with error metadata"() {
+    given:
+    def server = httpServer {
+      handlers {
+        post("/lambda/end-invocation") {
+          response
+            .status(200)
+            .send()
+        }
+      }
+    }
+    LambdaHandler.setExtensionBaseUrl(server.address.toString())
+    DDSpan span = Mock(DDSpan) {
+      getTraceId() >> DDTraceId.from("1234")
+      getSpanId() >> DDSpanId.from("5678")
+      getSamplingPriority() >> 2
+      getTag(DDTags.ERROR_MSG) >> "custom error message"
+      getTag(DDTags.ERROR_TYPE) >> "java.lang.Throwable"
+      getTag(DDTags.ERROR_STACK) >> "errorStack\n \ttest"
+    }
+
+    when:
+    LambdaHandler.notifyEndInvocation(span, {}, true)
+
+    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="
+
+    cleanup:
+    server.close()
+  }
+
   def "test moshi toJson SQSEvent"() {
     given:
     def myEvent = new SQSEvent()