Skip to content

Certain xray_recorder operations only work in the LambdaContext #18

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

Merged
merged 1 commit into from
Oct 3, 2019
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
18 changes: 18 additions & 0 deletions datadog_lambda/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import logging

from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core.lambda_launcher import LambdaContext

from ddtrace import patch, tracer
from datadog_lambda.constants import (
Expand Down Expand Up @@ -94,6 +95,10 @@ def get_dd_trace_context():
automatically, but this function can be used to manually inject the trace
context to an outgoing request.
"""
if not is_lambda_context():
logger.debug('get_dd_trace_context is only supported in LambdaContext')
return {}

global dd_trace_context
xray_trace_entity = xray_recorder.get_trace_entity() # xray (sub)segment
if dd_trace_context:
Expand Down Expand Up @@ -124,7 +129,12 @@ def set_correlation_ids():

TODO: Remove me when Datadog tracer is natively supported in Lambda.
"""
if not is_lambda_context():
logger.debug('set_correlation_ids is only supported in LambdaContext')
return

context = get_dd_trace_context()

span = tracer.trace('dummy.span')
span.trace_id = context[TraceHeader.TRACE_ID]
span.span_id = context[TraceHeader.PARENT_ID]
Expand Down Expand Up @@ -154,3 +164,11 @@ def inject_correlation_ids():
patch(logging=True)

logger.debug('logs injection configured')


def is_lambda_context():
"""
Return True if the X-Ray context is `LambdaContext`, rather than the
regular `Context` (e.g., when testing lambda functions locally).
"""
return type(xray_recorder.context) == LambdaContext
Copy link
Contributor

Choose a reason for hiding this comment

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

Hmm, could this throw if X-Ray is not enabled? The technique I've been using to detect lambda is to look for the 'AWS_LAMBDA_FUNCTION_NAME' environment variable.

Copy link
Collaborator Author

@tianchu tianchu Oct 3, 2019

Choose a reason for hiding this comment

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

I did consider checking the env var, but I felt that env var is only a proxy for truth. E.g., developers may have functions reading this env var (that's why it's provided by AWS), in this case, they still need to manually inject values to this env var for local testing as well.

BTW, I got this idea from the xray sdk: https://github.com/aws/aws-xray-sdk-python/blob/e9d3746155729a28766f866714ce082eb6612ad3/aws_xray_sdk/ext/flask/middleware.py#L23

Tested locally and in lambda, disabling AWS_XRAY_SDK_ENABLED does not cause issues.

Copy link
Contributor

Choose a reason for hiding this comment

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

Okay, makes sense.

10 changes: 10 additions & 0 deletions tests/test_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,11 @@ def setUp(self):
self.mock_current_subsegment
self.addCleanup(patcher.stop)

patcher = patch('datadog_lambda.tracing.is_lambda_context')
self.mock_is_lambda_context = patcher.start()
self.mock_is_lambda_context.return_value = True
self.addCleanup(patcher.stop)

def test_without_datadog_trace_headers(self):
extract_dd_trace_context({})
self.assertDictEqual(
Expand Down Expand Up @@ -165,6 +170,11 @@ def setUp(self):
}
self.addCleanup(patcher.stop)

patcher = patch('datadog_lambda.tracing.is_lambda_context')
self.mock_is_lambda_context = patcher.start()
self.mock_is_lambda_context.return_value = True
self.addCleanup(patcher.stop)

def test_set_correlation_ids(self):
set_correlation_ids()
trace_id, span_id = get_correlation_ids()
Expand Down