Tag traces with HTTP headers from DD_TRACE_HEADER_TAGS (#2782) (#2829)

* Tag traces with HTTP headers from DD_TRACE_HEADER_TAGS

Value must be either comma or space separated. e.g. Host:http.host,User-Agent,http.user_agent
or referer:http.referer Content-Type:http.content_type.

* Better documentation

* Fix urllib3 tests

* Fix urllib3 tests

* Update documentation

* Keep the header_is_traced method

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: Kyle Verhoog <[email protected]>
(cherry picked from commit fbb64a3)

Co-authored-by: Nicolas Vivet <[email protected]>

Author: mergify[bot]

ddtrace/contrib/trace_utils.py

@@ -97,10 +97,11 @@ def _store_headers(headers, span, integration_config, request_or_response):
         return
 
     for header_name, header_value in headers.items():
-        if not integration_config.header_is_traced(header_name):
+        tag_name = integration_config._header_tag_name(header_name)
+        if tag_name is None:
             continue
-        tag_name = _normalize_tag_name(request_or_response, header_name)
-        span.set_tag(tag_name, header_value)
+        # An empty tag defaults to a http.<request or response>.headers.<header name> tag
+        span.set_tag(tag_name or _normalize_tag_name(request_or_response, header_name), header_value)
 
 
 def _store_request_headers(headers, span, integration_config):

ddtrace/profiling/collector/stack.pyx

@@ -21,6 +21,7 @@ from ddtrace.profiling.collector import _traceback
 from ddtrace.utils import attr as attr_utils
 from ddtrace.utils import formats
 
+
 try:
     import gevent.thread
 except ImportError:

ddtrace/settings/config.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 import os
 from typing import List
+from typing import Optional
 from typing import Tuple
 
 from ddtrace.utils.cache import cachedmethod
@@ -106,7 +107,10 @@ def is_error_code(self, status_code):
     def __init__(self):
         # use a dict as underlying storing mechanism
         self._config = {}
-        self.http = HttpConfig()
+
+        header_tags = parse_tags_str(get_env("trace", "header_tags") or "")
+        self.http = HttpConfig(header_tags=header_tags)
+
         # Master switch for turning on and off trace search by default
         # this weird invocation of get_env is meant to read the DD_ANALYTICS_ENABLED
         # legacy environment variable. It should be removed in the future
@@ -203,6 +207,7 @@ def trace_headers(self, whitelist):
         return self
 
     def header_is_traced(self, header_name):
+        # type: (str) -> bool
         """
         Returns whether or not the current header should be traced.
         :param header_name: the header name
@@ -211,6 +216,10 @@ def header_is_traced(self, header_name):
         """
         return self.http.header_is_traced(header_name)
 
+    def _header_tag_name(self, header_name):
+        # type: (str) -> Optional[str]
+        return self.http._header_tag_name(header_name)
+
     def _get_service(self, default=None):
         """
         Returns the globally configured service.

ddtrace/settings/http.py

@@ -1,8 +1,10 @@
 from typing import List
+from typing import Mapping
 from typing import Optional
-from typing import Set
 from typing import Union
 
+import six
+
 from ..internal.logger import get_logger
 from ..utils.cache import cachedmethod
 from ..utils.http import normalize_header_name
@@ -17,15 +19,27 @@ class HttpConfig(object):
     related to the http context.
     """
 
-    def __init__(self):
-        # type: () -> None
-        self._whitelist_headers = set()  # type: Set[str]
+    def __init__(self, header_tags=None):
+        # type: (Optional[Mapping[str, str]]) -> None
+        self._header_tags = {normalize_header_name(k): v for k, v in header_tags.items()} if header_tags else {}
         self.trace_query_string = None
 
+    @cachedmethod()
+    def _header_tag_name(self, header_name):
+        # type: (str) -> Optional[str]
+        if not self._header_tags:
+            return None
+
+        normalized_header_name = normalize_header_name(header_name)
+        log.debug(
+            "Checking header '%s' tracing in whitelist %s", normalized_header_name, six.viewkeys(self._header_tags)
+        )
+        return self._header_tags.get(normalized_header_name)
+
     @property
     def is_header_tracing_configured(self):
         # type: () -> bool
-        return len(self._whitelist_headers) > 0
+        return len(self._header_tags) > 0
 
     def trace_headers(self, whitelist):
         # type: (Union[List[str], str]) -> Optional[HttpConfig]
@@ -44,14 +58,15 @@ def trace_headers(self, whitelist):
             normalized_header_name = normalize_header_name(whitelist_entry)
             if not normalized_header_name:
                 continue
-            self._whitelist_headers.add(normalized_header_name)
+            # Empty tag is replaced by the default tag for this header:
+            #  Host on the request defaults to http.request.headers.host
+            self._header_tags.setdefault(normalized_header_name, "")
 
         # Mypy can't catch cached method's invalidate()
-        self.header_is_traced.invalidate()  # type: ignore[attr-defined]
+        self._header_tag_name.invalidate()  # type: ignore[attr-defined]
 
         return self
 
-    @cachedmethod()
     def header_is_traced(self, header_name):
         # type: (str) -> bool
         """
@@ -60,14 +75,9 @@ def header_is_traced(self, header_name):
         :type header_name: str
         :rtype: bool
         """
-        if not self._whitelist_headers:
-            return False
-
-        normalized_header_name = normalize_header_name(header_name)
-        log.debug("Checking header '%s' tracing in whitelist %s", normalized_header_name, self._whitelist_headers)
-        return normalized_header_name in self._whitelist_headers
+        return self._header_tag_name(header_name) is not None
 
     def __repr__(self):
         return "<{} traced_headers={} trace_query_string={}>".format(
-            self.__class__.__name__, self._whitelist_headers, self.trace_query_string
+            self.__class__.__name__, self._header_tags.keys(), self.trace_query_string
         )

ddtrace/settings/integration.py

@@ -87,17 +87,21 @@ def is_header_tracing_configured(self):
         return self.http.is_header_tracing_configured or self.global_config.http.is_header_tracing_configured
 
     def header_is_traced(self, header_name):
+        # type: (str) -> bool
         """
         Returns whether or not the current header should be traced.
         :param header_name: the header name
         :type header_name: str
         :rtype: bool
         """
-        return (
-            self.http.header_is_traced(header_name)
-            if self.http.is_header_tracing_configured
-            else self.global_config.header_is_traced(header_name)
-        )
+        return self._header_tag_name(header_name) is not None
+
+    def _header_tag_name(self, header_name):
+        # type: (str) -> Optional[str]
+        tag_name = self.http._header_tag_name(header_name)
+        if tag_name is None:
+            return self.global_config._header_tag_name(header_name)
+        return tag_name
 
     def _is_analytics_enabled(self, use_global_config):
         # DEV: analytics flag can be None which should not be taken as

docs/advanced_usage.rst

@@ -387,7 +387,9 @@ Headers tracing
 
 For a selected set of integrations, it is possible to store http headers from both requests and responses in tags.
 
-Configuration can be provided both at the global level and at the integration level.
+The recommended method is to use the ``DD_TRACE_HEADER_TAGS`` environment variable.
+
+Alternatively, configuration can be provided both at the global level and at the integration level in your application code.
 
 Examples::
 

docs/configuration.rst

@@ -142,6 +142,13 @@ below:
      - 1.0
      - A float, f, 0.0 <= f <= 1.0. f*100% of traces will be sampled.
 
+       .. _dd-trace-header-tags:
+   * - ``DD_TRACE_HEADER_TAGS``
+     - String
+     -
+     - A map of case-insensitive header keys to tag names. Automatically applies matching header values as tags on root spans.
+       For example, ``User-Agent:http.useragent,content-type:http.content_type``.
+
        .. _dd-profiling-enabled:
    * - ``DD_PROFILING_ENABLED``
      - Boolean

releasenotes/notes/trace-header-tags-88711cc9a4466d8e.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Tag traces with HTTP headers specified on the ``DD_TRACE_HEADER_TAGS`` environment variable.
+    Value must be either comma or space separated. e.g. ``Host:http.host,User-Agent,http.user_agent``
+    or ``referer:http.referer Content-Type:http.content_type``.

scripts/profiles/django-simple/run.sh

@@ -30,6 +30,9 @@ export DJANGO_ALLOWED_HOSTS="127.0.0.1"
 export DJANGO_SECRET_KEY="SECRET_KEY"
 export DATABASE_URL="sqlite:///django.db"
 
+# Tag traces with HTTP headers to benchmark the related code
+export DD_TRACE_HEADER_TAGS="User-Agent:http.user_agent,Referer:http.referer,Content-Type:http.content_type,Etag:http.etag"
+
 # Baseline
 pushd ${PREFIX}/trace-examples/python/django/django-simple
     gunicorn config.wsgi --pid ${PREFIX}/gunicorn.pid > /dev/null &

tests/contrib/urllib3/test_urllib3.py

@@ -101,7 +101,7 @@ def test_args_kwargs(self):
 
         for args, kwargs in inputs:
 
-            with self.override_http_config("urllib3", {"_whitelist_headers": set()}):
+            with self.override_http_config("urllib3", {"_header_tags": dict()}):
                 config.urllib3.http.trace_headers(["accept"])
                 pool = urllib3.connectionpool.HTTPConnectionPool(HOST, PORT)
                 out = pool.urlopen(*args, **kwargs)
@@ -314,7 +314,7 @@ def test_request_and_response_headers(self):
         assert s.get_tag("http.response.headers.access-control-allow-origin") is None
 
         # Enabled when explicitly configured
-        with self.override_http_config("urllib3", {"_whitelist_headers": set()}):
+        with self.override_http_config("urllib3", {"_header_tags": dict()}):
             config.urllib3.http.trace_headers(["my-header", "access-control-allow-origin"])
             self.http.request("GET", URL_200, headers={"my-header": "my_value"})
             spans = self.pop_spans()