Propagate x-datadog-origin (#821)

* Propagate x-datadog-origin

* Add tests

* More tests

* Apply suggestions from code review (Brett)

Co-Authored-By: btoueg <[email protected]>

* Don't forget to tag

* Fix test

* Rename tag to _dd.origin

* Make use of constants

* Use new extract_header_value helper

Author: btoueg

ddtrace/constants.py

@@ -1,6 +1,7 @@
 FILTERS_KEY = 'FILTERS'
 SAMPLE_RATE_METRIC_KEY = '_sample_rate'
 SAMPLING_PRIORITY_KEY = '_sampling_priority_v1'
+ORIGIN_KEY = '_dd.origin'
 EVENT_SAMPLE_RATE_KEY = '_dd1.sr.eausr'
 
 NUMERIC_TAGS = (EVENT_SAMPLE_RATE_KEY, )

ddtrace/context.py

@@ -1,7 +1,7 @@
 import logging
 import threading
 
-from .constants import SAMPLING_PRIORITY_KEY
+from .constants import SAMPLING_PRIORITY_KEY, ORIGIN_KEY
 from .utils.formats import asbool, get_env
 
 log = logging.getLogger(__name__)
@@ -25,7 +25,7 @@ class Context(object):
     _partial_flush_enabled = asbool(get_env('tracer', 'partial_flush_enabled', 'false'))
     _partial_flush_min_spans = int(get_env('tracer', 'partial_flush_min_spans', 500))
 
-    def __init__(self, trace_id=None, span_id=None, sampled=True, sampling_priority=None):
+    def __init__(self, trace_id=None, span_id=None, sampled=True, sampling_priority=None, _dd_origin=None):
         """
         Initialize a new thread-safe ``Context``.
 
@@ -41,6 +41,7 @@ def __init__(self, trace_id=None, span_id=None, sampled=True, sampling_priority=
         self._parent_span_id = span_id
         self._sampled = sampled
         self._sampling_priority = sampling_priority
+        self._dd_origin = _dd_origin
 
     @property
     def trace_id(self):
@@ -184,6 +185,10 @@ def get(self):
                 # attach the sampling priority to the context root span
                 if sampled and sampling_priority is not None and trace:
                     trace[0].set_metric(SAMPLING_PRIORITY_KEY, sampling_priority)
+                origin = self._dd_origin
+                # attach the origin to the root span tag
+                if sampled and origin is not None and trace:
+                    trace[0].set_tag(ORIGIN_KEY, origin)
 
                 # clean the current state
                 self._trace = []
@@ -202,6 +207,10 @@ def get(self):
                 # attach the sampling priority to the context root span
                 if sampled and sampling_priority is not None and trace:
                     trace[0].set_metric(SAMPLING_PRIORITY_KEY, sampling_priority)
+                origin = self._dd_origin
+                # attach the origin to the root span tag
+                if sampled and origin is not None and trace:
+                    trace[0].set_tag(ORIGIN_KEY, origin)
 
                 # Any open spans will remain as `self._trace`
                 # Any finished spans will get returned to be flushed

ddtrace/propagation/http.py

@@ -11,6 +11,7 @@
 HTTP_HEADER_TRACE_ID = "x-datadog-trace-id"
 HTTP_HEADER_PARENT_ID = "x-datadog-parent-id"
 HTTP_HEADER_SAMPLING_PRIORITY = "x-datadog-sampling-priority"
+HTTP_HEADER_ORIGIN = "x-datadog-origin"
 
 
 # Note that due to WSGI spec we have to also check for uppercased and prefixed
@@ -24,6 +25,9 @@
 POSSIBLE_HTTP_HEADER_SAMPLING_PRIORITIES = frozenset(
     [HTTP_HEADER_SAMPLING_PRIORITY, get_wsgi_header(HTTP_HEADER_SAMPLING_PRIORITY)]
 )
+POSSIBLE_HTTP_HEADER_ORIGIN = frozenset(
+    [HTTP_HEADER_ORIGIN, get_wsgi_header(HTTP_HEADER_ORIGIN)]
+)
 
 
 class HTTPPropagator(object):
@@ -54,6 +58,9 @@ def parent_call():
         # Propagate priority only if defined
         if sampling_priority is not None:
             headers[HTTP_HEADER_SAMPLING_PRIORITY] = str(span_context.sampling_priority)
+        # Propagate origin only if defined
+        if span_context._dd_origin is not None:
+            headers[HTTP_HEADER_ORIGIN] = str(span_context._dd_origin)
 
     @staticmethod
     def extract_header_value(possible_header_names, headers, default=None):
@@ -86,6 +93,12 @@ def extract_sampling_priority(headers):
             POSSIBLE_HTTP_HEADER_SAMPLING_PRIORITIES, headers,
         )
 
+    @staticmethod
+    def extract_origin(headers):
+        return HTTPPropagator.extract_header_value(
+            POSSIBLE_HTTP_HEADER_ORIGIN, headers,
+        )
+
     def extract(self, headers):
         """Extract a Context from HTTP headers into a new Context.
 
@@ -111,6 +124,7 @@ def my_controller(url, headers):
             trace_id = HTTPPropagator.extract_trace_id(headers)
             parent_span_id = HTTPPropagator.extract_parent_span_id(headers)
             sampling_priority = HTTPPropagator.extract_sampling_priority(headers)
+            origin = HTTPPropagator.extract_origin(headers)
 
             if sampling_priority is not None:
                 sampling_priority = int(sampling_priority)
@@ -119,15 +133,17 @@ def my_controller(url, headers):
                 trace_id=trace_id,
                 span_id=parent_span_id,
                 sampling_priority=sampling_priority,
+                _dd_origin=origin,
             )
         # If headers are invalid and cannot be parsed, return a new context and log the issue.
         except Exception as error:
             try:
                 log.debug(
-                    "invalid x-datadog-* headers, trace-id: %s, parent-id: %s, priority: %s, error: %s",
+                    "invalid x-datadog-* headers, trace-id: %s, parent-id: %s, priority: %s, origin: %s, error: %s",
                     headers.get(HTTP_HEADER_TRACE_ID, 0),
                     headers.get(HTTP_HEADER_PARENT_ID, 0),
                     headers.get(HTTP_HEADER_SAMPLING_PRIORITY),
+                    headers.get(HTTP_HEADER_ORIGIN, ''),
                     error,
                 )
             # We might fail on string formatting errors ; in that case only format the first error

tests/contrib/pyramid/test_pyramid.py

@@ -1,5 +1,7 @@
 from nose.tools import eq_, ok_
 
+from ddtrace.constants import SAMPLING_PRIORITY_KEY, ORIGIN_KEY
+
 from .utils import PyramidTestCase, PyramidBase
 
 
@@ -32,6 +34,7 @@ def test_distributed_tracing(self):
             'x-datadog-trace-id': '100',
             'x-datadog-parent-id': '42',
             'x-datadog-sampling-priority': '2',
+            'x-datadog-origin': 'synthetics',
         }
         self.app.get('/', headers=headers, status=200)
         writer = self.tracer.writer
@@ -41,7 +44,8 @@ def test_distributed_tracing(self):
         span = spans[0]
         eq_(span.trace_id, 100)
         eq_(span.parent_id, 42)
-        eq_(span.get_metric('_sampling_priority_v1'), 2)
+        eq_(span.get_metric(SAMPLING_PRIORITY_KEY), 2)
+        eq_(span.get_tag(ORIGIN_KEY), 'synthetics')
 
 
 class TestPyramidDistributedTracingDisabled(PyramidBase):
@@ -58,6 +62,7 @@ def test_distributed_tracing_disabled(self):
             'x-datadog-trace-id': '100',
             'x-datadog-parent-id': '42',
             'x-datadog-sampling-priority': '2',
+            'x-datadog-origin': 'synthetics',
         }
         self.app.get('/', headers=headers, status=200)
         writer = self.tracer.writer
@@ -67,4 +72,5 @@ def test_distributed_tracing_disabled(self):
         span = spans[0]
         ok_(span.trace_id != 100)
         ok_(span.parent_id != 42)
-        ok_(span.get_metric('_sampling_priority_v1') != 2)
+        ok_(span.get_metric(SAMPLING_PRIORITY_KEY) != 2)
+        ok_(span.get_tag(ORIGIN_KEY) != 'synthetics')

tests/contrib/pyramid/test_pyramid_autopatch.py

@@ -27,6 +27,7 @@ def test_distributed_tracing(self):
             'x-datadog-trace-id': '100',
             'x-datadog-parent-id': '42',
             'x-datadog-sampling-priority': '2',
+            'x-datadog-origin': 'synthetics',
         }
         self.app.get('/', headers=headers, status=200)
         writer = self.tracer.writer

tests/contrib/tornado/test_tornado_web.py

@@ -3,7 +3,7 @@
 from .web.app import CustomDefaultHandler
 from .utils import TornadoTestCase
 
-from ddtrace.constants import SAMPLING_PRIORITY_KEY, EVENT_SAMPLE_RATE_KEY
+from ddtrace.constants import SAMPLING_PRIORITY_KEY, ORIGIN_KEY, EVENT_SAMPLE_RATE_KEY
 
 from opentracing.scope_managers.tornado import TornadoScopeManager
 from tests.opentracer.utils import init_tracer
@@ -321,7 +321,8 @@ def test_no_propagation(self):
         headers = {
             'x-datadog-trace-id': '1234',
             'x-datadog-parent-id': '4567',
-            'x-datadog-sampling-priority': '2'
+            'x-datadog-sampling-priority': '2',
+            'x-datadog-origin': 'synthetics',
         }
         response = self.fetch('/success/', headers=headers)
         eq_(200, response.code)
@@ -342,6 +343,7 @@ def test_no_propagation(self):
         assert request_span.trace_id != 1234
         assert request_span.parent_id != 4567
         assert request_span.get_metric(SAMPLING_PRIORITY_KEY) != 2
+        assert request_span.get_tag(ORIGIN_KEY) != 'synthetics'
 
 
 class TestCustomTornadoWeb(TornadoTestCase):

tests/propagation/test_http.py

@@ -7,6 +7,7 @@
     HTTP_HEADER_TRACE_ID,
     HTTP_HEADER_PARENT_ID,
     HTTP_HEADER_SAMPLING_PRIORITY,
+    HTTP_HEADER_ORIGIN,
 )
 
 
@@ -21,6 +22,7 @@ def test_inject(self):
 
         with tracer.trace("global_root_span") as span:
             span.context.sampling_priority = 2
+            span.context._dd_origin = "synthetics"
             headers = {}
             propagator = HTTPPropagator()
             propagator.inject(span.context, headers)
@@ -31,6 +33,10 @@ def test_inject(self):
                 int(headers[HTTP_HEADER_SAMPLING_PRIORITY]),
                 span.context.sampling_priority,
             )
+            eq_(
+                headers[HTTP_HEADER_ORIGIN],
+                span.context._dd_origin,
+            )
 
     def test_extract(self):
         tracer = get_dummy_tracer()
@@ -39,6 +45,7 @@ def test_extract(self):
             "x-datadog-trace-id": "1234",
             "x-datadog-parent-id": "5678",
             "x-datadog-sampling-priority": "1",
+            "x-datadog-origin": "synthetics",
         }
 
         propagator = HTTPPropagator()
@@ -49,6 +56,7 @@ def test_extract(self):
             eq_(span.trace_id, 1234)
             eq_(span.parent_id, 5678)
             eq_(span.context.sampling_priority, 1)
+            eq_(span.context._dd_origin, "synthetics")
 
     def test_WSGI_extract(self):
         """Ensure we support the WSGI formatted headers as well."""
@@ -58,6 +66,7 @@ def test_WSGI_extract(self):
             "HTTP_X_DATADOG_TRACE_ID": "1234",
             "HTTP_X_DATADOG_PARENT_ID": "5678",
             "HTTP_X_DATADOG_SAMPLING_PRIORITY": "1",
+            "HTTP_X_DATADOG_ORIGIN": "synthetics",
         }
 
         propagator = HTTPPropagator()
@@ -68,3 +77,4 @@ def test_WSGI_extract(self):
             eq_(span.trace_id, 1234)
             eq_(span.parent_id, 5678)
             eq_(span.context.sampling_priority, 1)
+            eq_(span.context._dd_origin, "synthetics")

tests/test_context.py

@@ -388,6 +388,7 @@ def test_clone(self):
         eq_(cloned_ctx._parent_span_id, ctx._parent_span_id)
         eq_(cloned_ctx._sampled, ctx._sampled)
         eq_(cloned_ctx._sampling_priority, ctx._sampling_priority)
+        eq_(cloned_ctx._dd_origin, ctx._dd_origin)
         eq_(cloned_ctx._current_span, ctx._current_span)
         eq_(cloned_ctx._trace, [])
         eq_(cloned_ctx._finished_spans, 0)