Fix distributed tracing (#28)

* Create span and set parent manually

* Close the scope

* Set trace and parent IDs correctly

* Skipped 0.0.10 by accident

* improve language around default value of DD_TRACE_ENABLED

Author: agocs

build.gradle

@@ -46,7 +46,7 @@ sourceCompatibility = 1.8
 targetCompatibility = 1.8
 
 group = 'com.datadoghq'
-version= '0.0.9'
+version= '0.0.10'
 
 allprojects {
     repositories {

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

@@ -1,5 +1,11 @@
 package com.datadoghq.datadog_lambda_java;
 
+import io.opentracing.Scope;
+import io.opentracing.SpanContext;
+import io.opentracing.Tracer;
+import io.opentracing.Tracer.SpanBuilder;
+import io.opentracing.propagation.Format.Builtin;
+import io.opentracing.propagation.TextMapAdapter;
 import java.io.IOException;
 import java.net.URL;
 import java.net.URLConnection;
@@ -28,8 +34,10 @@ public class DDLambda {
     private String MDC_TRACE_CONTEXT_FIELD = "dd.trace_context";
     private String JSON_TRACE_ID = "dd.trace_id";
     private String JSON_SPAN_ID = "dd.span_id";
+    private String TRACE_ENABLED_ENV = "DD_TRACE_ENABLED";
     private Tracing tracing;
     private boolean enhanced = true;
+    private Scope tracingScope;
 
     /**
      * Create a new DDLambda instrumenter given some Lambda context
@@ -41,7 +49,7 @@ public DDLambda(Context cxt) {
         this.enhanced = checkEnhanced();
         recordEnhanced(INVOCATION, cxt);
         addTraceContextToMDC();
-        addTagsToSpan(cxt);
+        startSpan(null, cxt);
     }
 
     /**
@@ -55,7 +63,7 @@ protected DDLambda(Context cxt, String xrayTraceInfo) {
         this.enhanced = checkEnhanced();
         recordEnhanced(INVOCATION, cxt);
         addTraceContextToMDC();
-        addTagsToSpan(cxt);
+        startSpan(null, cxt);
     }
 
     /**
@@ -71,7 +79,7 @@ public DDLambda(APIGatewayProxyRequestEvent req, Context cxt) {
         this.tracing = new Tracing(req);
         this.tracing.submitSegment();
         addTraceContextToMDC();
-        addTagsToSpan(cxt);
+        startSpan(req.getHeaders(), cxt);
     }
 
     /**
@@ -87,7 +95,7 @@ public DDLambda(APIGatewayV2ProxyRequestEvent req, Context cxt) {
         this.tracing = new Tracing(req);
         this.tracing.submitSegment();
         addTraceContextToMDC();
-        addTagsToSpan(cxt);
+        startSpan(req.getHeaders(), cxt);
     }
 
     /**
@@ -103,10 +111,69 @@ public DDLambda(Headerable req, Context cxt) {
         this.tracing = new Tracing(req);
         this.tracing.submitSegment();
         addTraceContextToMDC();
-        addTagsToSpan(cxt);
+        startSpan(req.getHeaders(), cxt);
     }
 
-    private void addTagsToSpan(Context cxt) {
+    /**
+     * startSpan is called by the DDLambda constructors. If there is a dd-agent-java present, it will start
+     * a span. If not, startSpan is a noop.
+     * @param headers are the headers from the Lambda request. If Headers are null or empty, distributed tracing
+     *                is impossible but a span will still start.
+     * @param cxt is the Lambda Context passed to the Handler function.
+     */
+    private void startSpan(Map<String,String> headers, Context cxt){
+        //If the user has not set DD_TRACE_ENABLED=true, don't start a span
+        if(!checkTraceEnabled()){
+            return;
+        }
+
+        String functionName = "";
+        if (cxt != null){
+            functionName = cxt.getFunctionName();
+        }
+
+        //Get the Datadog tracer, if it exists
+        Tracer tracer = GlobalTracer.get();
+        //Datadog tracer will be able to extract datadog trace headers from an incoming request
+        SpanContext parentContext = tracer.extract(Builtin.HTTP_HEADERS, new TextMapAdapter(headers));
+
+        SpanBuilder spanBuilder = tracer.buildSpan(functionName).asChildOf(parentContext);
+        spanBuilder = addDDTags(spanBuilder, cxt);
+        Span thisSpan = spanBuilder.start();
+
+        //Hang on to the scope, we'll need it later to close.
+        this.tracingScope = tracer.activateSpan(thisSpan);
+    }
+
+    /**
+     * Finish the active span. If you have installed the dd-trace-java Lambda
+     * layer, you MUST call DDLambda.finish() at the end of your Handler
+     * in order to finish spans.
+     */
+    public void finish(){
+        Span span = GlobalTracer.get().activeSpan();
+
+        if (this.tracingScope == null){
+            DDLogger.getLoggerImpl().debug("Unable to close tracing scope because it is null.");
+            return;
+        }
+        this.tracingScope.close();
+
+        if (span != null) {
+            span.finish();
+        } else {
+            DDLogger.getLoggerImpl().debug("Unable to finish span because it is null.");
+            return;
+        }
+    }
+
+    /**
+     * addDDTags adds Datadog tags to a span's tags.
+     * @param spanBuilder is the SpanBuilder being used to build the span to which these tags will be applied
+     * @param cxt is the Lambda Context that contains the information necessary to build these tags
+     * @return a SpanBuilder with the necessary tags.
+     */
+    private SpanBuilder addDDTags(SpanBuilder spanBuilder, Context cxt) {
         String requestId = "";
         String functionName = "";
         String functionArn = "";
@@ -116,17 +183,19 @@ private void addTagsToSpan(Context cxt) {
             functionName = cxt.getFunctionName();
             functionArn = santitizeFunctionArn(cxt.getInvokedFunctionArn());
             functionVersion = cxt.getFunctionVersion();
+        } else {
+            return spanBuilder;
         }
-        Span span = GlobalTracer.get().activeSpan();
-        if (span != null) {
-            span.setTag("request_id", requestId);
-            span.setTag("service", "aws.lambda");
-            span.setTag("function_arn", functionArn);
-            span.setTag("cold_start", ColdStart.getColdStart(cxt));
-            span.setTag("datadog_lambda", BuildConfig.datadog_lambda_version);
-            span.setTag("resource_names", functionName);
-            span.setTag("function_version", functionVersion);
+        if (spanBuilder != null) {
+            spanBuilder.withTag("request_id", requestId);
+            spanBuilder.withTag("service", "aws.lambda");
+            spanBuilder.withTag("function_arn", functionArn);
+            spanBuilder.withTag("cold_start", ColdStart.getColdStart(cxt));
+            spanBuilder.withTag("datadog_lambda", BuildConfig.datadog_lambda_version);
+            spanBuilder.withTag("resource_names", functionName);
+            spanBuilder.withTag("function_version", functionVersion);
         }
+        return spanBuilder;
     }
 
     protected String santitizeFunctionArn(String functionArn){
@@ -169,6 +238,22 @@ protected boolean checkEnhanced() {
         return true;
     }
 
+    /**
+     * Check to see if the user has set DD_TRACE_ENABLED
+     * @return true if DD_TRACE_ENABLED has been set to "true" (or "TRUE" or "tRuE" or ...), false otherwise
+     */
+    protected boolean checkTraceEnabled(){
+        String sysTraceEnabled = System.getenv(TRACE_ENABLED_ENV);
+        if (sysTraceEnabled == null) {
+            return false;
+        }
+
+        if (sysTraceEnabled.toLowerCase().equals("true")){
+            return true;
+        }
+        return false;
+    }
+
     /**
      * metric allows the user to record their own custom metric that will be sent to Datadog.
      *

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

@@ -11,7 +11,7 @@
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
 import com.google.gson.Gson;
 
-public class Tracing {
+public class Tracing{
 
     protected DDTraceContext cxt;
     protected XRayTraceContext xrt;