-
Notifications
You must be signed in to change notification settings - Fork 51
durable: extract trace context from checkpoints and input payload #818
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
20b0054
604a2a1
4db58ed
2aedd77
499aa45
f11ab2a
b3f83df
62d8b69
42247d0
958daf0
bb9e643
9c4e70c
806de56
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,2 +1,3 @@ | ||
| [flake8] | ||
| max-line-length = 100 | ||
| max-line-length = 100 | ||
| extend-ignore = E203,E701 |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -31,7 +31,7 @@ | |
|
|
||
| from ddtrace import patch | ||
| from ddtrace import __version__ as ddtrace_version | ||
| from ddtrace.propagation.http import HTTPPropagator | ||
| from ddtrace.propagation.http import HTTPPropagator, _DatadogMultiHeader | ||
| from ddtrace.trace import Context, Span, tracer | ||
|
|
||
| from datadog_lambda.config import config | ||
|
|
@@ -61,6 +61,7 @@ | |
| DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000" | ||
| HIGHER_64_BITS = "HIGHER_64_BITS" | ||
| LOWER_64_BITS = "LOWER_64_BITS" | ||
| _TRACE_CHECKPOINT_PREFIX = "_datadog_" | ||
|
|
||
|
|
||
| def _dsm_set_checkpoint(context_json, event_type, arn): | ||
|
|
@@ -546,6 +547,56 @@ def extract_context_from_step_functions(event, lambda_context): | |
| return extract_context_from_lambda_context(lambda_context) | ||
|
|
||
|
|
||
| def _extract_context_from_durable_checkpoint(operation): | ||
| if not isinstance(operation, dict): | ||
| return None | ||
|
|
||
| step_details = operation.get("StepDetails") | ||
| if not isinstance(step_details, dict): | ||
| return None | ||
|
|
||
| result = step_details.get("Result") | ||
| if isinstance(result, str): | ||
| try: | ||
| result = json.loads(result) | ||
| except Exception: | ||
| return None | ||
|
|
||
| if not isinstance(result, dict): | ||
| return None | ||
|
|
||
| # Checkpoints are written by dd-trace-py as x-datadog-* headers, so extract | ||
| # directly and bypass DD_TRACE_PROPAGATION_STYLE_EXTRACT on purpose. | ||
| return _DatadogMultiHeader._extract(result) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can this cause issues if tomorrow ddtrace refactors this class? |
||
|
|
||
|
|
||
| def extract_context_from_durable_execution(event): | ||
| operations = event.get("InitialExecutionState", {}).get("Operations") | ||
| if isinstance(operations, dict): | ||
| operations = list(operations.values()) | ||
| if not isinstance(operations, list) or not operations: | ||
| return None | ||
|
|
||
| highest = -1 | ||
| best_operation = None | ||
| for operation in operations: | ||
| if not isinstance(operation, dict): | ||
| continue | ||
| name = operation.get("Name") | ||
| if not isinstance(name, str) or not name.startswith(_TRACE_CHECKPOINT_PREFIX): | ||
| continue | ||
| suffix = name[len(_TRACE_CHECKPOINT_PREFIX) :] | ||
| try: | ||
| number = int(suffix) | ||
| except (TypeError, ValueError): | ||
| continue | ||
| if number > highest: | ||
| highest = number | ||
| best_operation = operation | ||
|
|
||
| return _extract_context_from_durable_checkpoint(best_operation) | ||
|
|
||
|
|
||
| def extract_context_custom_extractor(extractor, event, lambda_context): | ||
| """ | ||
| Extract Datadog trace context using a custom trace extractor function | ||
|
|
@@ -633,9 +684,12 @@ def extract_dd_trace_context( | |
| global dd_trace_context | ||
| trace_context_source = None | ||
| event_source = parse_event_source(event) | ||
| context = None | ||
|
|
||
| if extractor is not None: | ||
| context = extract_context_custom_extractor(extractor, event, lambda_context) | ||
| elif isinstance(event, dict) and "DurableExecutionArn" in event: | ||
| context = extract_context_from_durable_execution(event) | ||
|
Comment on lines
+691
to
+692
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. What happens in the first invocation here? Do we expect the first invocation to extract the context with the other methods? |
||
| elif isinstance(event, (set, dict)) and "request" in event: | ||
| context = extract_context_from_request_header_or_context( | ||
| event, lambda_context, event_source | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: should these functions live in datadog_lambda/durable.py?