cold start tracing (python)

datadog_lambda/__init__.py

@@ -1,3 +1,9 @@
+import os
+import logging
+from datadog_lambda.cold_start import initialize_cold_start_tracing
+
+initialize_cold_start_tracing()
+
 # The minor version corresponds to the Lambda layer version.
 # E.g.,, version 0.5.0 gets packaged into layer version 5.
 try:
@@ -7,8 +13,5 @@
 
 __version__ = importlib_metadata.version(__name__)
 
-import os
-import logging
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.getLevelName(os.environ.get("DD_LOG_LEVEL", "INFO").upper()))

datadog_lambda/cold_start.py

@@ -1,3 +1,10 @@
+import time
+import os
+from typing import List, Hashable
+import logging
+
+logger = logging.getLogger(__name__)
+
 _cold_start = True
 _lambda_container_initialized = False
 
@@ -21,3 +28,196 @@ def is_cold_start():
 def get_cold_start_tag():
     """Returns the cold start tag to be used in metrics"""
     return "cold_start:{}".format(str(is_cold_start()).lower())
+
+
+class ImportNode(object):
+    def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None):
+        self.module_name = module_name
+        self.full_file_path = full_file_path
+        self.start_time_ns = start_time_ns
+        self.end_time_ns = end_time_ns
+        self.children = []
+
+
+root_nodes: List[ImportNode] = []
+import_stack: List[ImportNode] = []
+already_wrapped_loaders = set()
+
+
+def reset_node_stacks():
+    global root_nodes
+    root_nodes = []
+    global import_stack
+    import_stack = []
+
+
+def push_node(module_name, file_path):
+    node = ImportNode(module_name, file_path, time.time_ns())
+    global import_stack
+    if import_stack:
+        import_stack[-1].children.append(node)
+    import_stack.append(node)
+
+
+def pop_node(module_name):
+    global import_stack
+    if not import_stack:
+        return
+    node = import_stack.pop()
+    if node.module_name != module_name:
+        return
+    end_time_ns = time.time_ns()
+    node.end_time_ns = end_time_ns
+    if not import_stack:  # import_stack empty, a root node has been found
+        global root_nodes
+        root_nodes.append(node)
+
+
+def wrap_exec_module(original_exec_module):
+    def wrapped_method(module):
+        should_pop = False
+        try:
+            spec = module.__spec__
+            push_node(spec.name, spec.origin)
+            should_pop = True
+        except Exception:
+            pass
+        try:
+            return original_exec_module(module)
+        finally:
+            if should_pop:
+                pop_node(spec.name)
+
+    return wrapped_method
+
+
+def wrap_find_spec(original_find_spec):
+    def wrapped_find_spec(*args, **kwargs):
+        spec = original_find_spec(*args, **kwargs)
+        if spec is None:
+            return None
+        loader = getattr(spec, "loader", None)
+        if (
+            loader is not None
+            and isinstance(loader, Hashable)
+            and loader not in already_wrapped_loaders
+        ):
+            if hasattr(loader, "exec_module"):
+                try:
+                    loader.exec_module = wrap_exec_module(loader.exec_module)
+                    already_wrapped_loaders.add(loader)
+                except Exception as e:
+                    logger.debug("Failed to wrap the loader. %s", e)
+        return spec
+
+    return wrapped_find_spec
+
+
+def initialize_cold_start_tracing():
+    if (
+        is_cold_start()
+        and os.environ.get("DD_TRACE_ENABLED", "true").lower() == "true"
+        and os.environ.get("DD_COLD_START_TRACING", "true").lower() == "true"
+    ):
+        from sys import version_info, meta_path
+
+        if version_info >= (3, 7):  # current implementation only support version > 3.7
+            for importer in meta_path:
+                try:
+                    importer.find_spec = wrap_find_spec(importer.find_spec)
+                except Exception:
+                    pass
+
+
+class ColdStartTracer(object):
+    def __init__(
+        self,
+        tracer,
+        function_name,
+        cold_start_span_finish_time_ns,
+        trace_ctx,
+        min_duration_ms: int,
+        ignored_libs: List[str] = [],
+    ):
+        self._tracer = tracer
+        self.function_name = function_name
+        self.cold_start_span_finish_time_ns = cold_start_span_finish_time_ns
+        self.min_duration_ms = min_duration_ms
+        self.trace_ctx = trace_ctx
+        self.ignored_libs = ignored_libs
+        self.need_to_reactivate_context = True
+
+    def trace(self, root_nodes: List[ImportNode] = root_nodes):
+        if not root_nodes:
+            return
+        cold_start_span_start_time_ns = root_nodes[0].start_time_ns
+        cold_start_span = self.create_cold_start_span(cold_start_span_start_time_ns)
+        while root_nodes:
+            root_node = root_nodes.pop()
+            self.trace_tree(root_node, cold_start_span)
+        self.finish_span(cold_start_span, self.cold_start_span_finish_time_ns)
+
+    def trace_tree(self, import_node: ImportNode, parent_span):
+        if (
+            import_node.end_time_ns - import_node.start_time_ns
+            < self.min_duration_ms * 1e6
+            or import_node.module_name in self.ignored_libs
+        ):
+            return
+
+        span = self.start_span(
+            "aws.lambda.import", import_node.module_name, import_node.start_time_ns
+        )
+        tags = {
+            "resource_names": import_node.module_name,
+            "resource.name": import_node.module_name,
+            "filename": import_node.full_file_path,
+            "operation_name": self.get_operation_name(import_node.full_file_path),
+        }
+        span.set_tags(tags)
+        if parent_span:
+            span.parent_id = parent_span.span_id
+        for child_node in import_node.children:
+            self.trace_tree(child_node, span)
+        self.finish_span(span, import_node.end_time_ns)
+
+    def create_cold_start_span(self, start_time_ns):
+        span = self.start_span("aws.lambda.load", self.function_name, start_time_ns)
+        tags = {
+            "resource_names": self.function_name,
+            "resource.name": self.function_name,
+            "operation_name": "aws.lambda.load",
+        }
+        span.set_tags(tags)
+        return span
+
+    def start_span(self, span_type, resource, start_time_ns):
+        if self.need_to_reactivate_context:
+            self._tracer.context_provider.activate(
+                self.trace_ctx
+            )  # reactivate required after each finish() call
+            self.need_to_reactivate_context = False
+        span_kwargs = {
+            "service": "aws.lambda",
+            "resource": resource,
+            "span_type": span_type,
+        }
+        span = self._tracer.trace(span_type, **span_kwargs)
+        span.start_ns = start_time_ns
+        return span
+
+    def finish_span(self, span, finish_time_ns):
+        span.finish(finish_time_ns / 1e9)
+        self.need_to_reactivate_context = True
+
+    def get_operation_name(self, filename: str):
+        if filename is None:
+            return "aws.lambda.import_core_module"
+        if not isinstance(filename, str):
+            return "aws.lambda.import"
+        if filename.startswith("/opt/"):
+            return "aws.lambda.import_layer"
+        elif filename.startswith("/var/lang/"):
+            return "aws.lambda.import_runtime"
+        else:
+            return "aws.lambda.import"

datadog_lambda/wrapper.py

@@ -2,7 +2,6 @@
 # under the Apache License Version 2.0.
 # This product includes software developed at Datadog (https://www.datadoghq.com/).
 # Copyright 2019 Datadog, Inc.
-
 import base64
 import os
 import logging
@@ -12,7 +11,7 @@
 from time import time_ns
 
 from datadog_lambda.extension import should_use_extension, flush_extension
-from datadog_lambda.cold_start import set_cold_start, is_cold_start
+from datadog_lambda.cold_start import set_cold_start, is_cold_start, ColdStartTracer
 from datadog_lambda.constants import (
     TraceContextSource,
     XraySubsegment,
@@ -38,6 +37,7 @@
     create_inferred_span,
     InferredSpanInfo,
     is_authorizer_response,
+    tracer,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -131,6 +131,28 @@ def __init__(self, func):
             self.decode_authorizer_context = (
                 os.environ.get("DD_DECODE_AUTHORIZER_CONTEXT", "true").lower() == "true"
             )
+            self.cold_start_tracing = (
+                os.environ.get("DD_COLD_START_TRACING", "true").lower() == "true"
+            )
+            self.min_cold_start_trace_duration = 3
+            if "DD_MIN_COLD_START_DURATION" in os.environ:
+                try:
+                    self.min_cold_start_trace_duration = int(
+                        os.environ["DD_MIN_COLD_START_DURATION"]
+                    )
+                except Exception:
+                    logger.debug("Malformatted env DD_MIN_COLD_START_DURATION")
+            self.cold_start_trace_skip_lib = [
+                "ddtrace.internal.compat",
+                "ddtrace.filters",
+            ]
+            if "DD_COLD_START_TRACE_SKIP_LIB" in os.environ:
+                try:
+                    self.cold_start_trace_skip_lib = os.environ[
+                        "DD_COLD_START_TRACE_SKIP_LIB"
+                    ].split(",")
+                except Exception:
+                    logger.debug("Malformatted for env DD_COLD_START_TRACE_SKIP_LIB")
             self.response = None
             if profiling_env_var:
                 self.prof = profiler.Profiler(env=env_env_var, service=service_env_var)
@@ -257,6 +279,11 @@ def _after(self, event, context):
                 create_dd_dummy_metadata_subsegment(
                     self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
                 )
+            should_trace_cold_start = (
+                dd_tracing_enabled and self.cold_start_tracing and is_cold_start()
+            )
+            if should_trace_cold_start:
+                trace_ctx = tracer.current_trace_context()
 
             if self.span:
                 if dd_capture_lambda_payload_enabled:
@@ -276,6 +303,20 @@ def _after(self, event, context):
                 else:
                     self.inferred_span.finish()
 
+            if should_trace_cold_start:
+                try:
+                    following_span = self.span or self.inferred_span
+                    ColdStartTracer(
+                        tracer,
+                        self.function_name,
+                        following_span.start_ns,
+                        trace_ctx,
+                        self.min_cold_start_trace_duration,
+                        self.cold_start_trace_skip_lib,
+                    ).trace()
+                except Exception as e:
+                    logger.debug("Failed to create cold start spans. %s", e)
+
             if not self.flush_to_log or should_use_extension:
                 flush_stats()
             if should_use_extension:

tests/integration/serverless.yml

@@ -11,6 +11,7 @@ provider:
     DD_TRACE_ENABLED: true
     DD_API_KEY: ${env:DD_API_KEY}
     DD_TRACE_MANAGED_SERVICES: true
+    DD_COLD_START_TRACING: false
   timeout: 15
   deploymentBucket:
     name: integration-tests-serververless-deployment-bucket