[core] Add config to set hostname tag on trace root span (#938)

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, )
 

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 = []

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()

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)

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

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'

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.