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

Jump to bottom

[SVLS-4647] Fix lambda authorizer trace context injection issue #453

Merged
merged 1 commit into from
Mar 27, 2024
Merged

[SVLS-4647] Fix lambda authorizer trace context injection issue #453

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
10 changes: 2 additions & 8 deletions datadog_lambda/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
TraceContextSource,
XraySubsegment,
Headers,
TraceHeader,
)
from datadog_lambda.metric import (
flush_stats,
Expand All @@ -44,6 +43,7 @@
InferredSpanInfo,
is_authorizer_response,
tracer,
propagator,
)
from datadog_lambda.trigger import (
extract_trigger_tags,
Expand Down Expand Up @@ -254,13 +254,7 @@ def _inject_authorizer_span_headers(self, request_id):
injected_headers = {}
source_span = self.inferred_span if self.inferred_span else self.span
span_context = source_span.context
injected_headers[TraceHeader.TRACE_ID] = str(span_context.trace_id)
injected_headers[TraceHeader.PARENT_ID] = str(span_context.span_id)
sampling_priority = span_context.sampling_priority
if sampling_priority is not None:
injected_headers[TraceHeader.SAMPLING_PRIORITY] = str(
span_context.sampling_priority
)
propagator.inject(span_context, injected_headers)
Copy link
Contributor

Choose a reason for hiding this comment

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

Great catch, I wonder if this methods should eventually be refactored into a tracing service and not the wrapper 😕

Copy link
Contributor

Choose a reason for hiding this comment

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

the propagator.inject ends up injecting into injected_headers, correct?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Good question!

Having the same thought, I went ahead trying to refactor it.
Then I found it's already in a helper method called _inject_authorizer_span_headers.
Then it felt that the injection is more relevant to the wrapper because the wrapper now has the inferred_span or span data and it also should be responsible for the final response data. And then at this point, the injection became more like prepare_response and the only tracer related part is already encapsulated by the propagator.inject helper.
Hope this makes sense.

injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns
if request_id is not None:
injected_headers[Headers.Authorizing_Request_Id] = request_id
Expand Down
19 changes: 9 additions & 10 deletions tests/test_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from datadog_lambda.metric import lambda_metric
from datadog_lambda.thread_stats_writer import ThreadStatsWriter
from ddtrace import Span, tracer
from ddtrace.internal.constants import MAX_UINT_64BITS


def get_mock_context(
Expand Down Expand Up @@ -543,19 +544,19 @@ def lambda_handler(event, context):
lambda_event = {}

lambda_context = get_mock_context()
mock_span = Span(name="my_inferred_span", span_id=123, trace_id=456)
mock_span.context.sampling_priority = "1"
mock_span.context.dd_origin = None
mock_span.start_ns = 1668127541671386817
mock_span.duration_ns = 1e8
lambda_handler.inferred_span = mock_span
test_span = tracer.trace("test_span")
trace_ctx = tracer.current_trace_context()
test_span.finish()
lambda_handler.inferred_span = test_span
lambda_handler.make_inferred_span = False
result = lambda_handler(lambda_event, lambda_context)
raw_inject_data = result["context"]["_datadog"]
self.assertIsInstance(raw_inject_data, str)
inject_data = json.loads(base64.b64decode(raw_inject_data))
self.assertEqual(inject_data[TraceHeader.PARENT_ID], "123")
self.assertEqual(inject_data[TraceHeader.TRACE_ID], "456")
self.assertEqual(inject_data[TraceHeader.PARENT_ID], str(trace_ctx.span_id))
self.assertEqual(
inject_data[TraceHeader.TRACE_ID], str(MAX_UINT_64BITS & trace_ctx.trace_id)
)
self.assertEqual(inject_data[TraceHeader.SAMPLING_PRIORITY], "1")
self.assertEqual(result["context"]["scope"], "still here")

Expand Down Expand Up @@ -662,7 +663,6 @@ def tearDown(self):

@patch("datadog_lambda.wrapper.should_use_extension", True)
def test_local_test_envvar_flushing(self):

flushes = []
lambda_event = {}
lambda_context = get_mock_context()
Expand All @@ -680,7 +680,6 @@ def flush():
({"DD_LOCAL_TEST": ""}, False),
({}, False),
):

os.environ = environ
flushes.clear()

Expand Down