Add DD_TRACE_ENABLED and DD_MERGE_XRAY_TRACES options

Author: DarcyRaynerDD

Dockerfile

@@ -10,6 +10,7 @@ WORKDIR /build
 # Install datadog_lambda and dependencies from local
 COPY . .
 RUN pip install . -t ./python/lib/$runtime/site-packages
+RUN pip install --find-links=https://s3.amazonaws.com/pypi.datadoghq.com/trace-dev/index.html ddtrace==0.35.1.dev5+g1e11b2dd -t ./python/lib/$runtime/site-packages
 
 # Remove *.pyc files
 RUN find ./python/lib/$runtime/site-packages -name \*.pyc -delete

datadog_lambda/trace_wrapper.py

@@ -4,22 +4,25 @@
 # Copyright 2019 Datadog, Inc.
 
 import logging
+import os
 
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.lambda_launcher import LambdaContext
 
-from ddtrace import patch, tracer
 from datadog_lambda.constants import (
     SamplingPriority,
     TraceHeader,
     XraySubsegment,
     Source,
 )
-from datadog_lambda.trace_wrapper import trace_wrapper
+from ddtrace import tracer, patch
 
 logger = logging.getLogger(__name__)
 
 dd_trace_context = {}
+dd_native_tracing_enabled = (
+    os.environ.get("DD_TRACE_ENABLED", "false").lower() == "true"
+)
 
 
 def _convert_xray_trace_id(xray_trace_id):
@@ -60,11 +63,26 @@ def _get_xray_trace_context():
     }
 
 
+def _get_dd_trace_native_context():
+    span = tracer.current_span()
+    if not span:
+        return None
+
+    parent_id = span.context.span_id
+    trace_id = span.context.trace_id
+    return {
+        "parent_id": str(parent_id),
+        "trace_id": str(trace_id),
+        "sampling_priority": SamplingPriority.AUTO_KEEP,
+        "source": Source.DDTRACE,
+    }
+
+
 def _context_obj_to_headers(obj):
     return {
-        TraceHeader.TRACE_ID: str(obj.get("trace_id")),
-        TraceHeader.PARENT_ID: str(obj.get("parent_id")),
-        TraceHeader.SAMPLING_PRIORITY: str(obj.get("sampling_priority")),
+        TraceHeader.TRACE_ID: str(obj.get("trace-id")),
+        TraceHeader.PARENT_ID: str(obj.get("parent-id")),
+        TraceHeader.SAMPLING_PRIORITY: str(obj.get("sampling-priority")),
     }
 
 
@@ -122,10 +140,11 @@ def get_dd_trace_context():
     """
     global dd_trace_context
 
-    native_trace_context = trace_wrapper.trace_context
-    if native_trace_context:
-        logger.info("get_dd_trace_context using dd-trace context")
-        return _context_obj_to_headers(native_trace_context)
+    if dd_native_tracing_enabled:
+        native_trace_context = _get_dd_trace_native_context()
+        if native_trace_context is not None:
+            logger.info("get_dd_trace_context using dd-trace context")
+            return _context_obj_to_headers(native_trace_context)
 
     try:
         trace_headers = _context_obj_to_headers(dd_trace_context)
@@ -155,6 +174,9 @@ def set_correlation_ids():
     if not is_lambda_context():
         logger.debug("set_correlation_ids is only supported in LambdaContext")
         return
+    if dd_native_tracing_enabled:
+        logger.debug("using ddtrace implementation for spans")
+        return
 
     context = get_dd_trace_context()
 

datadog_lambda/tracing.py

@@ -16,12 +16,14 @@
 from datadog_lambda.patch import patch_all
 from datadog_lambda.tracing import (
     extract_dd_trace_context,
-    set_correlation_ids,
     inject_correlation_ids,
     get_dd_trace_context,
+    dd_native_tracing_enabled,
+    set_correlation_ids,
 )
-from datadog_lambda.trace_wrapper import trace_wrapper
 from datadog_lambda.constants import Source
+from ddtrace import tracer, Span
+from ddtrace.propagation.http import HTTPPropagator
 
 
 logger = logging.getLogger(__name__)
@@ -84,8 +86,12 @@ def __init__(self, func):
             self.logs_injection = (
                 os.environ.get("DD_LOGS_INJECTION", "true").lower() == "true"
             )
+            self.merge_xray_traces = (
+                os.environ.get("DD_MERGE_XRAY_TRACES", "false").lower() == "true"
+            )
             self.handler_name = os.environ.get("_HANDLER", "handler")
             self.function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", "function")
+            self.propagator = HTTPPropagator()
 
             # Inject trace correlation ids to logs
             if self.logs_injection:
@@ -116,8 +122,9 @@ def _before(self, event, context):
             # Extract Datadog trace context from incoming requests
             dd_context = extract_dd_trace_context(event)
             span_context = None
-            if dd_context["source"] == Source.EVENT:
-                span_context = trace_wrapper.extract(dd_context)
+            if dd_context["source"] == Source.EVENT or self.merge_xray_traces:
+                headers = get_dd_trace_context()
+                span_context = self.propagator.extract(headers)
 
             tags = {}
             if context:
@@ -134,12 +141,14 @@ def _before(self, event, context):
                 "child_of": span_context,
             }
 
-            self.span = trace_wrapper.start_span("aws.lambda", **args)
-            if self.span:
-                self.span.set_tags(tags)
+            self.span = None
+            if dd_native_tracing_enabled:
+                self.span = tracer.start_span("aws.lambda", **args)
+                if self.span:
+                    self.span.set_tags(tags)
+            else:
+                set_correlation_ids()
 
-            # Set log correlation ids using extracted trace context
-            set_correlation_ids()
             logger.debug("datadog_lambda_wrapper _before() done")
         except Exception:
             traceback.print_exc()