Aj/proactive init (#343)

* feat: support proactive initialization

* feat: 10 sec max init time

* fix: specs

* fix: format, fix missing import

* feat: Only tag if it is a proactive init

* feat: add units

Author: astuyve

datadog_lambda/cold_start.py

@@ -6,17 +6,28 @@
 logger = logging.getLogger(__name__)
 
 _cold_start = True
+_proactive_initialization = False
 _lambda_container_initialized = False
 
 
-def set_cold_start():
+def set_cold_start(init_timestamp_ns):
     """Set the value of the cold start global
 
     This should be executed once per Lambda execution before the execution
     """
     global _cold_start
     global _lambda_container_initialized
-    _cold_start = not _lambda_container_initialized
+    global _proactive_initialization
+    if not _lambda_container_initialized:
+        now = time.time_ns()
+        if (now - init_timestamp_ns) // 1_000_000_000 > 10:
+            _cold_start = False
+            _proactive_initialization = True
+        else:
+            _cold_start = not _lambda_container_initialized
+    else:
+        _cold_start = False
+        _proactive_initialization = False
     _lambda_container_initialized = True
 
 
@@ -25,11 +36,25 @@ def is_cold_start():
     return _cold_start
 
 
+def is_proactive_init():
+    """Returns the value of the global proactive_initialization"""
+    return _proactive_initialization
+
+
+def is_new_sandbox():
+    return is_cold_start() or is_proactive_init()
+
+
 def get_cold_start_tag():
     """Returns the cold start tag to be used in metrics"""
     return "cold_start:{}".format(str(is_cold_start()).lower())
 
 
+def get_proactive_init_tag():
+    """Returns the proactive init tag to be used in metrics"""
+    return "proactive_initialization:{}".format(str(is_proactive_init()).lower())
+
+
 class ImportNode(object):
     def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None):
         self.module_name = module_name
@@ -115,7 +140,7 @@ def wrapped_find_spec(*args, **kwargs):
 
 def initialize_cold_start_tracing():
     if (
-        is_cold_start()
+        is_new_sandbox()
         and os.environ.get("DD_TRACE_ENABLED", "true").lower() == "true"
         and os.environ.get("DD_COLD_START_TRACING", "true").lower() == "true"
     ):

datadog_lambda/tracing.py

@@ -1211,6 +1211,7 @@ def create_function_execution_span(
     context,
     function_name,
     is_cold_start,
+    is_proactive_init,
     trace_context_source,
     merge_xray_traces,
     trigger_tags,
@@ -1235,6 +1236,8 @@ def create_function_execution_span(
             "dd_trace": ddtrace_version,
             "span.name": "aws.lambda",
         }
+    if is_proactive_init:
+        tags["proactive_initialization"] = str(is_proactive_init).lower()
     if trace_context_source == TraceContextSource.XRAY and merge_xray_traces:
         tags["_dd.parent_source"] = trace_context_source
     tags.update(trigger_tags)

datadog_lambda/wrapper.py

@@ -11,7 +11,13 @@
 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, ColdStartTracer
+from datadog_lambda.cold_start import (
+    set_cold_start,
+    is_cold_start,
+    is_proactive_init,
+    is_new_sandbox,
+    ColdStartTracer,
+)
 from datadog_lambda.constants import (
     TraceContextSource,
     XraySubsegment,
@@ -72,6 +78,8 @@
 
 env_env_var = os.environ.get(DD_ENV, None)
 
+init_timestamp_ns = time_ns()
+
 """
 Usage:
 
@@ -245,7 +253,7 @@ def _inject_authorizer_span_headers(self, request_id):
     def _before(self, event, context):
         try:
             self.response = None
-            set_cold_start()
+            set_cold_start(init_timestamp_ns)
             submit_invocations_metric(context)
             self.trigger_tags = extract_trigger_tags(event, context)
             # Extract Datadog trace context and source from incoming requests
@@ -272,14 +280,15 @@ def _before(self, event, context):
                     context,
                     self.function_name,
                     is_cold_start(),
+                    is_proactive_init(),
                     trace_context_source,
                     self.merge_xray_traces,
                     self.trigger_tags,
                     parent_span=self.inferred_span,
                 )
             else:
                 set_correlation_ids()
-            if profiling_env_var and is_cold_start():
+            if profiling_env_var and is_new_sandbox():
                 self.prof.start(stop_on_exit=False, profile_children=True)
             logger.debug("datadog_lambda_wrapper _before() done")
         except Exception:
@@ -299,7 +308,7 @@ def _after(self, event, context):
                     self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
                 )
             should_trace_cold_start = (
-                dd_tracing_enabled and self.cold_start_tracing and is_cold_start()
+                dd_tracing_enabled and self.cold_start_tracing and is_new_sandbox()
             )
             if should_trace_cold_start:
                 trace_ctx = tracer.current_trace_context()

tests/test_cold_start.py

@@ -1,3 +1,4 @@
+import time
 import unittest
 import datadog_lambda.cold_start as cold_start
 from sys import modules, meta_path
@@ -6,6 +7,34 @@
 
 
 class TestColdStartTracingSetup(unittest.TestCase):
+    def test_proactive_init(self):
+        cold_start._cold_start = True
+        cold_start._proactive_initialization = False
+        cold_start._lambda_container_initialized = False
+        fifteen_seconds_ago = time.time_ns() - 15_000_000_000
+        cold_start.set_cold_start(fifteen_seconds_ago)
+        self.assertTrue(cold_start.is_proactive_init())
+        self.assertTrue(cold_start.is_new_sandbox())
+        self.assertFalse(cold_start.is_cold_start())
+        self.assertEqual(
+            cold_start.get_proactive_init_tag(), "proactive_initialization:true"
+        )
+        self.assertEqual(cold_start.get_cold_start_tag(), "cold_start:false")
+
+    def test_cold_start(self):
+        cold_start._cold_start = True
+        cold_start._proactive_initialization = False
+        cold_start._lambda_container_initialized = False
+        one_second_ago = time.time_ns() - 1_000_000_000
+        cold_start.set_cold_start(one_second_ago)
+        self.assertFalse(cold_start.is_proactive_init())
+        self.assertTrue(cold_start.is_new_sandbox())
+        self.assertTrue(cold_start.is_cold_start())
+        self.assertEqual(
+            cold_start.get_proactive_init_tag(), "proactive_initialization:false"
+        )
+        self.assertEqual(cold_start.get_cold_start_tag(), "cold_start:true")
+
     def test_initialize_cold_start_tracing(self):
         cold_start._cold_start = True
         cold_start.initialize_cold_start_tracing()  # testing double wrapping

tests/test_tracing.py

@@ -526,7 +526,9 @@ def test_set_correlation_ids_handle_empty_trace_context(self):
 class TestFunctionSpanTags(unittest.TestCase):
     def test_function(self):
         ctx = get_mock_context()
-        span = create_function_execution_span(ctx, "", False, {"source": ""}, False, {})
+        span = create_function_execution_span(
+            ctx, "", False, False, {"source": ""}, False, {}
+        )
         self.assertEqual(span.get_tag("function_arn"), function_arn)
         self.assertEqual(span.get_tag("function_version"), "$LATEST")
         self.assertEqual(span.get_tag("resource_names"), "Function")
@@ -537,7 +539,9 @@ def test_function_with_version(self):
         ctx = get_mock_context(
             invoked_function_arn=function_arn + ":" + function_version
         )
-        span = create_function_execution_span(ctx, "", False, {"source": ""}, False, {})
+        span = create_function_execution_span(
+            ctx, "", False, False, {"source": ""}, False, {}
+        )
         self.assertEqual(span.get_tag("function_arn"), function_arn)
         self.assertEqual(span.get_tag("function_version"), function_version)
         self.assertEqual(span.get_tag("resource_names"), "Function")
@@ -546,7 +550,9 @@ def test_function_with_version(self):
     def test_function_with_alias(self):
         function_alias = "alias"
         ctx = get_mock_context(invoked_function_arn=function_arn + ":" + function_alias)
-        span = create_function_execution_span(ctx, "", False, {"source": ""}, False, {})
+        span = create_function_execution_span(
+            ctx, "", False, False, {"source": ""}, False, {}
+        )
         self.assertEqual(span.get_tag("function_arn"), function_arn)
         self.assertEqual(span.get_tag("function_version"), function_alias)
         self.assertEqual(span.get_tag("resource_names"), "Function")
@@ -558,6 +564,7 @@ def test_function_with_trigger_tags(self):
             ctx,
             "",
             False,
+            False,
             {"source": ""},
             False,
             {"function_trigger.event_source": "cloudwatch-logs"},