[Tracing] Adds SQS and LambdaContext Trace extraction (#105)

* adds sqs and lambda trace extraction

* format

* fix parsing error

* remove extra whitespace

* format again

* add integration tests update

* add lambda context fallback + int tests

* add fallbacks

* make client context extraction more resilient

* fix lint

* nit comment

Author: Czechh

datadog_lambda/tracing.py

@@ -5,6 +5,7 @@
 
 import logging
 import os
+import json
 
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.core.lambda_launcher import LambdaContext
@@ -89,7 +90,67 @@ def _context_obj_to_headers(obj):
     }
 
 
-def extract_dd_trace_context(event):
+def extract_context_from_lambda_context(lambda_context):
+    """
+    Extract Datadog trace context from the `client_context` attr
+    from the Lambda `context` object.
+
+    dd_trace libraries inject this trace context on synchronous invocations
+    """
+    client_context = lambda_context.client_context
+
+    if client_context and "custom" in client_context:
+        dd_data = client_context.get("custom", {}).get("_datadog", {})
+        trace_id = dd_data.get(TraceHeader.TRACE_ID)
+        parent_id = dd_data.get(TraceHeader.PARENT_ID)
+        sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
+
+        return trace_id, parent_id, sampling_priority
+
+    return None, None, None
+
+
+def extract_context_from_http_event_or_context(event, lambda_context):
+    """
+    Extract Datadog trace context from the `headers` key in from the Lambda
+    `event` object.
+
+    Falls back to lambda context if no trace data is found in the `headers`
+    """
+    headers = event.get("headers", {})
+    lowercase_headers = {k.lower(): v for k, v in headers.items()}
+
+    trace_id = lowercase_headers.get(TraceHeader.TRACE_ID)
+    parent_id = lowercase_headers.get(TraceHeader.PARENT_ID)
+    sampling_priority = lowercase_headers.get(TraceHeader.SAMPLING_PRIORITY)
+
+    if not trace_id or not parent_id or not sampling_priority:
+        return extract_context_from_lambda_context(lambda_context)
+
+    return trace_id, parent_id, sampling_priority
+
+
+def extract_context_from_sqs_event_or_context(event, lambda_context):
+    """
+    Extract Datadog trace context from the first SQS message attributes.
+
+    Falls back to lambda context if no trace data is found in the SQS message attributes.
+    """
+    try:
+        first_record = event["Records"][0]
+        msg_attributes = first_record.get("messageAttributes", {})
+        dd_json_data = msg_attributes.get("_datadog", {}).get("StringValue", r"{}")
+        dd_data = json.loads(dd_json_data)
+        trace_id = dd_data.get(TraceHeader.TRACE_ID)
+        parent_id = dd_data.get(TraceHeader.PARENT_ID)
+        sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY)
+
+        return trace_id, parent_id, sampling_priority
+    except Exception:
+        return extract_context_from_lambda_context(lambda_context)
+
+
+def extract_dd_trace_context(event, lambda_context):
     """
     Extract Datadog trace context from the Lambda `event` object.
 
@@ -101,14 +162,26 @@ def extract_dd_trace_context(event):
     the correct context.
     """
     global dd_trace_context
-    headers = event.get("headers", {})
-    lowercase_headers = {k.lower(): v for k, v in headers.items()}
 
-    trace_id = lowercase_headers.get(TraceHeader.TRACE_ID)
-    parent_id = lowercase_headers.get(TraceHeader.PARENT_ID)
-    sampling_priority = lowercase_headers.get(TraceHeader.SAMPLING_PRIORITY)
+    if "headers" in event:
+        (
+            trace_id,
+            parent_id,
+            sampling_priority,
+        ) = extract_context_from_http_event_or_context(event, lambda_context)
+    elif "Records" in event:
+        (
+            trace_id,
+            parent_id,
+            sampling_priority,
+        ) = extract_context_from_sqs_event_or_context(event, lambda_context)
+    else:
+        trace_id, parent_id, sampling_priority = extract_context_from_lambda_context(
+            lambda_context
+        )
+
     if trace_id and parent_id and sampling_priority:
-        logger.debug("Extracted Datadog trace context from headers")
+        logger.debug("Extracted Datadog trace context from event or context")
         metadata = {
             "trace-id": trace_id,
             "parent-id": parent_id,

datadog_lambda/wrapper.py

@@ -119,7 +119,7 @@ def _before(self, event, context):
             set_cold_start()
             submit_invocations_metric(context)
             # Extract Datadog trace context from incoming requests
-            dd_context = extract_dd_trace_context(event)
+            dd_context = extract_dd_trace_context(event, context)
 
             self.span = None
             if dd_tracing_enabled:

scripts/run_integration_tests.sh

@@ -148,7 +148,9 @@ for _sls_type in "${CONFIGS[@]}"; do
                     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" |
-                    sed -E "s/(\"runtime-id\"\: \")[a-z0-9\.\-]+/\1XXXX/g"
+                    sed -E "s/(\"runtime-id\"\: \")[a-z0-9\.\-]+/\1XXXX/g" |
+                    sed -E "s/(\"datadog_lambda\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g" |
+                    sed -E "s/(\"dd_trace\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g"
             )
 
             if [ ! -f $function_snapshot_path ]; then

scripts/run_tests.sh

@@ -23,5 +23,5 @@ do
     docker run -v `pwd`:/datadog-lambda-python \
         -w /datadog-lambda-python \
         datadog-lambda-python-test:$python_version \
-        flake8
+        flake8 datadog_lambda/
 done

setup.py

@@ -30,7 +30,7 @@
     install_requires=[
         "aws-xray-sdk==2.6.0",
         "datadog==0.39.0",
-        "ddtrace==0.44.0",
+        "ddtrace==0.45.0",
         "wrapt==1.11.2",
         "setuptools==42.0.2",
     ],

tests/integration/input_events/sqs.json

@@ -10,7 +10,11 @@
         "SenderId": "AIDAIENQZJOLO23YVJ4VO",
         "ApproximateFirstReceiveTimestamp": "1545082649185"
       },
-      "messageAttributes": {},
+      "messageAttributes": {
+        "_datadog": {
+          "StringValue": "{\"x-datadog-trace-id\":\"666\",\"x-datadog-parent-id\":\"777\",\"x-datadog-sampling-priority\":\"1\"}"
+        }
+      },
       "md5OfBody": "e4e68fb7bd0e697a0ae8f1bb342846b3",
       "eventSource": "aws:sqs",
       "eventSourceARN": "arn:aws:sqs:us-east-2:123456789012:my-queue",

tests/integration/snapshots/logs/async-metrics_python27.log

@@ -1,4 +1,4 @@
-/var/runtime/botocore/vendored/requests/__init__.py:91: RequestsDependencyWarning: urllib3 (1.26.2) or chardet (3.0.4) doesn't match a supported version!
+/var/runtime/botocore/vendored/requests/__init__.py:91: RequestsDependencyWarning: urllib3 (1.26.2) or chardet (4.0.0) doesn't match a supported version!
 RequestsDependencyWarning)
 START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-dev-async-metrics_python27", "resource:integration-dev-async-metrics_python27", "cold_start:true", "memorysize:1024", "runtime:python2.7", "datadog_lambda:vXX", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}

tests/integration/snapshots/logs/async-metrics_python27_with_plugin.log

@@ -1,26 +1,26 @@
-/var/runtime/botocore/vendored/requests/__init__.py:91: RequestsDependencyWarning: urllib3 (1.26.2) or chardet (3.0.4) doesn't match a supported version!
+/var/runtime/botocore/vendored/requests/__init__.py:91: RequestsDependencyWarning: urllib3 (1.26.2) or chardet (4.0.0) doesn't match a supported version!
 RequestsDependencyWarning)
 START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-plugin-dev-async-metrics_python27_with_plugin", "resource:integration-plugin-dev-async-metrics_python27_with_plugin", "cold_start:true", "memorysize:1024", "runtime:python2.7", "datadog_lambda:vXX", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 {"e": XXXX, "m": "hello.dog", "t": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 {"e": XXXX, "m": "tests.integration.count", "t": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 21}
-{"traces": [[{"resource": "integration-plugin-dev-async-metrics_python27_with_plugin", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 2, "system.pid": XXXX}, "parent_id": "XXXX", "meta": {"runtime-id": "XXXX", "request_id": "XXXX", "cold_start": "true", "datadog_lambda": "2.24.0", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-plugin-dev-async-metrics_python27_with_plugin", "dd_trace": "0.44.0", "_dd.origin": "lambda", "_dd.parent_source": "xray", "resource_names": "integration-plugin-dev-async-metrics_python27_with_plugin", "function_version": "$LATEST"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
+{"traces": [[{"resource": "integration-plugin-dev-async-metrics_python27_with_plugin", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 2, "system.pid": XXXX}, "parent_id": "XXXX", "meta": {"runtime-id": "XXXX", "request_id": "XXXX", "cold_start": "true", "datadog_lambda": "X.X.X", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-plugin-dev-async-metrics_python27_with_plugin", "dd_trace": "X.X.X", "_dd.origin": "lambda", "_dd.parent_source": "xray", "resource_names": "integration-plugin-dev-async-metrics_python27_with_plugin", "function_version": "$LATEST"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB	Init Duration: XXXX ms	
 XRAY TraceId: XXXX	SegmentId: XXXX	Sampled: true	
 START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-plugin-dev-async-metrics_python27_with_plugin", "resource:integration-plugin-dev-async-metrics_python27_with_plugin", "cold_start:false", "memorysize:1024", "runtime:python2.7", "datadog_lambda:vXX", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 {"e": XXXX, "m": "hello.dog", "t": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 {"e": XXXX, "m": "tests.integration.count", "t": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 21}
-{"traces": [[{"resource": "integration-plugin-dev-async-metrics_python27_with_plugin", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 2, "system.pid": XXXX}, "parent_id": "XXXX", "meta": {"runtime-id": "XXXX", "request_id": "XXXX", "cold_start": "false", "datadog_lambda": "2.24.0", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-plugin-dev-async-metrics_python27_with_plugin", "dd_trace": "0.44.0", "_dd.origin": "lambda", "_dd.parent_source": "xray", "resource_names": "integration-plugin-dev-async-metrics_python27_with_plugin", "function_version": "$LATEST"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
+{"traces": [[{"resource": "integration-plugin-dev-async-metrics_python27_with_plugin", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 2, "system.pid": XXXX}, "parent_id": "XXXX", "meta": {"runtime-id": "XXXX", "request_id": "XXXX", "cold_start": "false", "datadog_lambda": "X.X.X", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-plugin-dev-async-metrics_python27_with_plugin", "dd_trace": "X.X.X", "_dd.origin": "lambda", "_dd.parent_source": "xray", "resource_names": "integration-plugin-dev-async-metrics_python27_with_plugin", "function_version": "$LATEST"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB	
 XRAY TraceId: XXXX	SegmentId: XXXX	Sampled: true	
 START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-plugin-dev-async-metrics_python27_with_plugin", "resource:integration-plugin-dev-async-metrics_python27_with_plugin", "cold_start:false", "memorysize:1024", "runtime:python2.7", "datadog_lambda:vXX", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 {"e": XXXX, "m": "hello.dog", "t": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 {"e": XXXX, "m": "tests.integration.count", "t": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 21}
-{"traces": [[{"resource": "integration-plugin-dev-async-metrics_python27_with_plugin", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 2, "system.pid": XXXX}, "parent_id": "XXXX", "meta": {"runtime-id": "XXXX", "request_id": "XXXX", "cold_start": "false", "datadog_lambda": "2.24.0", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-plugin-dev-async-metrics_python27_with_plugin", "dd_trace": "0.44.0", "_dd.origin": "lambda", "_dd.parent_source": "xray", "resource_names": "integration-plugin-dev-async-metrics_python27_with_plugin", "function_version": "$LATEST"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
+{"traces": [[{"resource": "integration-plugin-dev-async-metrics_python27_with_plugin", "name": "aws.lambda", "service": "aws.lambda", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 1, "system.pid": XXXX}, "parent_id": "XXXX", "meta": {"runtime-id": "XXXX", "request_id": "XXXX", "cold_start": "false", "datadog_lambda": "X.X.X", "function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-plugin-dev-async-metrics_python27_with_plugin", "dd_trace": "X.X.X", "_dd.origin": "lambda", "resource_names": "integration-plugin-dev-async-metrics_python27_with_plugin", "function_version": "$LATEST"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB	
 XRAY TraceId: XXXX	SegmentId: XXXX	Sampled: true	

tests/integration/snapshots/logs/async-metrics_python36.log

@@ -1,4 +1,4 @@
-/var/runtime/botocore/vendored/requests/__init__.py:91: RequestsDependencyWarning: urllib3 (1.26.2) or chardet (3.0.4) doesn't match a supported version!
+/var/runtime/botocore/vendored/requests/__init__.py:91: RequestsDependencyWarning: urllib3 (1.26.2) or chardet (4.0.0) doesn't match a supported version!
 RequestsDependencyWarning)
 START RequestId: XXXX Version: $LATEST
 {"m": "aws.lambda.enhanced.invocations", "v": 1, "e": XXXX, "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-dev-async-metrics_python36", "resource:integration-dev-async-metrics_python36", "cold_start:true", "memorysize:1024", "runtime:python3.6", "datadog_lambda:vXX", "dd_lambda_layer:datadog-python36_2.XX.0"]}