From 21d207efe8c8233b6ff6499c5bd31a73c3b287ba Mon Sep 17 00:00:00 2001 From: Brett Langdon Date: Wed, 15 May 2019 08:47:48 -0400 Subject: [PATCH] [core] Add config to set hostname tag on trace root span (#938) --- ddtrace/constants.py | 1 + ddtrace/context.py | 14 +++++++- ddtrace/internal/hostname.py | 20 +++++++++++ ddtrace/settings/config.py | 4 +++ tests/base/__init__.py | 3 ++ tests/internal/test_hostname.py | 14 ++++++++ tests/test_context.py | 63 +++++++++++++++++++++++++++++++-- 7 files changed, 115 insertions(+), 4 deletions(-) create mode 100644 ddtrace/internal/hostname.py create mode 100644 tests/internal/test_hostname.py diff --git a/ddtrace/constants.py b/ddtrace/constants.py index 572d635fb6a..dec65da7072 100644 --- a/ddtrace/constants.py +++ b/ddtrace/constants.py @@ -3,6 +3,7 @@ SAMPLING_PRIORITY_KEY = '_sampling_priority_v1' ANALYTICS_SAMPLE_RATE_KEY = '_dd1.sr.eausr' ORIGIN_KEY = '_dd.origin' +HOSTNAME_KEY = '_dd.hostname' NUMERIC_TAGS = (ANALYTICS_SAMPLE_RATE_KEY, ) diff --git a/ddtrace/context.py b/ddtrace/context.py index 1b65e1ab74e..3a4c29d4e85 100644 --- a/ddtrace/context.py +++ b/ddtrace/context.py @@ -1,7 +1,9 @@ import threading -from .constants import SAMPLING_PRIORITY_KEY, ORIGIN_KEY +from .constants import HOSTNAME_KEY, SAMPLING_PRIORITY_KEY, ORIGIN_KEY from .internal.logger import get_logger +from .internal import hostname +from .settings import config from .utils.formats import asbool, get_env log = get_logger(__name__) @@ -190,6 +192,11 @@ def get(self): if sampled and origin is not None and trace: trace[0].set_tag(ORIGIN_KEY, origin) + # Set hostname tag if they requested it + if config.report_hostname: + # DEV: `get_hostname()` value is cached + trace[0].set_tag(HOSTNAME_KEY, hostname.get_hostname()) + # clean the current state self._trace = [] self._finished_spans = 0 @@ -212,6 +219,11 @@ def get(self): if sampled and origin is not None and trace: trace[0].set_tag(ORIGIN_KEY, origin) + # Set hostname tag if they requested it + if config.report_hostname: + # DEV: `get_hostname()` value is cached + trace[0].set_tag(HOSTNAME_KEY, hostname.get_hostname()) + # Any open spans will remain as `self._trace` # Any finished spans will get returned to be flushed opened_spans = [] diff --git a/ddtrace/internal/hostname.py b/ddtrace/internal/hostname.py new file mode 100644 index 00000000000..f5ce2e97298 --- /dev/null +++ b/ddtrace/internal/hostname.py @@ -0,0 +1,20 @@ +import functools +import socket + +_hostname = None + + +def _cached(func): + @functools.wraps(func) + def wrapper(): + global _hostname + if not _hostname: + _hostname = func() + + return _hostname + return wrapper + + +@_cached +def get_hostname(): + return socket.gethostname() diff --git a/ddtrace/settings/config.py b/ddtrace/settings/config.py index d0cfa7f674a..88ac02ae013 100644 --- a/ddtrace/settings/config.py +++ b/ddtrace/settings/config.py @@ -30,6 +30,10 @@ def __init__(self): get_env('trace', 'analytics_enabled', default=legacy_config_value) ) + self.report_hostname = asbool( + get_env('trace', 'report_hostname', default=False) + ) + def __getattr__(self, name): if name not in self._config: self._config[name] = IntegrationConfig(self, name) diff --git a/tests/base/__init__.py b/tests/base/__init__.py index e8ffbab740d..154b8e655fc 100644 --- a/tests/base/__init__.py +++ b/tests/base/__init__.py @@ -55,12 +55,15 @@ def override_global_config(values): """ # DEV: Uses dict as interface but internally handled as attributes on Config instance analytics_enabled_original = ddtrace.config.analytics_enabled + report_hostname_original = ddtrace.config.report_hostname ddtrace.config.analytics_enabled = values.get('analytics_enabled', analytics_enabled_original) + ddtrace.config.report_hostname = values.get('report_hostname', report_hostname_original) try: yield finally: ddtrace.config.analytics_enabled = analytics_enabled_original + ddtrace.config.report_hostname = report_hostname_original @staticmethod @contextlib.contextmanager diff --git a/tests/internal/test_hostname.py b/tests/internal/test_hostname.py new file mode 100644 index 00000000000..6ef048e1f4e --- /dev/null +++ b/tests/internal/test_hostname.py @@ -0,0 +1,14 @@ +import mock + +from ddtrace.internal.hostname import get_hostname + + +@mock.patch('socket.gethostname') +def test_get_hostname(socket_gethostname): + # Test that `get_hostname()` just returns `socket.gethostname` + socket_gethostname.return_value = 'test-hostname' + assert get_hostname() == 'test-hostname' + + # Change the value returned by `socket.gethostname` to test the cache + socket_gethostname.return_value = 'new-hostname' + assert get_hostname() == 'test-hostname' diff --git a/tests/test_context.py b/tests/test_context.py index 3a2fca0494c..f2505f4ba77 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -2,15 +2,16 @@ import mock import threading -from unittest import TestCase +from .base import BaseTestCase from tests.test_tracer import get_dummy_tracer from ddtrace.span import Span from ddtrace.context import Context, ThreadLocalContext +from ddtrace.constants import HOSTNAME_KEY from ddtrace.ext.priority import USER_REJECT, AUTO_REJECT, AUTO_KEEP, USER_KEEP -class TestTracingContext(TestCase): +class TestTracingContext(BaseTestCase): """ Tests related to the ``Context`` class that hosts the trace for the current execution flow. @@ -115,6 +116,62 @@ def test_get_trace_empty(self): assert trace is None assert sampled is None + @mock.patch('ddtrace.internal.hostname.get_hostname') + def test_get_report_hostname_enabled(self, get_hostname): + get_hostname.return_value = 'test-hostname' + + with self.override_global_config(dict(report_hostname=True)): + # Create a context and add a span and finish it + ctx = Context() + span = Span(tracer=None, name='fake_span') + ctx.add_span(span) + ctx.close_span(span) + + # Assert that we have not added the tag to the span yet + assert span.get_tag(HOSTNAME_KEY) is None + + # Assert that retrieving the trace sets the tag + trace, _ = ctx.get() + assert trace[0].get_tag(HOSTNAME_KEY) == 'test-hostname' + assert span.get_tag(HOSTNAME_KEY) == 'test-hostname' + + @mock.patch('ddtrace.internal.hostname.get_hostname') + def test_get_report_hostname_disabled(self, get_hostname): + get_hostname.return_value = 'test-hostname' + + with self.override_global_config(dict(report_hostname=False)): + # Create a context and add a span and finish it + ctx = Context() + span = Span(tracer=None, name='fake_span') + ctx.add_span(span) + ctx.close_span(span) + + # Assert that we have not added the tag to the span yet + assert span.get_tag(HOSTNAME_KEY) is None + + # Assert that retrieving the trace does not set the tag + trace, _ = ctx.get() + assert trace[0].get_tag(HOSTNAME_KEY) is None + assert span.get_tag(HOSTNAME_KEY) is None + + @mock.patch('ddtrace.internal.hostname.get_hostname') + def test_get_report_hostname_default(self, get_hostname): + get_hostname.return_value = 'test-hostname' + + # Create a context and add a span and finish it + ctx = Context() + span = Span(tracer=None, name='fake_span') + ctx.add_span(span) + ctx.close_span(span) + + # Assert that we have not added the tag to the span yet + assert span.get_tag(HOSTNAME_KEY) is None + + # Assert that retrieving the trace does not set the tag + trace, _ = ctx.get() + assert trace[0].get_tag(HOSTNAME_KEY) is None + assert span.get_tag(HOSTNAME_KEY) is None + def test_partial_flush(self): """ When calling `Context.get` @@ -393,7 +450,7 @@ def test_clone(self): assert cloned_ctx._finished_spans == 0 -class TestThreadContext(TestCase): +class TestThreadContext(BaseTestCase): """ Ensures that a ``ThreadLocalContext`` makes the Context local to each thread.