8000 Aj/proactive init (#343) · sagarkvv/datadog-lambda-python@9aef51d · GitHub
[go: up one dir, main page]
Skip to content

Commit 9aef51d

Browse files
astuyveastuyve
authored
Aj/proactive init (DataDog#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
1 parent ae76b33 commit 9aef51d

File tree

5 files changed

+83
-10
lines changed

5 files changed

+83
-10
lines changed

‎datadog_lambda/cold_start.py‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,28 @@
66
logger = logging.getLogger(__name__)
77

88
_cold_start = True
9+
_proactive_initialization = False
910
_lambda_container_initialized = False
1011

1112

12-
def set_cold_start():
13+
def set_cold_start(init_timestamp_ns):
1314
"""Set the value of the cold start global
1415
1516
This should be executed once per Lambda execution before the execution
1617
"""
1718
global _cold_start
1819
global _lambda_container_initialized
19-
_cold_start = not _lambda_container_initialized
20+
global _proactive_initialization
21+
if not _lambda_container_initialized:
22+
now = time.time_ns()
23+
if (now - init_timestamp_ns) // 1_000_000_000 > 10:
24+
_cold_start = False
25+
_proactive_initialization = True
26+
else:
27+
_cold_start = not _lambda_container_initialized
28+
else:
29+
_cold_start = False
30+
_proactive_initialization = False
2031
_lambda_container_initialized = True
2132

2233

@@ -25,11 +36,25 @@ def is_cold_start():
2536
return _cold_start
2637

2738

39+
def is_proactive_init():
40+
"""Returns the value of the global proactive_initialization"""
41+
return _proactive_initialization
42+
43+
44+
def is_new_sandbox():
45+
return is_cold_start() or is_proactive_init()
46+
47+
2848
def get_cold_start_tag():
2949
"""Returns the cold start tag to be used in metrics"""
3050
return "cold_start:{}".format(str(is_cold_start()).lower())
3151

3252

53+
def get_proactive_init_tag():
54+
"""Returns the proactive init tag to be used in metrics"""
55+
return "proactive_initialization:{}".format(str(is_proactive_init()).lower())
56+
57+
3358
class ImportNode(object):
3459
def __init__(self, module_name, full_file_path, start_time_ns, end_time_ns=None):
3560
self.module_name = module_name
@@ -115,7 +140,7 @@ def wrapped_find_spec(*args, **kwargs):
115140

116141
def initialize_cold_start_tracing():
117142
if (
118-
is_cold_start()
143+
is_new_sandbox()
119144
and os.environ.get("DD_TRACE_ENABLED", "true").lower() == "true"
120145
and os.environ.get("DD_COLD_START_TRACING", "true").lower() == "true"
121146
):

‎datadog_lambda/tracing.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1211,6 +1211,7 @@ def create_function_execution_span(
12111211
context,
12121212
function_name,
12131213
is_cold_start,
1214+
is_proactive_init,
12141215
trace_context_source,
12151216
merge_xray_traces,
12161217
trigger_tags,
@@ -1235,6 +1236,8 @@ def create_function_execution_span(
12351236
"dd_trace": ddtrace_version,
12361237
"span.name": "aws.lambda",
12371238
}
1239+
if is_proactive_init:
1240+
tags["proactive_initialization"] = str(is_proactive_init).lower()
12381241
if trace_context_source == TraceContextSource.XRAY and merge_xray_traces:
12391242
tags["_dd.parent_source"] = trace_context_source
12401243
tags.update(trigger_tags)

‎datadog_lambda/wrapper.py‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,13 @@
1111
from time import time_ns
1212

1313
from datadog_lambda.extension import should_use_extension, flush_extension
14-
from datadog_lambda.cold_start import set_cold_start, is_cold_start, ColdStartTracer
14+
from datadog_lambda.cold_start import (
15+
set_cold_start,
16+
is_cold_start,
17+
is_proactive_init,
18+
is_new_sandbox,
19+
ColdStartTracer,
20+
)
1521
from datadog_lambda.constants import (
1622
TraceContextSource,
1723
XraySubsegment,
@@ -72,6 +78,8 @@
7278

7379
env_env_var = os.environ.get(DD_ENV, None)
7480

81+
init_timestamp_ns = time_ns()
82+
7583
"""
7684
Usage:
7785
@@ -245,7 +253,7 @@ def _inject_authorizer_span_headers(self, request_id):
245253
def _before(self, event, context):
246254
try:
247255
self.response = None
248-
set_cold_start()
256+
set_cold_start(init_timestamp_ns)
249257
submit_invocations_metric(context)
250258
self.trigger_tags = extract_trigger_tags(event, context)
251259
# Extract Datadog trace context and source from incoming requests
@@ -272,14 +280,15 @@ def _before(self, event, context):
272280
context,
273281
self.function_name,
274282
is_cold_start(),
283+
is_proactive_init(),
275284
trace_context_source,
276285
self.merge_xray_traces,
277286
self.trigger_tags,
278287
parent_span=self.inferred_span,
279288
)
280289
else:
281290
set_correlation_ids()
282-
if profiling_env_var and is_cold_start():
291+
if profiling_env_var and is_new_sandbox():
283292
self.prof.start(stop_on_exit=False, profile_children=True)
284293
logger.debug("datadog_lambda_wrapper _before() done")
285294
except Exception:
@@ -299,7 +308,7 @@ def _after(self, event, context):
299308
self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
300309
)
301310
should_trace_cold_start = (
302-
dd_tracing_enabled and self.cold_start_tracing and is_cold_start()
311+
dd_tracing_enabled and self.cold_start_tracing and is_new_sandbox()
303312
)
304313
if should_trace_cold_start:
305314
trace_ctx = tracer.current_trace_context()

‎tests/test_cold_start.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import time
12
import unittest
23
import datadog_lambda.cold_start as cold_start
34
from sys import modules, meta_path
@@ -6,6 +7,34 @@
67

78

89
class TestColdStartTracingSetup(unittest.TestCase):
10+
def test_proactive_init(self):
11+
cold_start._cold_start = True
12+
cold_start._proactive_initialization = False
13+
cold_start._lambda_container_initialized = False
14+
fifteen_seconds_ago = time.time_ns() - 15_000_000_000
15+
cold_start.set_cold_start(fifteen_seconds_ago)
16+
self.assertTrue(cold_start.is_proactive_init())
17+
self.assertTrue(cold_start.is_new_sandbox())
18+
self.assertFalse(cold_start.is_cold_start())
19+
self.assertEqual(
20+
cold_start.get_proactive_init_tag(), "proactive_initialization:true"
21+
)
22+
self.assertEqual(cold_start.get_cold_start_tag(), "cold_start:false")
23+
24+
def test_cold_start(self):
25+
cold_start._cold_start = True
26+
cold_start._proactive_initialization = False
27+
cold_start._lambda_container_initialized = False
28+
one_second_ago = time.time_ns() - 1_000_000_000
29+
cold_start.set_cold_start(one_second_ago)
30+
self.assertFalse(cold_start.is_proactive_init())
31+
self.assertTrue(cold_start.is_new_sandbox())
32+
self.assertTrue(cold_start.is_cold_start())
33+
self.assertEqual(
34+
cold_start.get_proactive_init_tag(), "proactive_initialization:false"
35+
)
36+
self.assertEqual(cold_start.get_cold_start_tag(), "cold_start:true")
37+
938
def test_initialize_cold_start_tracing(self):
1039
cold_start._cold_start = True
1140
cold_start.initialize_cold_start_tracing() # testing double wrapping

‎tests/test_tracing.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -526,7 +526,9 @@ def test_set_correlation_ids_handle_empty_trace_context(self):
526526
class TestFunctionSpanTags(unittest.TestCase):
527527
def test_function(self):
528528
ctx = get_mock_context()
529-
span = create_function_execution_span(ctx, "", False, {"source": ""}, False, {})
529+
span = create_function_execution_span(
530+
ctx, "", False, False, {"source": ""}, False, {}
531+
)
530532
self.assertEqual(span.get_tag("function_arn"), function_arn)
531533
self.assertEqual(span.get_tag("function_version"), "$LATEST")
532534
self.assertEqual(span.get_tag("resource_names"), "Function")
@@ -537,7 +539,9 @@ def test_function_with_version(self):
537539
ctx = get_mock_context(
538540
invoked_function_arn=function_arn + ":" + function_version
539541
)
540-
span = create_function_execution_span(ctx, "", False, {"source": ""}, False, {})
542+
span = create_function_execution_span(
543+
ctx, "", False, False, {"source": ""}, False, {}
544+
)
541545
self.assertEqual(span.get_tag("function_arn"), function_arn)
542546
self.assertEqual(span.get_tag("function_version"), function_version)
543547
self.assertEqual(span.get_tag("resource_names"), "Function")
@@ -546,7 +550,9 @@ def test_function_with_version(self):
546550
def test_function_with_alias(self):
547551
function_alias = "alias"
548552
ctx = get_mock_context(invoked_function_arn=function_arn + ":" + function_alias)
549-
span = create_function_execution_span(ctx, "", False, {"source": ""}, False, {})
553+
span = create_function_execution_span(
554+
ctx, "", False, False, {"source": ""}, False, {}
555+
)
550556
self.assertEqual(span.get_tag("function_arn"), function_arn)
551557
self.assertEqual(span.get_tag("function_version"), function_alias)
552558
self.assertEqual(span.get_tag("resource_names"), "Function")
@@ -558,6 +564,7 @@ def test_function_with_trigger_tags(self):
558564
ctx,
559565
"",
560566
False,
567+
False,
561568
{"source": ""},
562569
False,
563570
{"function_trigger.event_source": "cloudwatch-logs"},

0 commit comments

Comments
 (0)
0