Implemented distributed tracing

This handles the w3c traceparent header (currently using our own naming,
as the spec for the w3c header isn't frozen yet), as well as propagating
it to other HTTP services.

Author: beniwohli

CHANGELOG.md

@@ -13,15 +13,14 @@ Further breaking changes:
  * Some settings now require a unit for duration or size. See documentation on
    configuration for more information.
 
-
 Other changes:
  * on Python 3.7, use [contextvars](https://docs.python.org/3/library/contextvars.html) instead of threadlocals for storing
    current transaction and span. This is a necessary precursor for full asyncio support. (#291)
 
 ## Unreleased
 
  * fixed an issue with detecting names of wrapped functions that are partials (#294)
-  
+
 ## v3.0.1
 
 [Check the diff](https://github.com/elastic/apm-agent-python/compare/v3.0.0...v3.0.1)

elasticapm/base.py

@@ -196,10 +196,10 @@ def queue(self, event_type, data, flush=False):
             flush = False
         self._transport.queue(event_type, data, flush)
 
-    def begin_transaction(self, transaction_type):
+    def begin_transaction(self, transaction_type, trace_parent=None):
         """Register the start of a transaction on the client
         """
-        return self.transaction_store.begin_transaction(transaction_type)
+        return self.transaction_store.begin_transaction(transaction_type, trace_parent=trace_parent)
 
     def end_transaction(self, name=None, result=""):
         transaction = self.transaction_store.end_transaction(result, name)
@@ -351,7 +351,10 @@ def _build_msg_for_logging(
 
         transaction = get_transaction()
         if transaction:
-            event_data["transaction"] = {"id": transaction.id}
+            if transaction.trace_parent:
+                event_data["trace_id"] = transaction.trace_parent.trace_id
+            event_data["parent_id"] = transaction.id
+            event_data["transaction_id"] = transaction.id
 
         return event_data
 

elasticapm/conf/__init__.py

@@ -220,6 +220,7 @@ class Config(_ConfigBase):
     framework_version = _ConfigValue("FRAMEWORK_VERSION", default=None)
     disable_send = _BoolConfigValue("DISABLE_SEND", default=False)
     instrument = _BoolConfigValue("DISABLE_INSTRUMENTATION", default=True)
+    enable_distributed_tracing = _BoolConfigValue("ENABLE_DISTRIBUTED_TRACING", default=True)
 
 
 def setup_logging(handler, exclude=["elasticapm", "gunicorn", "south", "elasticapm.errors"]):

elasticapm/conf/constants.py

@@ -1,5 +1,8 @@
 EVENTS_API_PATH = "/intake/v2/events"
 
+TRACE_CONTEXT_VERSION = 0
+TRACEPARENT_HEADER_NAME = "elastic-apm-traceparent"
+
 TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 KEYWORD_MAX_LENGTH = 1024

elasticapm/contrib/django/apps.py

@@ -3,12 +3,16 @@
 from django.apps import AppConfig
 from django.conf import settings as django_settings
 
+from elasticapm.conf import constants
 from elasticapm.contrib.django.client import get_client
+from elasticapm.utils.disttracing import TraceParent
 
 ERROR_DISPATCH_UID = "elasticapm-exceptions"
 REQUEST_START_DISPATCH_UID = "elasticapm-request-start"
 REQUEST_FINISH_DISPATCH_UID = "elasticapm-request-stop"
 
+TRACEPARENT_HEADER_NAME_WSGI = "HTTP_" + constants.TRACEPARENT_HEADER_NAME.upper().replace("-", "_")
+
 
 class ElasticAPMConfig(AppConfig):
     name = "elasticapm.contrib.django"
@@ -38,11 +42,7 @@ def register_handlers(client):
 
     request_started.disconnect(dispatch_uid=REQUEST_START_DISPATCH_UID)
     request_started.connect(
-        lambda sender, *args, **kwargs: client.begin_transaction("request")
-        if _should_start_transaction(client)
-        else None,
-        dispatch_uid=REQUEST_START_DISPATCH_UID,
-        weak=False,
+        partial(_request_started_handler, client), dispatch_uid=REQUEST_START_DISPATCH_UID, weak=False
     )
 
     request_finished.disconnect(dispatch_uid=REQUEST_FINISH_DISPATCH_UID)
@@ -65,6 +65,23 @@ def register_handlers(client):
         client.logger.debug("Not instrumenting Celery, couldn't import")
 
 
+def _request_started_handler(client, sender, *args, **kwargs):
+    if not _should_start_transaction(client):
+        return
+    # try to find trace id
+    traceparent_header = None
+    if "environ" in kwargs:
+        traceparent_header = kwargs["environ"].get(TRACEPARENT_HEADER_NAME_WSGI)
+    elif "scope" in kwargs:
+        # TODO handle Django Channels
+        traceparent_header = None
+    if traceparent_header:
+        trace_parent = TraceParent.from_string(traceparent_header)
+    else:
+        trace_parent = None
+    client.begin_transaction("request", trace_parent=trace_parent)
+
+
 def instrument(client):
     """
     Auto-instruments code to get nice spans

elasticapm/contrib/django/middleware/__init__.py

@@ -138,10 +138,6 @@ def instrument_middlewares(self):
                 except ImportError:
                     client.logger.info("Can't instrument middleware %s", middleware_path)
 
-    def process_request(self, request):
-        if not django_settings.DEBUG or self.client.config.debug:
-            self.client.begin_transaction("request")
-
     def process_view(self, request, view_func, view_args, view_kwargs):
         request._elasticapm_view_func = view_func
 

elasticapm/contrib/flask/__init__.py

@@ -19,10 +19,11 @@
 import elasticapm
 import elasticapm.instrumentation.control
 from elasticapm.base import Client
-from elasticapm.conf import setup_logging
+from elasticapm.conf import constants, setup_logging
 from elasticapm.contrib.flask.utils import get_data_from_request, get_data_from_response
 from elasticapm.handlers.logging import LoggingHandler
 from elasticapm.utils import build_name_with_http_method_prefix
+from elasticapm.utils.disttracing import TraceParent
 
 logger = logging.getLogger("elasticapm.errors.client")
 
@@ -132,7 +133,11 @@ def init_app(self, app, **defaults):
 
     def request_started(self, app):
         if not self.app.debug or self.client.config.debug:
-            self.client.begin_transaction("request")
+            if constants.TRACEPARENT_HEADER_NAME in request.headers:
+                trace_parent = TraceParent.from_string(request.headers[constants.TRACEPARENT_HEADER_NAME])
+            else:
+                trace_parent = None
+            self.client.begin_transaction("request", trace_parent=trace_parent)
 
     def request_finished(self, app, response):
         if not self.app.debug or self.client.config.debug:

elasticapm/instrumentation/packages/base.py

@@ -10,6 +10,7 @@
 
 class AbstractInstrumentedModule(object):
     name = None
+    mutates_unsampled_arguments = False
 
     instrument_list = [
         # List of (module, method) pairs to instrument. E.g.:
@@ -101,10 +102,41 @@ def uninstrument(self):
 
     def call_if_sampling(self, module, method, wrapped, instance, args, kwargs):
         transaction = get_transaction()
-        if not transaction or not transaction.is_sampled:
+        if not transaction:
+            return wrapped(*args, **kwargs)
+        elif not transaction.is_sampled:
+            args, kwargs = self.mutate_unsampled_call_args(module, method, wrapped, instance, args, kwargs, transaction)
             return wrapped(*args, **kwargs)
         else:
             return self.call(module, method, wrapped, instance, args, kwargs)
 
+    def mutate_unsampled_call_args(self, module, method, wrapped, instance, args, kwargs, transaction):
+        """
+        Method called for unsampled wrapped calls. This can e.g. be used to add traceparent headers to the
+        underlying http call for HTTP instrumentations.
+
+        :param module:
+        :param method:
+        :param wrapped:
+        :param instance:
+        :param args:
+        :param kwargs:
+        :param transaction:
+        :return:
+        """
+        return args, kwargs
+
     def call(self, module, method, wrapped, instance, args, kwargs):
+        """
+        Wrapped call. This method should gather all necessary data, then call `wrapped` in a `capture_span` context
+        manager.
+
+        :param module: Name of the wrapped module
+        :param method: Name of the wrapped method/function
+        :param wrapped: the wrapped method/function object
+        :param instance: the wrapped instance
+        :param args: arguments to the wrapped method/function
+        :param kwargs: keyword arguments to the wrapped method/function
+        :return: the result of calling the wrapped method/function
+        """
         raise NotImplemented

elasticapm/instrumentation/packages/urllib3.py

@@ -1,6 +1,8 @@
+from elasticapm.conf import constants
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
-from elasticapm.traces import capture_span
+from elasticapm.traces import DroppedSpan, capture_span, get_transaction
 from elasticapm.utils import default_ports
+from elasticapm.utils.disttracing import TracingOptions
 
 
 class Urllib3Instrumentation(AbstractInstrumentedModule):
@@ -14,6 +16,13 @@ def call(self, module, method, wrapped, instance, args, kwargs):
         else:
             method = args[0]
 
+        headers = None
+        if "headers" in kwargs:
+            headers = kwargs["headers"]
+            if headers is None:
+                headers = {}
+                kwargs["headers"] = headers
+
         host = instance.host
 
         if instance.port != default_ports.get(instance.scheme):
@@ -27,6 +36,30 @@ def call(self, module, method, wrapped, instance, args, kwargs):
         signature = method.upper() + " " + host
 
         url = instance.scheme + "://" + host + url
+        transaction = get_transaction()
+
+        with capture_span(signature, "ext.http.urllib3", {"url": url}, leaf=True) as span:
+            # if urllib3 has been called in a leaf span, this span might be a DroppedSpan.
+            leaf_span = span
+            while isinstance(leaf_span, DroppedSpan):
+                leaf_span = leaf_span.parent
 
-        with capture_span(signature, "ext.http.urllib3", {"url": url}, leaf=True):
+            if headers is not None:
+                trace_parent = transaction.trace_parent.copy_from(
+                    span_id=leaf_span.idx, trace_options=TracingOptions(recorded=True)
+                )
+                headers[constants.TRACEPARENT_HEADER_NAME] = trace_parent.to_ascii()
             return wrapped(*args, **kwargs)
+
+    def mutate_unsampled_call_args(self, module, method, wrapped, instance, args, kwargs, transaction):
+        # since we don't have a span, we set the span id to the transaction id
+        trace_parent = transaction.trace_parent.copy_from(
+            span_id=transaction.id, trace_options=TracingOptions(recorded=False)
+        )
+        if "headers" in kwargs:
+            headers = kwargs["headers"]
+            if headers is None:
+                headers = {}
+                kwargs["headers"] = headers
+            headers[constants.TRACEPARENT_HEADER_NAME] = trace_parent.to_ascii()
+        return args, kwargs

elasticapm/traces.py

@@ -5,11 +5,11 @@
 import re
 import threading
 import timeit
-import uuid
 
 from elasticapm.conf import constants
 from elasticapm.conf.constants import SPAN, TRANSACTION
 from elasticapm.utils import compat, encoding, get_name_from_func
+from elasticapm.utils.disttracing import TraceParent, TracingOptions
 
 __all__ = ("capture_span", "tag", "set_transaction_name", "set_custom_context", "set_user_context")
 
@@ -28,9 +28,9 @@
 
 
 class Transaction(object):
-    def __init__(self, store, transaction_type="custom", is_sampled=True):
-        self.id = str(uuid.uuid4())
-        self.trace_id = None  # for later use in distributed tracing
+    def __init__(self, store, transaction_type="custom", trace_parent=None, is_sampled=True):
+        self.id = "%016x" % random.getrandbits(64)
+        self.trace_parent = trace_parent
         self.timestamp = datetime.datetime.utcnow()
         self.start_time = _time_func()
         self.name = None
@@ -61,7 +61,8 @@ def begin_span(self, name, span_type, context=None, leaf=False):
             self._span_counter += 1
         else:
             start = _time_func() - self.start_time
-            span = Span(self._span_counter, self.id, self.trace_id, name, span_type, start, context, leaf)
+            span_id = "%016x" % random.getrandbits(64) if self.trace_parent else self._span_counter - 1
+            span = Span(span_id, self.id, self.trace_parent.trace_id, name, span_type, start, context, leaf=leaf)
             span.frames = store.frames_collector_func()
             span.parent = parent_span
             self._span_counter += 1
@@ -91,7 +92,7 @@ def to_dict(self):
         self.context["tags"] = self.tags
         result = {
             "id": self.id,
-            "trace_id": self.trace_id,
+            "trace_id": self.trace_parent.trace_id,
             "name": encoding.keyword_field(self.name or ""),
             "type": encoding.keyword_field(self.transaction_type),
             "duration": self.duration * 1000,  # milliseconds
@@ -100,6 +101,10 @@ def to_dict(self):
             "sampled": self.is_sampled,
             "span_count": {"started": self._span_counter - self.dropped_spans, "dropped": self.dropped_spans},
         }
+        if self.trace_parent:
+            result["trace_id"] = self.trace_parent.trace_id
+            if self.trace_parent.span_id:
+                result["parent_id"] = self.trace_parent.span_id
         if self.is_sampled:
             result["context"] = self.context
         return result
@@ -142,17 +147,19 @@ def __init__(self, idx, transaction_id, trace_id, name, span_type, start_time, c
         self.duration = None
         self.parent = None
         self.frames = None
+        self.trace_id = trace_id
+        self.transaction_id = transaction_id
 
     def to_dict(self):
         result = {
-            "id": compat.text_type(self.idx),
+            "id": self.idx,
             "transaction_id": self.transaction_id,
             "trace_id": self.trace_id,
+            "parent_id": self.parent.idx if self.parent else self.transaction_id,
             "name": encoding.keyword_field(self.name),
             "type": encoding.keyword_field(self.type),
             "start": self.start_time * 1000,  # milliseconds
             "duration": self.duration * 1000,  # milliseconds
-            "parent": self.parent.idx if self.parent else None,
             "context": self.context,
         }
         if self.frames:
@@ -212,14 +219,24 @@ def __len__(self):
         with self.cond:
             return len(self._transactions)
 
-    def begin_transaction(self, transaction_type):
+    def begin_transaction(self, transaction_type, trace_parent=None):
         """
         Start a new transactions and bind it in a thread-local variable
 
         :returns the Transaction object
         """
-        is_sampled = self._sample_rate == 1.0 or self._sample_rate > random.random()
-        transaction = Transaction(self, transaction_type, is_sampled=is_sampled)
+        if trace_parent:
+            is_sampled = bool(trace_parent.trace_options.recorded)
+        else:
+            is_sampled = self._sample_rate == 1.0 or self._sample_rate > random.random()
+        transaction = Transaction(self, transaction_type, trace_parent=trace_parent, is_sampled=is_sampled)
+        if trace_parent is None:
+            transaction.trace_parent = TraceParent(
+                constants.TRACE_CONTEXT_VERSION,
+                "%032x" % random.getrandbits(128),
+                transaction.id,
+                TracingOptions(recorded=is_sampled),
+            )
         set_transaction(transaction)
         return transaction
 
@@ -264,7 +281,7 @@ def decorated(*args, **kwds):
     def __enter__(self):
         transaction = get_transaction()
         if transaction and transaction.is_sampled:
-            transaction.begin_span(self.name, self.type, context=self.extra, leaf=self.leaf)
+            return transaction.begin_span(self.name, self.type, context=self.extra, leaf=self.leaf)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         transaction = get_transaction()