Skip to content

Support trace and log correlation #15

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
Sep 16, 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
30 changes: 30 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,14 @@ If your Lambda function powers a performance-critical task (e.g., a consumer-fac

- DD_FLUSH_TO_LOG

To connect logs and traces, set the environment variable below to `True`. The default format of the AWS provided `LambdaLoggerHandler` will be overridden to inject `dd.trace_id` and `dd.span_id`. The default Datadog lambda log integration pipeline will automatically parse them and map the `dd.trace_id` into the reserved [trace_id attribute](https://docs.datadoghq.com/logs/processing/#trace-id-attribute).

- DD_LOGS_INJECTION

To debug the Datadog Lambda Layer, set the environment variable below to `DEBUG`.

- DD_LOG_LEVEL

### The Serverless Framework

If your Lambda function is deployed using the Serverless Framework, refer to this sample `serverless.yml`.
Expand Down Expand Up @@ -211,6 +219,28 @@ If your Lambda function is triggered by API Gateway via [the non-proxy integrati

If your Lambda function is deployed by the Serverless Framework, such a mapping template gets created by default.

## Log and Trace Correlations

To connect logs and traces, set the environment variable `DD_LOGS_INJECTION` to `True`. The log format of the AWS provided `LambdaLoggerHandler` will be overridden to inject `dd.trace_id` and `dd.span_id`. The default Datadog lambda log integration pipeline will automatically parse them and map the `dd.trace_id` into the reserved attribute [trace_id](https://docs.datadoghq.com/logs/processing/#trace-id-attribute).

If you use a custom logger handler to log in json, you can manually inject the ids using the helper function `get_correlation_ids`.

```python
from datadog_lambda.wrapper import datadog_lambda_wrapper
from ddtrace.helpers import get_correlation_ids

@datadog_lambda_wrapper
def lambda_handler(event, context):
trace_id, span_id = get_correlation_ids()
logger.info({
"message": "hello world",
"dd": {
"trace_id": trace_id,
"span_id": span_id
}
})
```

## Opening Issues

If you encounter a bug with this package, we want to hear about it. Before opening a new issue, search the existing issues to avoid duplicates.
Expand Down
69 changes: 56 additions & 13 deletions datadog_lambda/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

from aws_xray_sdk.core import xray_recorder

from ddtrace import patch, tracer
from datadog_lambda.constants import (
SamplingPriority,
TraceHeader,
Expand All @@ -18,6 +19,28 @@
dd_trace_context = {}


def _convert_xray_trace_id(xray_trace_id):
"""
Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
"""
return str(0x7FFFFFFFFFFFFFFF & int(xray_trace_id[-16:], 16))


def _convert_xray_entity_id(xray_entity_id):
"""
Convert X-Ray (sub)segement id (hex) to a Datadog span id (int).
"""
return str(int(xray_entity_id, 16))


def _convert_xray_sampling(xray_sampled):
"""
Convert X-Ray sampled (True/False) to its Datadog counterpart.
"""
return str(SamplingPriority.USER_KEEP) if xray_sampled \
else str(SamplingPriority.USER_REJECT)


def extract_dd_trace_context(event):
"""
Extract Datadog trace context from the Lambda `event` object.
Expand Down Expand Up @@ -56,6 +79,8 @@ def extract_dd_trace_context(event):
# reset to avoid using the context from the last invocation.
dd_trace_context = {}

logger.debug('extracted dd trace context %s', dd_trace_context)


def get_dd_trace_context():
"""
Expand Down Expand Up @@ -91,23 +116,41 @@ def get_dd_trace_context():
}


def _convert_xray_trace_id(xray_trace_id):
def set_correlation_ids():
"""
Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
"""
return str(0x7FFFFFFFFFFFFFFF & int(xray_trace_id[-16:], 16))
Create a dummy span, and overrides its trace_id and span_id, to make
ddtrace.helpers.get_correlation_ids() return the correct ids for both
auto and manual log correlations.


def _convert_xray_entity_id(xray_entity_id):
TODO: Remove me when Datadog tracer is natively supported in Lambda.
"""
Convert X-Ray (sub)segement id (hex) to a Datadog span id (int).
"""
return str(int(xray_entity_id, 16))
context = get_dd_trace_context()
span = tracer.trace('dummy.span')
span.trace_id = context[TraceHeader.TRACE_ID]
span.span_id = context[TraceHeader.PARENT_ID]

logger.debug('correlation ids set')

def _convert_xray_sampling(xray_sampled):

def inject_correlation_ids():
"""
Convert X-Ray sampled (True/False) to its Datadog counterpart.
Override the formatter of LambdaLoggerHandler to inject datadog trace and
span id for log correlation.

For manual injections to custom log handlers, use `ddtrace.helpers.get_correlation_ids`
to retrieve correlation ids (trace_id, span_id).
"""
return str(SamplingPriority.USER_KEEP) if xray_sampled \
else str(SamplingPriority.USER_REJECT)
# Override the log format of the AWS provided LambdaLoggerHandler
root_logger = logging.getLogger()
for handler in root_logger.handlers:
if handler.__class__.__name__ == 'LambdaLoggerHandler':
handler.setFormatter(logging.Formatter(
'[%(levelname)s]\t%(asctime)s.%(msecs)dZ\t%(aws_request_id)s\t'
'[dd.trace_id=%(dd.trace_id)s dd.span_id=%(dd.span_id)s]\t%(message)s\n',
'%Y-%m-%dT%H:%M:%S'
))

# Patch `logging.Logger.makeRecord` to actually inject correlation ids
patch(logging=True)

logger.debug('logs injection configured')
18 changes: 15 additions & 3 deletions datadog_lambda/wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,12 @@
import traceback

from datadog_lambda.metric import lambda_stats
from datadog_lambda.tracing import extract_dd_trace_context
from datadog_lambda.patch import patch_all
from datadog_lambda.tracing import (
extract_dd_trace_context,
set_correlation_ids,
inject_correlation_ids,
)

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -37,15 +41,23 @@ class _LambdaDecorator(object):
def __init__(self, func):
self.func = func
self.flush_to_log = os.environ.get('DD_FLUSH_TO_LOG', '').lower() == 'true'
self.logs_injection = os.environ.get('DD_LOGS_INJECTION', '').lower() == 'true'

# Inject trace correlation ids to logs
if self.logs_injection:
inject_correlation_ids()

# Patch HTTP clients to propagate Datadog trace context
patch_all()
logger.debug('datadog_lambda_wrapper initialized')

def _before(self, event, context):
try:
# Extract Datadog trace context from incoming requests
extract_dd_trace_context(event)

# Patch HTTP clients to propogate Datadog trace context
patch_all()
# Set log correlation ids using extracted trace context
set_correlation_ids()
except Exception:
traceback.print_exc()

Expand Down
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
install_requires=[
'aws-xray-sdk==2.4.2',
'datadog==0.28.0',
'ddtrace==0.28.0',
'wrapt==1.11.1',
'setuptools==40.8.0',
'boto3==1.9.160'
Expand Down
20 changes: 20 additions & 0 deletions tests/test_tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
except ImportError:
from mock import MagicMock, patch

from ddtrace.helpers import get_correlation_ids
from datadog_lambda.constants import (
SamplingPriority,
TraceHeader,
Expand All @@ -12,6 +13,7 @@
from datadog_lambda.tracing import (
extract_dd_trace_context,
get_dd_trace_context,
set_correlation_ids,
_convert_xray_trace_id,
_convert_xray_entity_id,
_convert_xray_sampling,
Expand Down Expand Up @@ -150,3 +152,21 @@ def test_convert_xray_sampling(self):
_convert_xray_sampling(False),
str(SamplingPriority.USER_REJECT)
)


class TestLogsInjection(unittest.TestCase):

def setUp(self):
patcher = patch('datadog_lambda.tracing.get_dd_trace_context')
self.mock_get_dd_trace_context = patcher.start()
self.mock_get_dd_trace_context.return_value = {
TraceHeader.TRACE_ID: '123',
TraceHeader.PARENT_ID: '456',
}
self.addCleanup(patcher.stop)

def test_set_correlation_ids(self):
set_correlation_ids()
trace_id, span_id = get_correlation_ids()
self.assertEqual(trace_id, '123')
self.assertEqual(span_id, '456')
28 changes: 26 additions & 2 deletions tests/test_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,14 @@ def setUp(self):
self.mock_extract_dd_trace_context = patcher.start()
self.addCleanup(patcher.stop)

patcher = patch('datadog_lambda.wrapper.set_correlation_ids')
self.mock_set_correlation_ids = patcher.start()
self.addCleanup(patcher.stop)

patcher = patch('datadog_lambda.wrapper.inject_correlation_ids')
self.mock_inject_correlation_ids = patcher.start()
self.addCleanup(patcher.stop)

patcher = patch('datadog_lambda.wrapper.patch_all')
self.mock_patch_all = patcher.start()
self.addCleanup(patcher.stop)
Expand All @@ -42,6 +50,8 @@ def lambda_handler(event, context):
])
self.mock_wrapper_lambda_stats.flush.assert_called()
self.mock_extract_dd_trace_context.assert_called_with(lambda_event)
self.mock_set_correlation_ids.assert_called()
self.mock_inject_correlation_ids.assert_not_called()
self.mock_patch_all.assert_called()

def test_datadog_lambda_wrapper_flush_to_log(self):
Expand All @@ -57,7 +67,21 @@ def lambda_handler(event, context):

self.mock_metric_lambda_stats.distribution.assert_not_called()
self.mock_wrapper_lambda_stats.flush.assert_not_called()
self.mock_extract_dd_trace_context.assert_called_with(lambda_event)
self.mock_patch_all.assert_called()

del os.environ["DD_FLUSH_TO_LOG"]

def test_datadog_lambda_wrapper_inject_correlation_ids(self):
os.environ["DD_LOGS_INJECTION"] = 'True'

@datadog_lambda_wrapper
def lambda_handler(event, context):
lambda_metric("test.metric", 100)

lambda_event = {}
lambda_context = {}
lambda_handler(lambda_event, lambda_context)

self.mock_set_correlation_ids.assert_called()
self.mock_inject_correlation_ids.assert_called()

del os.environ["DD_LOGS_INJECTION"]