diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py index a10a2ad7a..0a3ba34d1 100644 --- a/datadog_lambda/cold_start.py +++ b/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" ): diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index f356e13fa..198332a3e 100644 --- a/datadog_lambda/tracing.py +++ b/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) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 6e8384214..7b9e87f9f 100644 --- a/datadog_lambda/wrapper.py +++ b/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,6 +280,7 @@ 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, @@ -279,7 +288,7 @@ def _before(self, event, context): ) 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() diff --git a/tests/test_cold_start.py b/tests/test_cold_start.py index 56636deca..65193e1de 100644 --- a/tests/test_cold_start.py +++ b/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 diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 74cebac4c..f5f41f94e 100644 --- a/tests/test_tracing.py +++ b/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"},