Implement dd-trace logic

Author: DarcyRaynerDD

datadog_lambda/constants.py

@@ -14,13 +14,20 @@ class SamplingPriority(object):
 
 # Datadog trace headers
 class TraceHeader(object):
-    TRACE_ID = 'x-datadog-trace-id'
-    PARENT_ID = 'x-datadog-parent-id'
-    SAMPLING_PRIORITY = 'x-datadog-sampling-priority'
+    TRACE_ID = "x-datadog-trace-id"
+    PARENT_ID = "x-datadog-parent-id"
+    SAMPLING_PRIORITY = "x-datadog-sampling-priority"
 
 
 # X-Ray subsegment to save Datadog trace metadata
 class XraySubsegment(object):
-    NAME = 'datadog-metadata'
-    KEY = 'trace'
-    NAMESPACE = 'datadog'
+    NAME = "datadog-metadata"
+    KEY = "trace"
+    NAMESPACE = "datadog"
+
+
+# Source of datadog context
+class Source(object):
+    XRAY = "xray"
+    EVENT = "event"
+    DDTRACE = "ddtrace"

datadog_lambda/trace_wrapper.py

@@ -0,0 +1,66 @@
+import sys
+from datadog_lambda.constants import SamplingPriority, TraceHeader, Source
+
+
+class TraceWrapper:
+    """
+    TraceWrapper wraps dd-trace, to make this library usable when 
+    dd-trace hasn't been installed/initialised
+    """
+
+    def __init__(self):
+        self._tracer = None
+
+    def extract(self, event):
+        tracer = self._load_tracer()
+        if not tracer:
+            return None
+        return _propagator.extract(event)
+
+    def start_span(self, name, **kwargs):
+        tracer = self._load_tracer()
+        if not tracer:
+            return None
+
+        return self._tracer.start_span(name, **kwargs)
+
+    def _load_tracer(self):
+        if not TraceWrapper.tracer_enabled():
+            return None
+        try:
+            if not self._tracer:
+                from ddtrace import tracer
+                from ddtrace.propagation.http import HTTPPropagator
+
+                self._tracer = tracer
+                self._propagator = HTTPPropagator()
+        except:
+            pass
+        return self._tracer
+
+    @property
+    def trace_context(self):
+        tracer = self._load_tracer()
+        if not tracer:
+            return None
+        span = tracer.current_span()
+        if not span:
+            return None
+
+        parent_id = span.context.span_id
+        trace_id = span.context.trace_id
+        return {
+            "parent_id": parent_id,
+            "trace_id": trace_id,
+            "sampling_priority": SamplingPriority.AUTO_KEEP,
+            "source": Source.DDTRACE,
+        }
+
+    @staticmethod
+    def tracer_enabled():
+        mods = sys.modules.keys()
+        # Check whether user has imported ddtrace
+        return "ddtrace" in mods
+
+
+trace_wrapper = TraceWrapper()

datadog_lambda/tracing.py

@@ -13,7 +13,9 @@
     SamplingPriority,
     TraceHeader,
     XraySubsegment,
+    Source,
 )
+from datadog_lambda.trace_wrapper import trace_wrapper
 
 logger = logging.getLogger(__name__)
 
@@ -38,8 +40,32 @@ def _convert_xray_sampling(xray_sampled):
     """
     Convert X-Ray sampled (True/False) to its Datadog counterpart.
     """
-    return str(SamplingPriority.USER_KEEP) if xray_sampled \
+    return (
+        str(SamplingPriority.USER_KEEP)
+        if xray_sampled
         else str(SamplingPriority.USER_REJECT)
+    )
+
+
+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": Source.XRAY,
+    }
+
+
+def _context_obj_to_headers(obj):
+    return {
+        TraceHeader.TRACE_ID: obj.get("trace_id"),
+        TraceHeader.PARENT_ID: obj.get("parent_id"),
+        TraceHeader.SAMPLING_PRIORITY: obj.get("sampling_priority"),
+    }
 
 
 def extract_dd_trace_context(event):
@@ -54,33 +80,32 @@ def extract_dd_trace_context(event):
     the correct context.
     """
     global dd_trace_context
-    headers = event.get('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 trace_id and parent_id and sampling_priority:
-        logger.debug('Extracted Datadog trace context from headers')
+        logger.debug("Extracted Datadog trace context from headers")
         dd_trace_context = {
-            'trace-id': trace_id,
-            'parent-id': parent_id,
-            'sampling-priority': sampling_priority,
+            "trace-id": trace_id,
+            "parent-id": parent_id,
+            "sampling-priority": sampling_priority,
+            "source": Source.EVENT,
         }
         xray_recorder.begin_subsegment(XraySubsegment.NAME)
         subsegment = xray_recorder.current_subsegment()
         subsegment.put_metadata(
-            XraySubsegment.KEY,
-            dd_trace_context,
-            XraySubsegment.NAMESPACE
+            XraySubsegment.KEY, dd_trace_context, XraySubsegment.NAMESPACE
         )
         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 = {}
-
-    logger.debug('extracted dd trace context %s', 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():
@@ -95,30 +120,25 @@ def get_dd_trace_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),
-        }
+
+    if not dd_trace_context:
+        return None
+    trace_context = _context_obj_to_headers(dd_trace_context)
+    datadog_context = trace_wrapper.trace_context
+    if datadog_context:
+        logger.debug("get_dd_trace_context using dd-trace context")
+        return datadog_context
+    try:
+        xray_context = _get_xray_trace_context()  # xray (sub)segment
+        if xray_context:
+            trace_context[TraceHeader.PARENT_ID] = xray_context["parent_id"]
+    except Exception as e:
+        logger.debug(
+            "get_dd_trace_context couldn't read from segment from x-ray, with error %s"
+            % e
+        )
+    return trace_context
 
 
 def set_correlation_ids():
@@ -130,16 +150,16 @@ def set_correlation_ids():
     TODO: Remove me when Datadog tracer is natively supported in Lambda.
     """
     if not is_lambda_context():
-        logger.debug('set_correlation_ids is only supported in LambdaContext')
+        logger.debug("set_correlation_ids is only supported in LambdaContext")
         return
 
     context = get_dd_trace_context()
 
-    span = tracer.trace('dummy.span')
+    span = tracer.trace("dummy.span")
     span.trace_id = context[TraceHeader.TRACE_ID]
     span.span_id = context[TraceHeader.PARENT_ID]
 
-    logger.debug('correlation ids set')
+    logger.debug("correlation ids set")
 
 
 def inject_correlation_ids():
@@ -153,17 +173,19 @@ def inject_correlation_ids():
     # Override the log format of the AWS provided LambdaLoggerHandler
     root_logger = logging.getLogger()
     for handler in root_logger.handlers:
-        if handler.__class__.__name__ == 'LambdaLoggerHandler':
-            handler.setFormatter(logging.Formatter(
-                '[%(levelname)s]\t%(asctime)s.%(msecs)dZ\t%(aws_request_id)s\t'
-                '[dd.trace_id=%(dd.trace_id)s dd.span_id=%(dd.span_id)s]\t%(message)s\n',
-                '%Y-%m-%dT%H:%M:%S'
-            ))
+        if handler.__class__.__name__ == "LambdaLoggerHandler":
+            handler.setFormatter(
+                logging.Formatter(
+                    "[%(levelname)s]\t%(asctime)s.%(msecs)dZ\t%(aws_request_id)s\t"
+                    "[dd.trace_id=%(dd.trace_id)s dd.span_id=%(dd.span_id)s]\t%(message)s\n",
+                    "%Y-%m-%dT%H:%M:%S",
+                )
+            )
 
     # Patch `logging.Logger.makeRecord` to actually inject correlation ids
     patch(logging=True)
 
-    logger.debug('logs injection configured')
+    logger.debug("logs injection configured")
 
 
 def is_lambda_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,
@@ -18,7 +18,10 @@
     extract_dd_trace_context,
     set_correlation_ids,
     inject_correlation_ids,
+    get_dd_trace_context,
 )
+from datadog_lambda.trace_wrapper import trace_wrapper
+from datadog_lambda.constants import Source
 
 
 logger = logging.getLogger(__name__)
@@ -39,7 +42,6 @@ def my_lambda_handle(event, context):
 
 
 class _NoopDecorator(object):
-
     def __init__(self, func):
         self.func = func
 
@@ -82,6 +84,8 @@ def __init__(self, func):
             self.logs_injection = (
                 os.environ.get("DD_LOGS_INJECTION", "true").lower() == "true"
             )
+            self.handler_name = os.environ.get("_HANDLER", "handler")
+            self.function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", "function")
 
             # Inject trace correlation ids to logs
             if self.logs_injection:
@@ -106,10 +110,33 @@ 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)
+            span_context = None
+            if dd_context["source"] == Source.EVENT:
+                span_context = trace_wrapper.extract(dd_context)
+
+            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 = trace_wrapper.start_span("aws.lambda", **args)
+            if self.span:
+                self.span.set_tags(tags)
 
             # Set log correlation ids using extracted trace context
             set_correlation_ids()
@@ -119,6 +146,8 @@ def _before(self, event, context):
 
     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")