Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Redriven Step Functions Trace Merging #545

Merged
merged 8 commits into from
Jan 2, 2025
Merged

Redriven Step Functions Trace Merging #545

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
57bb8a6
use redrive count to generate step functions parent ID
avedmala Dec 4, 2024
bec850f
lint
avedmala Dec 4, 2024
0076581
lint
avedmala Dec 4, 2024
4ec08d3
dont use redriveCount of 0
avedmala Dec 4, 2024
06bffa4
read redriveCount as int
avedmala Dec 16, 2024
9c77aa8
update redriveCount value in tests
avedmala Dec 16, 2024
93fd23f
match test case with logs-backend snapshot
avedmala Dec 17, 2024
b7462df
account for missing redrive count
avedmala Dec 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion datadog_lambda/tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -384,12 +384,23 @@ def _parse_high_64_bits(trace_tags: str) -> str:


def _generate_sfn_parent_id(context: dict) -> int:
"""
The upstream Step Function can propagate its execution context to downstream Lambdas. The
Lambda can use these details to share the same traceID and infer its parent's spanID.

Excluding redriveCount when its 0 to account for cases where customers are using an old
version of the Lambda layer that doesn't use this value for its parentID generation.
"""
execution_id = context.get("Execution").get("Id")
redrive_count = context.get("Execution").get("RedriveCount")
state_name = context.get("State").get("Name")
state_entered_time = context.get("State").get("EnteredTime")

redrive_postfix = "" if redrive_count == 0 else f"#{redrive_count}"
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible that redrive _count is null?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Execution.RedriveCount will always be in the SFNs context object, defaulting to 0

Only Execution.RedriveTime has a risk of being missing, only appearing when redrive count is 1 or more

https://docs.aws.amazon.com/step-functions/latest/dg/input-output-contextobject.html

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Even though this is supposed to be true, I found some instances in customer logs where RedriveCount is null. Wasn't able to reproduce it myself. I'll add some null-check here and also for the node layer change accordingly

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we'll need to assume all the new fields could be null in the future. There may be some edge cases that AWS does not handle well.


return _deterministic_sha256_hash(
f"{execution_id}#{state_name}#{state_entered_time}", HIGHER_64_BITS
f"{execution_id}#{state_name}#{state_entered_time}{redrive_postfix}",
HIGHER_64_BITS,
)


Expand Down
74 changes: 65 additions & 9 deletions tests/test_tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -619,30 +619,84 @@ def test_step_function_trace_data(self):
lambda_ctx = get_mock_context()
sfn_event = {
"Execution": {
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
"Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
"Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
"RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
"StartTime": "2024-12-04T19:38:04.069Z",
"RedriveCount": 0,
},
"StateMachine": {},
"State": {
"Name": "my-awesome-state",
"EnteredTime": "Mon Nov 13 12:43:33 PST 2023",
"Name": "Lambda Invoke",
"EnteredTime": "2024-12-04T19:38:04.118Z",
"RetryCount": 0,
},
"StateMachine": {
"Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
"Name": "abhinav-activity-state-machine",
},
}
ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx)
self.assertEqual(source, "event")
expected_context = Context(
trace_id=3675572987363469717,
span_id=6880978411788117524,
trace_id=435175499815315247,
span_id=3929055471293792800,
sampling_priority=1,
meta={"_dd.p.tid": "3e7a89d1b7310603"},
)
self.assertEqual(ctx, expected_context)
self.assertEqual(
get_dd_trace_context(),
{
TraceHeader.TRACE_ID: "435175499815315247",
TraceHeader.PARENT_ID: "10713633173203262661",
TraceHeader.SAMPLING_PRIORITY: "1",
TraceHeader.TAGS: "_dd.p.tid=3e7a89d1b7310603",
},
)
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
self.mock_send_segment.assert_called_with(
XraySubsegment.TRACE_KEY,
expected_context,
)

# https://github.com/DataDog/logs-backend/blob/c17618cb552fc369ca40282bae0a65803f82f694/domains/serverless/apps/logs-to-traces-reducer/src/test/resources/test-json-files/stepfunctions/RedriveTest/snapshots/RedriveLambdaSuccessTraceMerging.json#L46
@with_trace_propagation_style("datadog")
def test_step_function_trace_data_redrive(self):
lambda_ctx = get_mock_context()
sfn_event = {
"Execution": {
"Id": "arn:aws:states:sa-east-1:425362996713:execution:abhinav-activity-state-machine:72a7ca3e-901c-41bb-b5a3-5f279b92a316",
"Name": "72a7ca3e-901c-41bb-b5a3-5f279b92a316",
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated these cases to match the logs-backend and also the node tests added in DataDog/datadog-lambda-js#598

"RoleArn": "arn:aws:iam::425362996713:role/service-role/StepFunctions-abhinav-activity-state-machine-role-22jpbgl6j",
"StartTime": "2024-12-04T19:38:04.069Z",
"RedriveCount": 1,
},
"State": {
"Name": "Lambda Invoke",
"EnteredTime": "2024-12-04T19:38:04.118Z",
"RetryCount": 0,
},
"StateMachine": {
"Id": "arn:aws:states:sa-east-1:425362996713:stateMachine:abhinav-activity-state-machine",
"Name": "abhinav-activity-state-machine",
},
}
ctx, source, event_source = extract_dd_trace_context(sfn_event, lambda_ctx)
self.assertEqual(source, "event")
expected_context = Context(
trace_id=435175499815315247,
span_id=5063839446130725204,
sampling_priority=1,
meta={"_dd.p.tid": "e987c84b36b11ab"},
meta={"_dd.p.tid": "3e7a89d1b7310603"},
)
self.assertEqual(ctx, expected_context)
self.assertEqual(
get_dd_trace_context(),
{
TraceHeader.TRACE_ID: "3675572987363469717",
TraceHeader.TRACE_ID: "435175499815315247",
TraceHeader.PARENT_ID: "10713633173203262661",
TraceHeader.SAMPLING_PRIORITY: "1",
TraceHeader.TAGS: "_dd.p.tid=e987c84b36b11ab",
TraceHeader.TAGS: "_dd.p.tid=3e7a89d1b7310603",
},
)
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
Expand All @@ -658,6 +712,7 @@ def test_step_function_trace_data_lambda_root(self):
"_datadog": {
"Execution": {
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
"RedriveCount": 0,
},
"StateMachine": {},
"State": {
Expand Down Expand Up @@ -700,6 +755,7 @@ def test_step_function_trace_data_sfn_root(self):
"_datadog": {
"Execution": {
"Id": "665c417c-1237-4742-aaca-8b3becbb9e75",
"RedriveCount": 0,
},
"StateMachine": {},
"State": {
Expand Down
Loading