Add trace log correlation (#21)

* adjust version code and add version tag to readme

* Add trace/log correlation info

* Add instructions for how to add trace ids to your logs

* Made some decisions about some log levels

* Add span ID, add more detail in the README

* Update README and version number

* Explicitly mention cloning the lambda pipeline

Author: agocs

README.md

@@ -38,7 +38,7 @@ Once [installed](#installation), you should be able to submit custom metrics fro
 
 Check out the instructions for [submitting custom metrics from AWS Lambda functions](https://docs.datadoghq.com/integrations/amazon_lambda/?tab=java#custom-metrics).
 
-## Tracing
+## Distributed Tracing
 
 Wrap your outbound HTTP requests with trace headers to see your lambda in context in APM.
 The Lambda Java Client Library provides instrumented HTTP connection objects as well as helper methods for
@@ -168,6 +168,59 @@ public class Handler implements RequestHandler<APIGatewayV2ProxyRequestEvent, AP
 }
 ```
 
+### Trace/Log Correlations
+
+In order to correlate your traces with your logs, you must inject the trace context
+into your log messages. We've added the these into the slf4j MDC under the key `dd.trace_context`
+and provided convenience methods to get it automatically. The trace context is added to the MDC as a side
+effect of instantiating any `new DDLambda(...)`.
+
+This is an example trace context: `[dd.trace_id=3371139772690049666 dd.span_id=13006875827893840236]`
+
+#### JSON Logs
+
+If you are using JSON logs, add the trace ID and span ID to each log message with the keys 
+`dd.trace_id` and `dd.span_id` respectively. To get a map containing trace and span IDs,
+ call `DDLambda.getTraceContext()`. Union this map with the JSON data being logged.
+
+#### Plain text logs
+
+If you are using plain text logs, then you must create a new [Parser](https://docs.datadoghq.com/logs/processing/parsing/?tab=matcher)
+by cloning the existing Lambda Pipeline. The new parser can extract the trace context from the correct position in the logs. 
+Use the helper `_trace_context` to extract the trace context. For example, if your log line looked like:
+
+```
+INFO 2020-11-11T14:00:00Z LAMBDA_REQUEST_ID [dd.trace_id=12345 dd.span_id=67890] This is a log message
+```
+
+Then your parser rule would look like:
+
+```
+my_custom_rule \[%{word:level}\]?\s+%{_timestamp}\s+%{notSpace:lambda.request_id}%{_trace_context}?.*
+```
+
+#### Log4j / SLF4J
+
+We have added the Trace ID into the slf4j MDC under the key `dd.trace_context`. That can be accessed
+using the `%X{dd.trace_context}` operator. Here is an example `log4j.properties`: 
+
+```
+log = .
+log4j.rootLogger = DEBUG, LAMBDA
+
+log4j.appender.LAMBDA=com.amazonaws.services.lambda.runtime.log4j.LambdaAppender
+log4j.appender.LAMBDA.layout=org.apache.log4j.PatternLayout
+log4j.appender.LAMBDA.layout.conversionPattern=%d{yyyy-MM-dd HH:mm:ss} %X{dd.trace_context} %-5p %c:%L - %m%n
+```
+
+would result in log lines looking like `2020-11-13 19:21:53 [dd.trace_id=1168910694192328743 dd.span_id=3204959397041471598] INFO  com.serverless.Handler:20 - Test Log Message`
+
+#### Other logging solutions
+
+If you are using a different logging solution, the trace ID can be accessed using the method
+`DDLambda.getTraceContextString()`. That returns your trace ID as a string that can be added
+to any log message.
+
 ## Opening Issues
 
 If you encounter a bug with this package, we want to hear about it. Before opening a new issue, 

build.gradle

@@ -28,18 +28,20 @@ dependencies {
     implementation 'org.apache.httpcomponents:httpclient:4.5.11'
     implementation 'com.squareup.okhttp3:okhttp:4.4.0'
     implementation 'org.apache.commons:commons-lang3:3.10'
-
+    implementation 'org.slf4j:slf4j-log4j12:1.7.28'
+    implementation 'org.jetbrains:annotations:15.0'
 
     // Use JUnit test framework
     testImplementation 'junit:junit:4.12'
-    implementation 'org.jetbrains:annotations:15.0'
+    testImplementation 'org.apache.logging.log4j:log4j-api:2.13.3'
+    testImplementation 'org.apache.logging.log4j:log4j-core:2.13.3'
 }
 
 sourceCompatibility = 1.8
 targetCompatibility = 1.8
 
 group = 'com.datadoghq'
-version= '0.0.6'
+version= '0.0.7'
 
 allprojects {
     repositories {

src/main/java/com/datadoghq/datadog_lambda_java/DDLambda.java

@@ -1,5 +1,6 @@
 package com.datadoghq.datadog_lambda_java;
 
+import org.slf4j.MDC;
 import java.io.IOException;
 import java.net.URL;
 import java.net.URLConnection;
@@ -20,6 +21,7 @@ public class DDLambda {
     private String ENHANCED_PREFIX = "aws.lambda.enhanced.";
     private String INVOCATION = "invocations";
     private String ERROR = "errors";
+    private String MDC_TRACE_CONTEXT_FIELD = "dd.trace_context";
     private Tracing tracing;
     private boolean enhanced = true;
 
@@ -31,6 +33,19 @@ public DDLambda(Context cxt){
         this.tracing = new Tracing();
         this.enhanced = checkEnhanced();
         recordEnhanced(INVOCATION, cxt);
+        MDC.put(MDC_TRACE_CONTEXT_FIELD, getTraceContextString());
+    }
+
+    /**
+     * Testing only: create a DDLambda instrumenter with a given context and xrayTraceInfo
+     * @param cxt Enhanced Metrics pulls information from the Lambda context.
+     * @param xrayTraceInfo This would normally be the contents of the "_X_AMZN_TRACE_ID" env var
+     */
+    protected DDLambda(Context cxt, String xrayTraceInfo){
+        this.tracing = new Tracing(xrayTraceInfo);
+        this.enhanced = checkEnhanced();
+        recordEnhanced(INVOCATION, cxt);
+        MDC.put(MDC_TRACE_CONTEXT_FIELD, getTraceContextString());
     }
 
     /**
@@ -44,6 +59,7 @@ public DDLambda(APIGatewayProxyRequestEvent req, Context cxt){
         recordEnhanced(INVOCATION, cxt);
         this.tracing = new Tracing(req);
         this.tracing.submitSegment();
+        MDC.put(MDC_TRACE_CONTEXT_FIELD, getTraceContextString());
     }
 
     /**
@@ -57,6 +73,7 @@ public DDLambda(APIGatewayV2ProxyRequestEvent req, Context cxt){
         recordEnhanced(INVOCATION, cxt);
         this.tracing = new Tracing(req);
         this.tracing.submitSegment();
+        MDC.put(MDC_TRACE_CONTEXT_FIELD, getTraceContextString());
     }
 
     /**
@@ -70,6 +87,7 @@ public DDLambda(Headerable req, Context cxt){
         recordEnhanced(INVOCATION, cxt);
         this.tracing = new Tracing(req);
         this.tracing.submitSegment();
+        MDC.put(MDC_TRACE_CONTEXT_FIELD, getTraceContextString());
     }
 
     private boolean checkEnhanced(){
@@ -212,4 +230,37 @@ public Request addTraceHeaders(Request request){
 
         return rb.build();
     }
+
+    /**
+     * Get the trace context for trace/log correlation. Inject this into your logs in order to correlate logs with traces.
+     * @return a map of the current trace context
+     */
+    public Map<String,String> getTraceContext(){
+        if (this.tracing == null){
+            DDLogger.getLoggerImpl().debug("No tracing context; unable to get Trace ID");
+            return null;
+        }
+        return this.tracing.getLogCorrelationTraceAndSpanIDsMap();
+    }
+
+    /**
+     * Get the trace context in string form. Inject this into your logs in order to correlate logs with traces.
+     * @return a string representation of the current trace context
+     */
+    public String getTraceContextString(){
+        Map<String,String> traceInfo = getTraceContext();
+        if (traceInfo == null){
+            DDLogger.getLoggerImpl().debug("No Trace/Log correlation IDs returned");
+            return "";
+        }
+
+        String traceID = traceInfo.get(this.tracing.TRACE_ID_KEY);
+        String spanID = traceInfo.get(this.tracing.SPAN_ID_KEY);
+        return formatTraceContext(this.tracing.TRACE_ID_KEY, traceID, this.tracing.SPAN_ID_KEY, spanID);
+    }
+
+    private String formatTraceContext(String traceKey, String trace, String spanKey, String span){
+        return String.format("[%s=%s %s=%s]", traceKey, trace, spanKey, span);
+    }
+
 }

src/main/java/com/datadoghq/datadog_lambda_java/Tracing.java

@@ -16,6 +16,9 @@ public class Tracing {
     protected DDTraceContext cxt;
     protected XRayTraceContext xrt;
 
+    protected String TRACE_ID_KEY = "dd.trace_id";
+    protected String SPAN_ID_KEY = "dd.span_id";
+
     public Tracing(){
         this.xrt = new XRayTraceContext();
     }
@@ -35,6 +38,14 @@ public Tracing(Headerable req){
         this.xrt = new XRayTraceContext();
     }
 
+    /**
+     * Test constructor that can take a dummy _X_AMZN_TRACE_ID value
+     * @param xrayTraceInfo
+     */
+    protected Tracing(String xrayTraceInfo){
+        this.xrt = new XRayTraceContext(xrayTraceInfo);
+    }
+
 
     public DDTraceContext getDDContext() {
         if (this.cxt == null) {
@@ -50,6 +61,31 @@ public XRayTraceContext getXrayContext() {
         return xrt;
     }
 
+    public Map<String,String> getLogCorrelationTraceAndSpanIDsMap(){
+        if (this.cxt != null){
+            String traceID = this.cxt.getTraceID();
+            String spanID = this.cxt.getParentID();
+            Map<String, String> out  = new HashMap<String, String>();
+            out.put(TRACE_ID_KEY, traceID);
+            out.put(SPAN_ID_KEY, spanID);
+            return out;
+        }
+        if (this.xrt != null){
+            String traceID = this.xrt.getAPMTraceID();
+            String spanID = this.xrt.getAPMParentID();
+            Map<String, String> out  = new HashMap<String, String>();
+            out.put(TRACE_ID_KEY, traceID);
+            out.put(SPAN_ID_KEY, spanID);
+            return out;
+        }
+        DDLogger.getLoggerImpl().debug("No DD trace context or XRay trace context set!");
+        return null;
+    }
+
+    private String formatLogCorrelation(String trace, String span){
+        return String.format("[dd.trace_id=%s dd.span_id=%s]", trace, span);
+    }
+
     private static DDTraceContext populateDDContext(Map<String,String> headers){
         DDTraceContext ctx = null;
         try{
@@ -317,7 +353,7 @@ class XRayTraceContext{
     public XRayTraceContext(){
         //Root=1-5e41a79d-e6a0db584029dba86a594b7e;Parent=8c34f5ad8f92d510;Sampled=1
         String traceId = System.getenv("_X_AMZN_TRACE_ID");
-        if (traceId == null){
+        if (traceId == null || traceId == ""){
             DDLogger.getLoggerImpl().debug("Unable to find _X_AMZN_TRACE_ID");
             return;
         }
@@ -337,6 +373,27 @@ public XRayTraceContext(){
         this.traceIdHeader = traceId;
     }
 
+    /**
+     * Test constructor that can take a dummy _X_AMZN_TRACE_ID value rather than reading from env vars
+     * @param traceId
+     */
+    protected XRayTraceContext(String traceId){
+        String[] traceParts = traceId.split(";");
+        if(traceParts.length != 3){
+            DDLogger.getLoggerImpl().error("Malformed _X_AMZN_TRACE_ID value: "+ traceId);
+            return;
+        }
+
+        try {
+            this.traceId = traceParts[0].split("=")[1];
+            parentId = traceParts[1].split("=")[1];
+        } catch (Exception e){
+            DDLogger.getLoggerImpl().error("Malformed _X_AMZN_TRACE_ID value: "+ traceId);
+            return;
+        }
+        this.traceIdHeader = traceId;
+    }
+
     public void setTraceIdHeader(String traceIdHeader) {
         this.traceIdHeader = traceIdHeader;
     }
@@ -375,4 +432,35 @@ public String getAPMParentID(){
             return null;
         }
     }
+
+    public String getAPMTraceID(){
+        //trace ID looks like 1-5e41a79d-e6a0db584029dba86a594b7e
+        String bigid = "";
+        try {
+            bigid = this.traceId.split("-")[2];
+        }
+        catch (ArrayIndexOutOfBoundsException | NullPointerException ai){
+            DDLogger.getLoggerImpl().debug("Unexpected format for the trace ID. Unable to parse it. " + this.traceId);
+            return "";
+        }
+
+        //just to verify
+        if (bigid.length() != 24){
+            DDLogger.getLoggerImpl().debug("Got an unusual traceid from x-ray. Unable to convert that to an APM id. " + this.traceId);
+            return "";
+        }
+
+        String last16 = bigid.substring(bigid.length()-16); // should be the last 16 characters of the big id
+
+        Long parsed = 0L;
+        try {
+            parsed = Long.parseUnsignedLong(last16, 16); //unsigned because parseLong throws a numberformatexception at anything greater than 0x7FFFF...
+        }
+        catch (NumberFormatException ne){
+            DDLogger.getLoggerImpl().debug("Got a NumberFormatException trying to parse the traceID. Unable to convert to an APM id. " + this.traceId);
+            return "";
+        }
+        parsed = parsed & 0x7FFFFFFFFFFFFFFFL; //take care of that pesky first bit...
+        return parsed.toString();
+    }
 }

src/test/java/com/datadoghq/datadog_lambda_java/TraceLogCorrelationTest.java

@@ -0,0 +1,52 @@
+package com.datadoghq.datadog_lambda_java;
+
+import java.io.ByteArrayOutputStream;
+import java.io.PrintStream;
+import java.util.Map;
+import org.junit.Assert;
+import org.junit.Test;
+import org.apache.log4j.Logger;
+import org.apache.log4j.ConsoleAppender;
+import org.apache.log4j.WriterAppender;
+import org.apache.log4j.PatternLayout;
+
+
+public class TraceLogCorrelationTest {
+    private final ByteArrayOutputStream outContent = new ByteArrayOutputStream();
+    private final PrintStream originalErr = System.err;
+
+    @Test
+    public void TestWithLog4J(){
+        Logger logger = Logger.getLogger(TraceLogCorrelationTest.class.getName());
+        PatternLayout layout = new PatternLayout("%d{ISO8601} %X{dd.trace_context} [%t] %-5p %c %x - %m%n");
+        logger.addAppender(new ConsoleAppender((layout)));
+
+        logger.info("Test log message");
+
+        EnhancedMetricTest.MockContext mc = new EnhancedMetricTest.MockContext();
+
+        //try to grab the contents of StdErr
+        logger.addAppender(new WriterAppender(layout, outContent));
+        DDLambda ddl = new DDLambda(mc, "Root=1-5fa98677-5dda00b12ec8b1c31f61aa82;Parent=b481b4a6ef3a296c;Sampled=1");
+        outContent.reset();
+        logger.info("Test log message 2");
+        logger.removeAllAppenders();
+
+        String logOutput  = outContent.toString();
+        String[] logFields = logOutput.split(" ");
+        Assert.assertEquals("[dd.trace_id=3371139772690049666", logFields[2]);
+        Assert.assertEquals("dd.span_id=13006875827893840236]", logFields[3]);
+    }
+
+    @Test
+    public void TestGetTraceLogCorrelationId(){
+        EnhancedMetricTest.MockContext mc = new EnhancedMetricTest.MockContext();
+        DDLambda ddl = new DDLambda(mc, "Root=1-5fa98677-5dda00b12ec8b1c31f61aa82;Parent=b481b4a6ef3a296c;Sampled=1");
+
+        Map<String,String> traceInfo = ddl.getTraceContext();
+        Assert.assertEquals("3371139772690049666", traceInfo.get("dd.trace_id"));
+        Assert.assertEquals("13006875827893840236", traceInfo.get("dd.span_id"));
+    }
+
+}
+