[Tracing] Adds SQS and LambdaContext Trace extraction

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"]}