feat: lambda support for DSM #622
Conversation
![datadog-datadog-prod-us1[bot]](https://avatars.githubusercontent.com/u/365230?s=60&v=4){kind=small}
datadog_lambda/tracing.py
Outdated
| if config.data_streams_enabled: | ||
| from ddtrace.data_streams import PROPAGATION_KEY_BASE_64 | ||
|
|
||
| data_streams_ctx = { |
There was a problem hiding this comment.
I know the else is redundant but datadog gets mad if i just do the if too many indents
datadog_lambda/tracing.py
Outdated
| except Exception as e: | ||
| logger.debug("The trace extractor returned with error %s", e) | ||
| return extract_context_from_lambda_context(lambda_context) | ||
| return extract_context_from_lambda_context(lambda_context), None |
datadog_lambda/tracing.py
Outdated
| @@ -265,15 +265,27 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): | |||
| if dd_json_data: | |||
| dd_data = json.loads(dd_json_data) | |||
|
|
|||
| data_streams_ctx = {} | |||
| if config.data_streams_enabled: | |||
| from ddtrace.data_streams import PROPAGATION_KEY_BASE_64 | |||
There was a problem hiding this comment.
My main concerns are
- Creating dictionary objects and bound methods for every invocation is inefficient.
- It is very hard to follow the logic and hard to maintain and may introduce unexpected behaviors that are hard to debug in the future.
May I suggest the following alternative implementation? Let me know what do you think.
def _create_dsm_carrier_func(dd_data):
"""Create a carrier function for DSM context extraction."""
def carrier_get(key):
return dd_data.get(key) if dd_data else None
return carrier_get
# then in In the extraction functions:
if config.data_streams_enabled:
dsm_carrier = _create_dsm_carrier_func(dd_data) # Pass the original dd_data
else:
dsm_carrier = NoneThere was a problem hiding this comment.
I agree with the justifications you made for this change will change the code now!
michael-zhao459
force-pushed
the
michael.zhao/dsm-lambda
branch
from
66fdbb3 to
c4fa49b
Compare
pyproject.toml
Outdated
| @@ -28,7 +28,7 @@ classifiers = [ | |||
| python = ">=3.8.0,<4" | |||
| datadog = ">=0.51.0,<1.0.0" | |||
| wrapt = "^1.11.2" | |||
| ddtrace = ">=2.20.0,<4" | |||
| ddtrace = ">=3.10.0" | |||
There was a problem hiding this comment.
this changes the major version. Is that what we want to do? Also, should we keep <4?
There was a problem hiding this comment.
My mistake on the <4. The code will break without ddtrace version 3.10.0.
tests/test_tracing.py
Outdated
| ) | ||
|
|
||
| @patch("datadog_lambda.tracing._dsm_set_checkpoint") | ||
| def test_sqs_incorrect_datadog_message_attribute(self, mock_dsm_set_checkpoint): |
tests/test_tracing.py
Outdated
|
|
||
| @patch("datadog_lambda.tracing._dsm_set_checkpoint") | ||
| @patch("datadog_lambda.tracing.logger") | ||
| def test_sqs_invalid_datadog_message_attribute_raises_exception( |
There was a problem hiding this comment.
nit: remove raises_exception from the test name. We care about what the test case tests, not the logic under the hood. Maybe in the future, the code won't raise an exception, and it's still OK because the function accepts invalid datadog message attributes.
tests/test_tracing.py
Outdated
| event, self.lambda_context, parse_event_source(event) | ||
| ) | ||
|
|
||
| mock_dsm_set_checkpoint.assert_called_once_with(None, "sqs", "") |
There was a problem hiding this comment.
I don't understand, are we testing mock_dsm_set_checkpoint, or self.mock_checkpoint. We should test only one of them, and it should be consistent across all tests. If possible, it's better to test the lower level one (self.mock_checkpoint)
What does this PR do?
This PR adds lambda support for Data Streams Monitoring (DSM) and reworks the original implementation.
DSM context is passed through the trace propagation headers, code is refactored to use existing extraction logic (deleted dsm.py, reinventing the wheel here).
If DSM is enabled, add custom DSM logic to the extracted context afterwards
Motivation
Remove redundant code. DSM customers wanted to have Lambda support, currently context is not propagated correctly with lambdas.
Testing Guidelines
The tests go through all of the SQS case first, then all of the SNS case, then all of the SNS->SQS case, then all of the Kinesis case
Additional Notes
Types of Changes
Check all that apply