Darcy.rayner/dd trace support

README.md

@@ -10,7 +10,7 @@ Datadog Lambda Layer for Python (2.7, 3.6, 3.7 and 3.8) enables custom metric su
 
 ## IMPORTANT NOTE
 
-AWS Lambda is expected to recieve a [breaking change](https://aws.amazon.com/blogs/compute/upcoming-changes-to-the-python-sdk-in-aws-lambda/) on **January 30, 2021**. If you are using Datadog Python Lambda layer version 7 or below, please upgrade to version 11. 
+AWS Lambda is expected to recieve a [breaking change](https://aws.amazon.com/blogs/compute/upcoming-changes-to-the-python-sdk-in-aws-lambda/) on **January 30, 2021**. If you are using Datadog Python Lambda layer version 7 or below, please upgrade to version 11.
 
 ## Installation
 
@@ -21,6 +21,7 @@ arn:aws:lambda:<AWS_REGION>:464622532012:layer:Datadog-<PYTHON_RUNTIME>:<VERSION
 ```
 
 Replace `<AWS_REGION>` with the AWS region where your Lambda function is published to. Replace `<PYTHON_RUNTIME>` with one of the following that matches your Lambda's Python runtime:
+
 - `Datadog-Python27`
 - `Datadog-Python36`
 - `Datadog-Python37`
@@ -81,7 +82,7 @@ If `DD_FLUSH_TO_LOG` is set to false (not recommended), the Datadog API Key must
 - DD_KMS_API_KEY - the KMS-encrypted API Key, requires the `kms:Decrypt` permission
 - DD_API_KEY_SECRET_ARN - the Secret ARN to fetch API Key from the Secrets Manager, requires the `secretsmanager:GetSecretValue` permission (and `kms:Decrypt` if using a customer managed CMK)
 
-You can also supply or override the API key at runtime (not recommended): 
+You can also supply or override the API key at runtime (not recommended):
 
 ```python
 # Override DD API Key after importing datadog_lambda packages
@@ -243,6 +244,7 @@ If your Lambda function is triggered by API Gateway via [the non-proxy integrati
 If your Lambda function is deployed by the Serverless Framework, such a mapping template gets created by default.
 
 ## Log and Trace Correlations
+
 By default, the Datadog trace id gets automatically injected into the logs for correlation, if using the standard python `logging` library.
 
 If you use a custom logger handler to log in json, you can inject the ids using the helper function `get_correlation_ids` [manually](https://docs.datadoghq.com/tracing/connect_logs_and_traces/?tab=python#manual-trace-id-injection).
@@ -265,6 +267,35 @@ def lambda_handler(event, context):
   })
 ```
 
+## Datadog Tracer (**Experimental**)
+
+You can now trace Lambda functions using Datadog APM's tracing libraries ([dd-trace-py](https://github.com/DataDog/dd-trace-py)).
+
+1. If you are using the Lambda layer, upgrade it to at least version 15.
+1. If you are using the pip package `datadog-lambda-python`, upgrade it to at least version `v0.X.0`.
+1. Install (or update to) the latest version of [Datadog forwarder Lambda function](https://docs.datadoghq.com/integrations/amazon_web_services/?tab=allpermissions#set-up-the-datadog-lambda-function). Ensure the trace forwarding layer is attached to the forwarder, e.g., ARN for Python 2.7 `arn:aws:lambda:<AWS_REGION>:464622532012:layer:Datadog-Trace-Forwarder-Python27:4`.
+1. Set the environment variable `DD_TRACE_ENABLED` to true on your function.
+1. Instrument your function using `dd-trace`.
+
+```py
+from datadog_lambda.metric import lambda_metric
+from datadog_lambda.wrapper import datadog_lambda_wrapper
+
+from ddtrace import tracer
+@datadog_lambda_wrapper
+def hello(event, context):
+  return {
+    "statusCode": 200,
+    "body": get_message()
+  }
+
+@tracer.wrap()
+def get_message():
+  return "hello world"
+```
+
+You can also use `dd-trace` and the X-Ray tracer together and merge the traces into one, using the environment variable `DD_MERGE_XRAY_TRACES` to true on your function.
+
 ## Opening Issues
 
 If you encounter a bug with this package, we want to hear about it. Before opening a new issue, search the existing issues to avoid duplicates.

datadog_lambda/__init__.py

@@ -1,6 +1,6 @@
 # The minor version corresponds to the Lambda layer version.
 # E.g.,, version 0.5.0 gets packaged into layer version 5.
-__version__ = "2.14.0"
+__version__ = "2.15.0"
 
 
 import os

datadog_lambda/constants.py

@@ -24,3 +24,12 @@ class XraySubsegment(object):
     NAME = "datadog-metadata"
     KEY = "trace"
     NAMESPACE = "datadog"
+
+
+# TraceContextSource of datadog context. The DD_MERGE_XRAY_TRACES
+# feature uses this to determine when to use X-Ray as the parent
+# trace.
+class TraceContextSource(object):
+    XRAY = "xray"
+    EVENT = "event"
+    DDTRACE = "ddtrace"

datadog_lambda/tracing.py

@@ -4,16 +4,23 @@
 # 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
+from datadog_lambda.constants import (
+    SamplingPriority,
+    TraceHeader,
+    XraySubsegment,
+    TraceContextSource,
+)
+from ddtrace import tracer, patch
 
 logger = logging.getLogger(__name__)
 
 dd_trace_context = {}
+dd_tracing_enabled = os.environ.get("DD_TRACE_ENABLED", "false").lower() == "true"
 
 
 def _convert_xray_trace_id(xray_trace_id):
@@ -41,6 +48,43 @@ def _convert_xray_sampling(xray_sampled):
     )
 
 
+def _get_xray_trace_context():
+    if not is_lambda_context():
+        return None
+
+    xray_trace_entity = xray_recorder.get_trace_entity()  # xray (sub)segment
+    return {
+        "trace-id": _convert_xray_trace_id(xray_trace_entity.trace_id),
+        "parent-id": _convert_xray_entity_id(xray_trace_entity.id),
+        "sampling-priority": _convert_xray_sampling(xray_trace_entity.sampled),
+        "source": TraceContextSource.XRAY,
+    }
+
+
+def _get_dd_trace_py_context():
+    span = tracer.current_span()
+    if not span:
+        return None
+
+    parent_id = span.context.span_id
+    trace_id = span.context.trace_id
+    sampling_priority = span.context.sampling_priority
+    return {
+        "parent-id": str(parent_id),
+        "trace-id": str(trace_id),
+        "sampling-priority": str(sampling_priority),
+        "source": TraceContextSource.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")),
+    }
+
+
 def extract_dd_trace_context(event):
     """
     Extract Datadog trace context from the Lambda `event` object.
@@ -61,23 +105,24 @@ def extract_dd_trace_context(event):
     sampling_priority = lowercase_headers.get(TraceHeader.SAMPLING_PRIORITY)
     if trace_id and parent_id and sampling_priority:
         logger.debug("Extracted Datadog trace context from headers")
-        dd_trace_context = {
+        metadata = {
             "trace-id": trace_id,
             "parent-id": parent_id,
             "sampling-priority": sampling_priority,
         }
         xray_recorder.begin_subsegment(XraySubsegment.NAME)
         subsegment = xray_recorder.current_subsegment()
-        subsegment.put_metadata(
-            XraySubsegment.KEY, dd_trace_context, XraySubsegment.NAMESPACE
-        )
+
+        subsegment.put_metadata(XraySubsegment.KEY, metadata, XraySubsegment.NAMESPACE)
+        dd_trace_context = metadata.copy()
+        dd_trace_context["source"] = TraceContextSource.EVENT
         xray_recorder.end_subsegment()
     else:
         # AWS Lambda runtime caches global variables between invocations,
         # reset to avoid using the context from the last invocation.
-        dd_trace_context = {}
-
+        dd_trace_context = _get_xray_trace_context()
     logger.debug("extracted dd trace context %s", dd_trace_context)
+    return dd_trace_context
 
 
 def get_dd_trace_context():
@@ -86,32 +131,38 @@ def get_dd_trace_context():
 
     If the Lambda function is invoked by a Datadog-traced service, a Datadog
     trace context may already exist, and it should be used. Otherwise, use the
-    current X-Ray trace entity.
+    current X-Ray trace entity, or the dd-trace-py context if DD_TRACE_ENABLED is true.
 
     Most of widely-used HTTP clients are patched to inject the context
     automatically, but this function can be used to manually inject the trace
     context to an outgoing request.
     """
-    if not is_lambda_context():
-        logger.debug("get_dd_trace_context is only supported in LambdaContext")
-        return {}
-
     global dd_trace_context
-    xray_trace_entity = xray_recorder.get_trace_entity()  # xray (sub)segment
-    if dd_trace_context:
-        return {
-            TraceHeader.TRACE_ID: dd_trace_context["trace-id"],
-            TraceHeader.PARENT_ID: _convert_xray_entity_id(xray_trace_entity.id),
-            TraceHeader.SAMPLING_PRIORITY: dd_trace_context["sampling-priority"],
-        }
-    else:
-        return {
-            TraceHeader.TRACE_ID: _convert_xray_trace_id(xray_trace_entity.trace_id),
-            TraceHeader.PARENT_ID: _convert_xray_entity_id(xray_trace_entity.id),
-            TraceHeader.SAMPLING_PRIORITY: _convert_xray_sampling(
-                xray_trace_entity.sampled
-            ),
-        }
+
+    context = None
+    xray_context = None
+
+    try:
+        xray_context = _get_xray_trace_context()  # xray (sub)segment
+    except Exception as e:
+        logger.debug(
+            "get_dd_trace_context couldn't read from segment from x-ray, with error %s"
+            % e
+        )
+
+    if xray_context and not dd_trace_context:
+        context = xray_context
+    elif xray_context and dd_trace_context:
+        context = dd_trace_context.copy()
+        context["parent-id"] = xray_context["parent-id"]
+
+    if dd_tracing_enabled:
+        dd_trace_py_context = _get_dd_trace_py_context()
+        if dd_trace_py_context is not None:
+            logger.debug("get_dd_trace_context using dd-trace context")
+            context = dd_trace_py_context
+
+    return _context_obj_to_headers(context) if context is not None else {}
 
 
 def set_correlation_ids():
@@ -125,6 +176,9 @@ def set_correlation_ids():
     if not is_lambda_context():
         logger.debug("set_correlation_ids is only supported in LambdaContext")
         return
+    if dd_tracing_enabled:
+        logger.debug("using ddtrace implementation for spans")
+        return
 
     context = get_dd_trace_context()
 

datadog_lambda/wrapper.py

@@ -7,7 +7,7 @@
 import logging
 import traceback
 
-from datadog_lambda.cold_start import set_cold_start
+from datadog_lambda.cold_start import set_cold_start, is_cold_start
 from datadog_lambda.metric import (
     lambda_stats,
     submit_invocations_metric,
@@ -16,9 +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_tracing_enabled,
+    set_correlation_ids,
 )
+from datadog_lambda.constants import TraceContextSource
+from ddtrace import tracer
+from ddtrace.propagation.http import HTTPPropagator
 
 
 logger = logging.getLogger(__name__)
@@ -81,13 +86,21 @@ 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:
                 inject_correlation_ids()
 
-            # Patch HTTP clients to propagate Datadog trace context
-            patch_all()
+            if not dd_tracing_enabled:
+                # When using dd_trace_py it will patch all the http clients for us,
+                # Patch HTTP clients to propagate Datadog trace context
+                patch_all()
             logger.debug("datadog_lambda_wrapper initialized")
         except Exception:
             traceback.print_exc()
@@ -105,24 +118,59 @@ def __call__(self, event, context, **kwargs):
 
     def _before(self, event, context):
         try:
+
             set_cold_start()
             submit_invocations_metric(context)
             # Extract Datadog trace context from incoming requests
-            extract_dd_trace_context(event)
+            dd_context = extract_dd_trace_context(event)
+
+            self.span = None
+            if dd_tracing_enabled:
+                self.span = self._create_dd_trace_py_span(context, dd_context)
+            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()
 
     def _after(self, event, context):
         try:
+            if self.span:
+                self.span.finish()
             if not self.flush_to_log:
                 lambda_stats.flush(float("inf"))
             logger.debug("datadog_lambda_wrapper _after() done")
         except Exception:
             traceback.print_exc()
 
+    def _create_dd_trace_py_span(self, context, trace_context):
+        span_context = None
+        if (
+            trace_context["source"] == TraceContextSource.EVENT
+            or self.merge_xray_traces
+        ):
+            headers = get_dd_trace_context()
+            span_context = self.propagator.extract(headers)
+
+        tags = {}
+        if context:
+            tags = {
+                "cold_start": is_cold_start(),
+                "function_arn": context.invoked_function_arn,
+                "request_id": context.aws_request_id,
+                "resource_names": context.function_name,
+            }
+        args = {
+            "service": self.function_name,
+            "resource": self.handler_name,
+            "span_type": "serverless",
+            "child_of": span_context,
+        }
+        span = tracer.start_span("aws.lambda", **args)
+        if span:
+            span.set_tags(tags)
+        return span
+
 
 datadog_lambda_wrapper = _LambdaDecorator

scripts/publish_staging.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+set -e
+
+./scripts/build_layers.sh
+./scripts/publish_layers.sh us-east-1

scripts/run_integration_tests.sh

@@ -116,7 +116,15 @@ for handler_name in "${LAMBDA_HANDLERS[@]}"; do
                 # Strip API key from logged requests
                 sed -E "s/(api_key=|'api_key': ')[a-z0-9\.\-]+/\1XXXX/g" |
                 # Normalize minor package version so that these snapshots aren't broken on version bumps
-                sed -E "s/(dd_lambda_layer:datadog-python[0-9]+_2\.)[0-9]+\.0/\1XX\.0/g"
+                sed -E "s/(dd_lambda_layer:datadog-python[0-9]+_2\.)[0-9]+\.0/\1XX\.0/g" |
+                # Strip out trace/span/parent/timestamps
+                sed -E "s/(\"trace_id\"\: \")[A-Z0-9\.\-]+/\1XXXX/g" |
+                sed -E "s/(\"span_id\"\: \")[A-Z0-9\.\-]+/\1XXXX/g" |
+                sed -E "s/(\"parent_id\"\: \")[A-Z0-9\.\-]+/\1XXXX/g" |
+                sed -E "s/(\"request_id\"\: \")[a-z0-9\.\-]+/\1XXXX/g" |
+                sed -E "s/(\"duration\"\: )[0-9\.\-]+/\1XXXX/g" |
+                sed -E "s/(\"start\"\: )[0-9\.\-]+/\1XXXX/g" |
+                sed -E "s/(\"system\.pid\"\: )[0-9\.\-]+/\1XXXX/g"
         )
 
         if [ ! -f $function_snapshot_path ]; then