From c64baa8e73a456db2cb8fcd6d67447d4e75e62b3 Mon Sep 17 00:00:00 2001 From: Yiming Luo Date: Fri, 13 Mar 2026 16:12:49 -0400 Subject: [PATCH 1/5] Add durable_function_first_invocation tag to aws.lambda spans Extends extract_durable_function_tags() to accept context.state and sets durable_function_first_invocation ("true"/"false") using not state.is_replaying(). wrapper._before() passes context.state to the function. Co-Authored-By: Claude Sonnet 4.6 --- datadog_lambda/durable.py | 12 ++++++++++-- datadog_lambda/wrapper.py | 6 ++++-- tests/test_durable.py | 38 +++++++++++++++++++++++++++++++++++++- 3 files changed, 51 insertions(+), 5 deletions(-) diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index e9443f92..f39f7b45 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -24,7 +24,7 @@ def _parse_durable_execution_arn(arn): return execution_name, execution_id -def extract_durable_function_tags(event): +def extract_durable_function_tags(event, state=None): """ Extracts durable function tags from the Lambda event payload. Returns a dict with durable function tags, or an empty dict if the event @@ -43,7 +43,15 @@ def extract_durable_function_tags(event): return {} execution_name, execution_id = parsed - return { + try: + first_invocation = str(not state.is_replaying()).lower() + except Exception: + logger.debug("Failed to compute durable_function_first_invocation") + first_invocation = None + tags = { "durable_function_execution_name": execution_name, "durable_function_execution_id": execution_id, } + if first_invocation is not None: + tags["durable_function_first_invocation"] = first_invocation + return tags diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index c174a501..30f8ffa4 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -153,7 +153,7 @@ def __init__(self, func): if config.trace_extractor: extractor_parts = config.trace_extractor.rsplit(".", 1) if len(extractor_parts) == 2: - (mod_name, extractor_name) = extractor_parts + mod_name, extractor_name = extractor_parts modified_extractor_name = modify_module_name(mod_name) extractor_module = import_module(modified_extractor_name) self.trace_extractor = getattr(extractor_module, extractor_name) @@ -244,7 +244,9 @@ def _before(self, event, context): submit_invocations_metric(context) self.trigger_tags = extract_trigger_tags(event, context) - self.durable_function_tags = extract_durable_function_tags(event) + self.durable_function_tags = extract_durable_function_tags( + event, context.state + ) # Extract Datadog trace context and source from incoming requests dd_context, trace_context_source, event_source = extract_dd_trace_context( event, diff --git a/tests/test_durable.py b/tests/test_durable.py index 60914934..d0f742d4 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -3,6 +3,7 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. import unittest +from unittest.mock import MagicMock from datadog_lambda.durable import ( _parse_durable_execution_arn, @@ -51,12 +52,47 @@ def test_extracts_tags_from_event_with_durable_execution_arn(self): "CheckpointToken": "some-token", "InitialExecutionState": {"Operations": []}, } - result = extract_durable_function_tags(event) + state = MagicMock() + state.is_replaying.return_value = True + result = extract_durable_function_tags(event, state) + self.assertEqual( + result, + { + "durable_function_execution_name": "my-execution", + "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", + "durable_function_first_invocation": "false", + }, + ) + + def test_sets_first_invocation_true_when_not_replaying(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + } + state = MagicMock() + state.is_replaying.return_value = False + result = extract_durable_function_tags(event, state) + self.assertEqual( + result, + { + "durable_function_execution_name": "my-execution", + "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", + "durable_function_first_invocation": "true", + }, + ) + + def test_sets_first_invocation_false_when_replaying(self): + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + } + state = MagicMock() + state.is_replaying.return_value = True + result = extract_durable_function_tags(event, state) self.assertEqual( result, { "durable_function_execution_name": "my-execution", "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", + "durable_function_first_invocation": "false", }, ) From 781646a51c3404d401e05453138f103434f50a96 Mon Sep 17 00:00:00 2001 From: Yiming Luo Date: Thu, 19 Mar 2026 14:11:33 +0100 Subject: [PATCH 2/5] Add durable_function_first_invocation tag to aws.lambda spans Extends extract_durable_function_tags() to set durable_function_first_invocation ("true"/"false") by checking len(InitialExecutionState.Operations) <= 1, mirroring the SDK's own replay detection logic (ReplayStatus.REPLAY when len > 1). Co-Authored-By: Claude Sonnet 4.6 --- datadog_lambda/durable.py | 15 ++++-------- datadog_lambda/wrapper.py | 6 ++--- tests/test_durable.py | 51 +++++++++++++++++++++++++++------------ 3 files changed, 43 insertions(+), 29 deletions(-) diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index f39f7b45..c900f1de 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -24,7 +24,7 @@ def _parse_durable_execution_arn(arn): return execution_name, execution_id -def extract_durable_function_tags(event, state=None): +def extract_durable_function_tags(event): """ Extracts durable function tags from the Lambda event payload. Returns a dict with durable function tags, or an empty dict if the event @@ -43,15 +43,10 @@ def extract_durable_function_tags(event, state=None): return {} execution_name, execution_id = parsed - try: - first_invocation = str(not state.is_replaying()).lower() - except Exception: - logger.debug("Failed to compute durable_function_first_invocation") - first_invocation = None - tags = { + operations = event.get("InitialExecutionState", {}).get("Operations", []) + is_first_invocation = len(operations) <= 1 + return { "durable_function_execution_name": execution_name, "durable_function_execution_id": execution_id, + "durable_function_first_invocation": str(is_first_invocation).lower(), } - if first_invocation is not None: - tags["durable_function_first_invocation"] = first_invocation - return tags diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 30f8ffa4..c174a501 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -153,7 +153,7 @@ def __init__(self, func): if config.trace_extractor: extractor_parts = config.trace_extractor.rsplit(".", 1) if len(extractor_parts) == 2: - mod_name, extractor_name = extractor_parts + (mod_name, extractor_name) = extractor_parts modified_extractor_name = modify_module_name(mod_name) extractor_module = import_module(modified_extractor_name) self.trace_extractor = getattr(extractor_module, extractor_name) @@ -244,9 +244,7 @@ def _before(self, event, context): submit_invocations_metric(context) self.trigger_tags = extract_trigger_tags(event, context) - self.durable_function_tags = extract_durable_function_tags( - event, context.state - ) + self.durable_function_tags = extract_durable_function_tags(event) # Extract Datadog trace context and source from incoming requests dd_context, trace_context_source, event_source = extract_dd_trace_context( event, diff --git a/tests/test_durable.py b/tests/test_durable.py index d0f742d4..dda36cc9 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -3,7 +3,6 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. import unittest -from unittest.mock import MagicMock from datadog_lambda.durable import ( _parse_durable_execution_arn, @@ -46,31 +45,31 @@ def test_works_with_numeric_version_qualifier(self): class TestExtractDurableFunctionTags(unittest.TestCase): - def test_extracts_tags_from_event_with_durable_execution_arn(self): + def test_sets_first_invocation_true_when_only_execution_operation(self): + # One operation (the EXECUTION op itself) → not replaying → first invocation event = { "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", "CheckpointToken": "some-token", - "InitialExecutionState": {"Operations": []}, + "InitialExecutionState": {"Operations": [{"OperationType": "EXECUTION"}]}, } - state = MagicMock() - state.is_replaying.return_value = True - result = extract_durable_function_tags(event, state) + result = extract_durable_function_tags(event) self.assertEqual( result, { "durable_function_execution_name": "my-execution", "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", - "durable_function_first_invocation": "false", + "durable_function_first_invocation": "true", }, ) - def test_sets_first_invocation_true_when_not_replaying(self): + def test_sets_first_invocation_true_when_no_operations(self): + # Empty operations list → not replaying → first invocation event = { "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + "CheckpointToken": "some-token", + "InitialExecutionState": {"Operations": []}, } - state = MagicMock() - state.is_replaying.return_value = False - result = extract_durable_function_tags(event, state) + result = extract_durable_function_tags(event) self.assertEqual( result, { @@ -80,13 +79,19 @@ def test_sets_first_invocation_true_when_not_replaying(self): }, ) - def test_sets_first_invocation_false_when_replaying(self): + def test_sets_first_invocation_false_when_multiple_operations(self): + # More than one operation → replaying → not first invocation event = { "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + "CheckpointToken": "some-token", + "InitialExecutionState": { + "Operations": [ + {"OperationType": "EXECUTION"}, + {"OperationType": "STEP"}, + ] + }, } - state = MagicMock() - state.is_replaying.return_value = True - result = extract_durable_function_tags(event, state) + result = extract_durable_function_tags(event) self.assertEqual( result, { @@ -96,6 +101,22 @@ def test_sets_first_invocation_false_when_replaying(self): }, ) + def test_sets_first_invocation_true_when_initial_execution_state_absent(self): + # No InitialExecutionState key → treated as empty → first invocation + event = { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", + "CheckpointToken": "some-token", + } + result = extract_durable_function_tags(event) + self.assertEqual( + result, + { + "durable_function_execution_name": "my-execution", + "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", + "durable_function_first_invocation": "true", + }, + ) + def test_returns_empty_dict_for_regular_lambda_event(self): event = { "body": '{"key": "value"}', From 7fee216078f870a449e0dda081c8bedc6de1b07d Mon Sep 17 00:00:00 2001 From: Yiming Luo Date: Thu, 19 Mar 2026 14:59:07 +0100 Subject: [PATCH 3/5] Add comment --- datadog_lambda/durable.py | 2 ++ tests/test_durable.py | 33 --------------------------------- 2 files changed, 2 insertions(+), 33 deletions(-) diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index c900f1de..3ebc725c 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -43,6 +43,8 @@ def extract_durable_function_tags(event): return {} execution_name, execution_id = parsed + # Use the number of operations to determine if it's the first invocation. This is + # what the durable execution SDK does to determine the replay status. operations = event.get("InitialExecutionState", {}).get("Operations", []) is_first_invocation = len(operations) <= 1 return { diff --git a/tests/test_durable.py b/tests/test_durable.py index dda36cc9..2cef9050 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -62,23 +62,6 @@ def test_sets_first_invocation_true_when_only_execution_operation(self): }, ) - def test_sets_first_invocation_true_when_no_operations(self): - # Empty operations list → not replaying → first invocation - event = { - "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", - "CheckpointToken": "some-token", - "InitialExecutionState": {"Operations": []}, - } - result = extract_durable_function_tags(event) - self.assertEqual( - result, - { - "durable_function_execution_name": "my-execution", - "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", - "durable_function_first_invocation": "true", - }, - ) - def test_sets_first_invocation_false_when_multiple_operations(self): # More than one operation → replaying → not first invocation event = { @@ -101,22 +84,6 @@ def test_sets_first_invocation_false_when_multiple_operations(self): }, ) - def test_sets_first_invocation_true_when_initial_execution_state_absent(self): - # No InitialExecutionState key → treated as empty → first invocation - event = { - "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", - "CheckpointToken": "some-token", - } - result = extract_durable_function_tags(event) - self.assertEqual( - result, - { - "durable_function_execution_name": "my-execution", - "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004", - "durable_function_first_invocation": "true", - }, - ) - def test_returns_empty_dict_for_regular_lambda_event(self): event = { "body": '{"key": "value"}', From 20c28a4a707dc746b107c84ebe15239b1f905466 Mon Sep 17 00:00:00 2001 From: Yiming Luo Date: Thu, 19 Mar 2026 15:26:24 +0100 Subject: [PATCH 4/5] Change <= 1 to == 1 --- datadog_lambda/durable.py | 4 ++-- tests/test_durable.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index 3ebc725c..f725602d 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -41,12 +41,12 @@ def extract_durable_function_tags(event): if not parsed: logger.error("Failed to parse DurableExecutionArn: %s", durable_execution_arn) return {} - + print(f"event: {event}") execution_name, execution_id = parsed # Use the number of operations to determine if it's the first invocation. This is # what the durable execution SDK does to determine the replay status. operations = event.get("InitialExecutionState", {}).get("Operations", []) - is_first_invocation = len(operations) <= 1 + is_first_invocation = len(operations) == 1 return { "durable_function_execution_name": execution_name, "durable_function_execution_id": execution_id, diff --git a/tests/test_durable.py b/tests/test_durable.py index 2cef9050..53fee61d 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -46,7 +46,7 @@ def test_works_with_numeric_version_qualifier(self): class TestExtractDurableFunctionTags(unittest.TestCase): def test_sets_first_invocation_true_when_only_execution_operation(self): - # One operation (the EXECUTION op itself) → not replaying → first invocation + # One operation → not replaying → first invocation event = { "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", "CheckpointToken": "some-token", From d184ba7d9f2b17bfa8efb6fc16a6ec4105a96681 Mon Sep 17 00:00:00 2001 From: Yiming Luo Date: Thu, 19 Mar 2026 15:34:04 +0100 Subject: [PATCH 5/5] Update comments --- datadog_lambda/durable.py | 1 - tests/test_durable.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index f725602d..3acc7c0c 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -41,7 +41,6 @@ def extract_durable_function_tags(event): if not parsed: logger.error("Failed to parse DurableExecutionArn: %s", durable_execution_arn) return {} - print(f"event: {event}") execution_name, execution_id = parsed # Use the number of operations to determine if it's the first invocation. This is # what the durable execution SDK does to determine the replay status. diff --git a/tests/test_durable.py b/tests/test_durable.py index 53fee61d..0887e2d9 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -46,7 +46,7 @@ def test_works_with_numeric_version_qualifier(self): class TestExtractDurableFunctionTags(unittest.TestCase): def test_sets_first_invocation_true_when_only_execution_operation(self): - # One operation → not replaying → first invocation + # One operation (the current EXECUTION operation itself) → not replaying → first invocation event = { "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004", "CheckpointToken": "some-token",