Implement PR feedback

Author: DarcyRaynerDD

datadog_lambda/constants.py

@@ -26,8 +26,10 @@ class XraySubsegment(object):
     NAMESPACE = "datadog"
 
 
-# Source of datadog context
-class Source(object):
+# 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

@@ -13,16 +13,14 @@
     SamplingPriority,
     TraceHeader,
     XraySubsegment,
-    Source,
+    TraceContextSource,
 )
 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"
-)
+dd_tracing_enabled = os.environ.get("DD_TRACE_ENABLED", "false").lower() == "true"
 
 
 def _convert_xray_trace_id(xray_trace_id):
@@ -59,22 +57,23 @@ def _get_xray_trace_context():
         "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": Source.XRAY,
+        "source": TraceContextSource.XRAY,
     }
 
 
-def _get_dd_trace_native_context():
+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": SamplingPriority.AUTO_KEEP,
-        "source": Source.DDTRACE,
+        "sampling-priority": str(sampling_priority),
+        "source": TraceContextSource.DDTRACE,
     }
 
 
@@ -116,7 +115,7 @@ def extract_dd_trace_context(event):
 
         subsegment.put_metadata(XraySubsegment.KEY, metadata, XraySubsegment.NAMESPACE)
         dd_trace_context = metadata.copy()
-        dd_trace_context["source"] = Source.EVENT
+        dd_trace_context["source"] = TraceContextSource.EVENT
         xray_recorder.end_subsegment()
     else:
         # AWS Lambda runtime caches global variables between invocations,
@@ -132,35 +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.
     """
     global dd_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)
+    context = None
+    xray_context = None
 
     try:
-        trace_headers = _context_obj_to_headers(dd_trace_context)
         xray_context = _get_xray_trace_context()  # xray (sub)segment
-        if xray_context and not trace_headers:
-            return _context_obj_to_headers(xray_context)
-        if xray_context and trace_headers:
-            trace_headers[TraceHeader.PARENT_ID] = xray_context["parent-id"]
-            return trace_headers
     except Exception as e:
         logger.debug(
             "get_dd_trace_context couldn't read from segment from x-ray, with error %s"
             % e
         )
 
-    return {}
+    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():
@@ -174,7 +176,7 @@ 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:
+    if dd_tracing_enabled:
         logger.debug("using ddtrace implementation for spans")
         return
 

datadog_lambda/wrapper.py

@@ -18,10 +18,10 @@
     extract_dd_trace_context,
     inject_correlation_ids,
     get_dd_trace_context,
-    dd_native_tracing_enabled,
+    dd_tracing_enabled,
     set_correlation_ids,
 )
-from datadog_lambda.constants import Source
+from datadog_lambda.constants import TraceContextSource
 from ddtrace import tracer
 from ddtrace.propagation.http import HTTPPropagator
 
@@ -97,8 +97,10 @@ def __init__(self, func):
             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()
@@ -121,31 +123,10 @@ def _before(self, event, context):
             submit_invocations_metric(context)
             # Extract Datadog trace context from incoming requests
             dd_context = extract_dd_trace_context(event)
-            span_context = None
-            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:
-                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,
-            }
 
             self.span = None
-            if dd_native_tracing_enabled:
-                self.span = tracer.start_span("aws.lambda", **args)
-                if self.span:
-                    self.span.set_tags(tags)
+            if dd_tracing_enabled:
+                self.span = self._create_dd_trace_py_span(context, dd_context)
             else:
                 set_correlation_ids()
 
@@ -163,5 +144,33 @@ def _after(self, event, context):
         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/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

tests/integration/snapshots/logs/http-requests_python27.log

@@ -2,7 +2,7 @@ START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python27", "cold_start:true", "memorysize:1024", "runtime:python2.7", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 HTTP GET https://ip-ranges.datadoghq.com/ Headers: ["x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {}
 HTTP GET https://ip-ranges.datadoghq.eu/ Headers: ["x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {}
-{"traces": [[{"resource": "http_requests.handle", "name": "aws.lambda", "service": "integration-tester-dev-http-requests_python27", "start": 1585077483318449000, "trace_id": "06C0653FC881D935", "metrics": {"_sampling_priority_v1": 1, "system.pid": 1, "_dd.agent_psr": 1.0}, "parent_id": "0000000000000000", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "request_id": "da026ab0-1d6c-4980-9720-1b4bdd843314", "cold_start": "True", "resource_names": "integration-tester-dev-http-requests_python27"}, "error": 0, "duration": 389814000, "type": "serverless", "span_id": "3DA8AA128268BFA1"}]]}
+{"traces": [[{"resource": "http_requests.handle", "name": "aws.lambda", "service": "integration-tester-dev-http-requests_python27", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 1, "system.pid": XXXX, "_dd.agent_psr": 1.0}, "parent_id": "XXXX", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "request_id": "XXXX", "cold_start": "True", "resource_names": "integration-tester-dev-http-requests_python27"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
 HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {"series": [{"tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "hello.dog", "interval": 10, "host": null, "points": [[XXXX, [1.0]]], "device": null, "type": "distribution"}, {"tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "tests.integration.count", "interval": 10, "host": null, "points": [[XXXX, [21.0]]], "device": null, "type": "distribution"}]}
 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	
@@ -11,7 +11,7 @@ START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python27", "cold_start:false", "memorysize:1024", "runtime:python2.7", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 HTTP GET https://ip-ranges.datadoghq.com/ Headers: ["x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {}
 HTTP GET https://ip-ranges.datadoghq.eu/ Headers: ["x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {}
-{"traces": [[{"resource": "http_requests.handle", "name": "aws.lambda", "service": "integration-tester-dev-http-requests_python27", "start": 1585077485080991000, "trace_id": "6BDE1FAE35B84C6C", "metrics": {"_sampling_priority_v1": 1, "system.pid": 1, "_dd.agent_psr": 1.0}, "parent_id": "0000000000000000", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "request_id": "f1b1655b-ce14-4540-bf54-8e128af7b8ab", "cold_start": "False", "resource_names": "integration-tester-dev-http-requests_python27"}, "error": 0, "duration": 80971000, "type": "serverless", "span_id": "393ECB596ED13F61"}]]}
+{"traces": [[{"resource": "http_requests.handle", "name": "aws.lambda", "service": "integration-tester-dev-http-requests_python27", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 1, "system.pid": XXXX, "_dd.agent_psr": 1.0}, "parent_id": "XXXX", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "request_id": "XXXX", "cold_start": "False", "resource_names": "integration-tester-dev-http-requests_python27"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
 HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {"series": [{"tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "hello.dog", "interval": 10, "host": null, "points": [[XXXX, [1.0]]], "device": null, "type": "distribution"}, {"tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "tests.integration.count", "interval": 10, "host": null, "points": [[XXXX, [21.0]]], "device": null, "type": "distribution"}]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB	
@@ -20,7 +20,7 @@ START RequestId: XXXX Version: $LATEST
 {"e": XXXX, "m": "aws.lambda.enhanced.invocations", "t": ["region:us-east-1", "account_id:XXXX", "functionname:integration-tester-dev-http-requests_python27", "cold_start:false", "memorysize:1024", "runtime:python2.7", "dd_lambda_layer:datadog-python27_2.XX.0"], "v": 1}
 HTTP GET https://ip-ranges.datadoghq.com/ Headers: ["x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {}
 HTTP GET https://ip-ranges.datadoghq.eu/ Headers: ["x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {}
-{"traces": [[{"resource": "http_requests.handle", "name": "aws.lambda", "service": "integration-tester-dev-http-requests_python27", "start": 1585077486596975000, "trace_id": "0A7F7B686FDCC310", "metrics": {"_sampling_priority_v1": 1, "system.pid": 1, "_dd.agent_psr": 1.0}, "parent_id": "0000000000000000", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "request_id": "6f90344e-1fa5-48cb-8933-d7f523868746", "cold_start": "False", "resource_names": "integration-tester-dev-http-requests_python27"}, "error": 0, "duration": 100314000, "type": "serverless", "span_id": "58E9AB1D40B41342"}]]}
+{"traces": [[{"resource": "http_requests.handle", "name": "aws.lambda", "service": "integration-tester-dev-http-requests_python27", "start": XXXX, "trace_id": "XXXX", "metrics": {"_sampling_priority_v1": 1, "system.pid": XXXX, "_dd.agent_psr": 1.0}, "parent_id": "XXXX", "meta": {"function_arn": "arn:aws:lambda:us-east-1:601427279990:function:integration-tester-dev-http-requests_python27", "request_id": "XXXX", "cold_start": "False", "resource_names": "integration-tester-dev-http-requests_python27"}, "error": 0, "duration": XXXX, "type": "serverless", "span_id": "XXXX"}]]}
 HTTP POST https://api.datadoghq.com/api/v1/distribution_points Headers: ["Content-Type:application/json", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:2", "x-datadog-trace-id:XXXX"] Data: {"series": [{"tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "hello.dog", "interval": 10, "host": null, "points": [[XXXX, [1.0]]], "device": null, "type": "distribution"}, {"tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python27_2.XX.0"], "metric": "tests.integration.count", "interval": 10, "host": null, "points": [[XXXX, [21.0]]], "device": null, "type": "distribution"}]}
 END RequestId: XXXX
 REPORT RequestId: XXXX	Duration: XXXX ms	Billed Duration: XXXX ms	Memory Size: 1024 MB	Max Memory Used: XXXX MB