Skip to content

[POC] Add DD_TRACE_EXTRACTOR env var to customize Trace extraction customization #112

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

Merged
merged 6 commits into from
Feb 19, 2021
Merged

[POC] Add DD_TRACE_EXTRACTOR env var to customize Trace extraction customization #112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
41d49ef
add extractor and tests
Czechh Jan 15, 2021
c7d34ea
lint and other fixes
Czechh Jan 15, 2021
342cc21
fix lint
Czechh Jan 15, 2021
075bbb9
Merge branch 'main' into sergio.prada/trace-extractor
tianchu Feb 17, 2021
15a1132
Update tests
tianchu Feb 17, 2021
b0fd4a0
Fix format
tianchu Feb 17, 2021
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
21 changes: 19 additions & 2 deletions datadog_lambda/tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,20 @@ def extract_context_from_sqs_event_or_context(event, lambda_context):
return extract_context_from_lambda_context(lambda_context)


def extract_dd_trace_context(event, lambda_context):
def extract_context_custom_extractor(extractor, event, lambda_context):
"""
Extract Datadog trace context using a custom trace extractor function
"""
try:
(trace_id, parent_id, sampling_priority,) = extractor(event, lambda_context)
return trace_id, parent_id, sampling_priority
except Exception as e:
logger.debug("The trace extractor returned with error %s", e)

return None, None, None


def extract_dd_trace_context(event, lambda_context, extractor=None):
"""
Extract Datadog trace context from the Lambda `event` object.

Expand All @@ -175,7 +188,11 @@ def extract_dd_trace_context(event, lambda_context):
global dd_trace_context
trace_context_source = None

if "headers" in event:
if extractor is not None:
(trace_id, parent_id, sampling_priority,) = extract_context_custom_extractor(
extractor, event, lambda_context
)
elif "headers" in event:
(
trace_id,
parent_id,
Expand Down
16 changes: 15 additions & 1 deletion datadog_lambda/wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import os
import logging
import traceback
from importlib import import_module

from datadog_lambda.extension import should_use_extension, flush_extension
from datadog_lambda.cold_start import set_cold_start, is_cold_start
Expand All @@ -15,6 +16,7 @@
submit_invocations_metric,
submit_errors_metric,
)
from datadog_lambda.module_name import modify_module_name
from datadog_lambda.patch import patch_all
from datadog_lambda.tracing import (
extract_dd_trace_context,
Expand Down Expand Up @@ -91,6 +93,16 @@ def __init__(self, func):
os.environ.get("DD_MERGE_XRAY_TRACES", "false").lower() == "true"
)
self.function_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", "function")
self.extractor_env = os.environ.get("DD_TRACE_EXTRACTOR", None)
self.trace_extractor = None

if self.extractor_env:
Copy link
Contributor

Choose a reason for hiding this comment

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

If someone isn't using the redirected handler, could this cause a cyclical dependency. Eg. user does this:

from datadog_lambda.wrapper import datadog_lambda_wrapper

@datadog_lambda_wrapper
def my_lambda_handler(event, context):
   pass

def my_lambda_extractor(event, context):
   //... some extractor logic here

In this case, the import chain would be my_function.py -> datadog_lambda.wrapper -> my_function.py -> etc etc .

Copy link
Contributor

Choose a reason for hiding this comment

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

In node we addressed this by only supporting the env var for the extractor in the redirected handler. When you aren't using the redirected handler, it was possible to pass the extractor in as a config variable to the wrapper.

Copy link
Collaborator

Choose a reason for hiding this comment

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

I think I can probably delay the import to when the extractor needs to be invoked. But I'd rather simply instruct users to put the extractor in a separate file (from the one where the handler is defined). I think that's probably good enough for a poc and given the fact that most of applications should be using the redirected handler now.

extractor_parts = self.extractor_env.rsplit(".", 1)
if len(extractor_parts) == 2:
(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)

# Inject trace correlation ids to logs
if self.logs_injection:
Expand Down Expand Up @@ -125,7 +137,9 @@ def _before(self, event, context):
set_cold_start()
submit_invocations_metric(context)
# Extract Datadog trace context and source from incoming requests
dd_context, trace_context_source = extract_dd_trace_context(event, context)
dd_context, trace_context_source = extract_dd_trace_context(
event, context, extractor=self.trace_extractor
)
# Create a Datadog X-Ray subsegment with the trace context
if dd_context and trace_context_source == TraceContextSource.EVENT:
create_dd_dummy_metadata_subsegment(
Expand Down
1 change: 1 addition & 0 deletions tests/integration/serverless-plugin.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ provider:
DD_INTEGRATION_TEST: true
DD_API_KEY: ${env:DD_API_KEY}
WITH_PLUGIN: true
lambdaHashingVersion: 20201221

layers:
python27:
Expand Down
1 change: 1 addition & 0 deletions tests/integration/serverless.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ provider:
environment:
DD_INTEGRATION_TEST: true
DD_API_KEY: ${env:DD_API_KEY}
lambdaHashingVersion: 20201221

layers:
python27:
Expand Down
64 changes: 64 additions & 0 deletions tests/test_tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,70 @@ def test_with_complete_datadog_trace_headers(self):
XraySubsegment.NAMESPACE,
)

def test_with_extractor_function(self):
def extractor_foo(event, context):
foo = event.get("foo", {})
lowercase_foo = {k.lower(): v for k, v in foo.items()}

trace_id = lowercase_foo.get(TraceHeader.TRACE_ID)
parent_id = lowercase_foo.get(TraceHeader.PARENT_ID)
sampling_priority = lowercase_foo.get(TraceHeader.SAMPLING_PRIORITY)
return trace_id, parent_id, sampling_priority

lambda_ctx = get_mock_context()
ctx, ctx_source = extract_dd_trace_context(
{
"foo": {
TraceHeader.TRACE_ID: "123",
TraceHeader.PARENT_ID: "321",
TraceHeader.SAMPLING_PRIORITY: "1",
}
},
lambda_ctx,
extractor=extractor_foo,
)
self.assertEquals(ctx_source, "event")
self.assertDictEqual(
ctx, {"trace-id": "123", "parent-id": "321", "sampling-priority": "1",},
)
self.assertDictEqual(
get_dd_trace_context(),
{
TraceHeader.TRACE_ID: "123",
TraceHeader.PARENT_ID: "65535",
TraceHeader.SAMPLING_PRIORITY: "1",
},
)

def test_graceful_fail_of_extractor_function(self):
def extractor_raiser(event, context):
raise Exception("kreator")

lambda_ctx = get_mock_context()
ctx, ctx_source = extract_dd_trace_context(
{
"foo": {
TraceHeader.TRACE_ID: "123",
TraceHeader.PARENT_ID: "321",
TraceHeader.SAMPLING_PRIORITY: "1",
}
},
lambda_ctx,
extractor=extractor_raiser,
)
self.assertEquals(ctx_source, "xray")
self.assertDictEqual(
ctx, {"trace-id": "4369", "parent-id": "65535", "sampling-priority": "2",},
)
self.assertDictEqual(
get_dd_trace_context(),
{
TraceHeader.TRACE_ID: "4369",
TraceHeader.PARENT_ID: "65535",
TraceHeader.SAMPLING_PRIORITY: "2",
},
)

def test_with_sqs_distributed_datadog_trace_data(self):
lambda_ctx = get_mock_context()
sqs_event = {
Expand Down
2 changes: 1 addition & 1 deletion tests/test_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ def lambda_handler(event, context):
)
self.mock_wrapper_lambda_stats.flush.assert_called()
self.mock_extract_dd_trace_context.assert_called_with(
lambda_event, lambda_context
lambda_event, lambda_context, extractor=None
)
self.mock_set_correlation_ids.assert_called()
self.mock_inject_correlation_ids.assert_called()
Expand Down