Store span ids in profiling events (#1043)

* feat(tracer): add a on_start_span hook

This allows to register a function to be called when a new span is being
started by a Tracer.

* chore(black): blackify ddtrace/settings

* feat(profiling/stack): collect span ids for running thread and stack frames

Author: jd

ddtrace/profiling/collector/stack.pyx

@@ -1,8 +1,11 @@
 """CPU profiling collector."""
 from __future__ import absolute_import
 
+import collections
+import logging
 import sys
 import threading
+import weakref
 
 from ddtrace import compat
 from ddtrace.profiling import _attr
@@ -12,6 +15,23 @@ from ddtrace.profiling import event
 from ddtrace.profiling.collector import _traceback
 from ddtrace.utils import formats
 from ddtrace.vendor import attr
+from ddtrace.vendor import six
+
+
+_LOG = logging.getLogger(__name__)
+
+
+if "gevent" in sys.modules:
+    try:
+        import gevent.monkey
+    except ImportError:
+        _LOG.error("gevent loaded but unable to import gevent.monkey")
+        from ddtrace.vendor.six.moves._thread import get_ident as _thread_get_ident
+    else:
+        _thread_get_ident = gevent.monkey.get_original("thread" if six.PY2 else "_thread", "get_ident")
+else:
+    from ddtrace.vendor.six.moves._thread import get_ident as _thread_get_ident
+
 
 # NOTE: Do not use LOG here. This code runs under a real OS thread and is unable to acquire any lock of the `logging`
 # module without having gevent crashing our dedicated thread.
@@ -112,6 +132,7 @@ class StackBasedEvent(event.SampleEvent):
     thread_name = attr.ib(default=None)
     frames = attr.ib(default=None)
     nframes = attr.ib(default=None)
+    trace_ids = attr.ib(default=None)
 
 
 @event.event_class
@@ -198,7 +219,7 @@ cdef get_thread_name(thread_id):
             return "Anonymous Thread %d" % thread_id
 
 
-cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_time):
+cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_time, thread_span_links):
     current_exceptions = []
 
     IF PY_MAJOR_VERSION >= 3 and PY_MINOR_VERSION >= 7:
@@ -241,6 +262,8 @@ cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_tim
 
     running_thread_ids = {t[0] for t in running_threads}
 
+    thread_span_links.clear_threads(running_thread_ids)
+
     if ignore_profiler:
         running_thread_ids -= _periodic.PERIODIC_THREAD_IDS
 
@@ -250,11 +273,13 @@ cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_tim
     for tid, frame in running_threads:
         if ignore_profiler and tid in _periodic.PERIODIC_THREAD_IDS:
             continue
+        spans = thread_span_links.get_active_leaf_spans_from_thread_id(tid)
         frames, nframes = _traceback.pyframe_to_frames(frame, max_nframes)
         stack_events.append(
             StackSampleEvent(
                 thread_id=tid,
                 thread_name=get_thread_name(tid),
+                trace_ids=set(span.trace_id for span in spans),
                 nframes=nframes, frames=frames,
                 wall_time_ns=wall_time,
                 cpu_time_ns=cpu_time[tid],
@@ -281,6 +306,59 @@ cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_tim
     return stack_events, exc_events
 
 
+@attr.s(slots=True, eq=False)
+class _ThreadSpanLinks(object):
+
+    _thread_id_to_spans = attr.ib(factory=lambda: collections.defaultdict(weakref.WeakSet), repr=False, init=False)
+
+    # We do not use a lock because:
+    # - When adding new items, it's not possible for a same thread to call `link_span` at different time
+    #   Since the WeakSet are per-thread, there's no chance of creating 2 WeakSet for the same thread
+    # - When reading the span WeakSet for a thread, the set is copied which means that if it was ever mutated during
+    #   our reading, we wouldn't care.
+    #   In practice, here, it won't even mutate during the reading since this only used by the stack collector which
+    #   already owns the GIL.
+
+    def link_span(self, span):
+        """Link a span to its running environment.
+
+        Track threads, tasks, etc.
+        """
+        self._thread_id_to_spans[_thread_get_ident()].add(span)
+
+    def clear_threads(self, existing_thread_ids):
+        """Clear the stored list of threads based on the list of existing thread ids.
+
+        If any thread that is part of this list was stored, its data will be deleted.
+
+        :param existing_thread_ids: A set of thread ids to keep.
+        """
+        # Iterate over a copy of the list of keys in case it's mutated during our iteration.
+        for thread_id in list(self._thread_id_to_spans.keys()):
+            if thread_id not in existing_thread_ids:
+                del self._thread_id_to_spans[thread_id]
+
+    def get_active_leaf_spans_from_thread_id(self, thread_id):
+        """Return the latest active spans for a thread.
+
+        In theory this should return a single span, though if multiple children span are active without being finished,
+        there can be several spans returned.
+
+        :param thread_id: The thread id.
+        :return: A set with the active spans.
+        """
+        spans = set(self._thread_id_to_spans.get(thread_id, ()))
+        # Iterate over a copy so we can modify the original
+        for span in set(spans):
+            if not span.finished:
+                try:
+                    spans.remove(span._parent)
+                except KeyError:
+                    pass
+
+        return {span for span in spans if not span.finished}
+
+
 @attr.s(slots=True)
 class StackCollector(collector.PeriodicCollector):
     """Execution stacks collector."""
@@ -295,8 +373,10 @@ class StackCollector(collector.PeriodicCollector):
     max_time_usage_pct = attr.ib(factory=_attr.from_env("DD_PROFILING_MAX_TIME_USAGE_PCT", 2, float))
     nframes = attr.ib(factory=_attr.from_env("DD_PROFILING_MAX_FRAMES", 64, int))
     ignore_profiler = attr.ib(factory=_attr.from_env("DD_PROFILING_IGNORE_PROFILER", True, formats.asbool))
+    tracer = attr.ib(default=None)
     _thread_time = attr.ib(init=False, repr=False)
     _last_wall_time = attr.ib(init=False, repr=False)
+    _thread_span_links = attr.ib(factory=_ThreadSpanLinks, init=False, repr=False)
 
     @max_time_usage_pct.validator
     def _check_max_time_usage(self, attribute, value):
@@ -306,8 +386,15 @@ class StackCollector(collector.PeriodicCollector):
     def start(self):
         self._thread_time = ThreadTime()
         self._last_wall_time = compat.monotonic_ns()
+        if self.tracer is not None:
+            self.tracer.on_start_span(self._thread_span_links.link_span)
         super(StackCollector, self).start()
 
+    def stop(self):
+        super(StackCollector, self).stop()
+        if self.tracer is not None:
+            self.tracer.deregister_on_start_span(self._thread_span_links.link_span)
+
     def _compute_new_interval(self, used_wall_time_ns):
         interval = (used_wall_time_ns / (self.max_time_usage_pct / 100.0)) - used_wall_time_ns
         return max(interval / 1e9, self.MIN_INTERVAL_TIME)
@@ -319,7 +406,7 @@ class StackCollector(collector.PeriodicCollector):
         self._last_wall_time = now
 
         all_events = stack_collect(
-            self.ignore_profiler, self._thread_time, self.nframes, self.interval, wall_time,
+            self.ignore_profiler, self._thread_time, self.nframes, self.interval, wall_time, self._thread_span_links,
         )
 
         used_wall_time_ns = compat.monotonic_ns() - now

ddtrace/profiling/profiler.py

@@ -2,6 +2,7 @@
 import logging
 import os
 
+import ddtrace
 from ddtrace.profiling import recorder
 from ddtrace.profiling import scheduler
 from ddtrace.vendor import attr
@@ -34,7 +35,7 @@ def _build_default_exporters():
 def _build_default_collectors():
     r = recorder.Recorder()
     return [
-        stack.StackCollector(r),
+        stack.StackCollector(r, tracer=ddtrace.tracer),
         memory.MemoryCollector(r),
         exceptions.UncaughtExceptionCollector(r),
         threading.LockCollector(r),

ddtrace/settings/__init__.py

@@ -8,10 +8,10 @@
 config = Config()
 
 __all__ = [
-    'config',
-    'Config',
-    'ConfigException',
-    'HttpConfig',
-    'Hooks',
-    'IntegrationConfig',
+    "config",
+    "Config",
+    "ConfigException",
+    "HttpConfig",
+    "Hooks",
+    "IntegrationConfig",
 ]

ddtrace/settings/config.py

@@ -38,18 +38,17 @@ class Config(object):
     this instance to register their defaults, so that they're public
     available and can be updated by users.
     """
+
     def __init__(self):
         # use a dict as underlying storing mechanism
         self._config = {}
         self._http = HttpConfig()
         # 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
-        legacy_config_value = get_env('analytics', 'enabled', default=False)
+        legacy_config_value = get_env("analytics", "enabled", default=False)
 
-        self.analytics_enabled = asbool(
-            get_env('trace', 'analytics_enabled', default=legacy_config_value)
-        )
+        self.analytics_enabled = asbool(get_env("trace", "analytics_enabled", default=legacy_config_value))
 
         # DEV: we don't use `self._get_service()` here because {DD,DATADOG}_SERVICE and
         # {DD,DATADOG}_SERVICE_NAME (deprecated) are distinct functionalities.
@@ -59,13 +58,9 @@ def __init__(self):
         self.version = get_env("version")
         self.logs_injection = asbool(get_env("logs", "injection", default=False))
 
-        self.report_hostname = asbool(
-            get_env('trace', 'report_hostname', default=False)
-        )
+        self.report_hostname = asbool(get_env("trace", "report_hostname", default=False))
 
-        self.health_metrics_enabled = asbool(
-            get_env('trace', 'health_metrics_enabled', default=False)
-        )
+        self.health_metrics_enabled = asbool(get_env("trace", "health_metrics_enabled", default=False))
 
     def __getattr__(self, name):
         if name not in self._config:
@@ -81,7 +76,7 @@ def get_from(self, obj):
         """
         pin = Pin.get_from(obj)
         if pin is None:
-            log.debug('No configuration found for %s', obj)
+            log.debug("No configuration found for %s", obj)
             return {}
 
         return pin._config
@@ -155,5 +150,5 @@ def _get_service(self, default=None):
 
     def __repr__(self):
         cls = self.__class__
-        integrations = ', '.join(self._config.keys())
-        return '{}.{}({})'.format(cls.__module__, cls.__name__, integrations)
+        integrations = ", ".join(self._config.keys())
+        return "{}.{}({})".format(cls.__module__, cls.__name__, integrations)

ddtrace/settings/exceptions.py

@@ -2,4 +2,5 @@ class ConfigException(Exception):
     """Configuration exception when an integration that is not available
     is called in the `Config` object.
     """
+
     pass

ddtrace/settings/http.py

@@ -46,9 +46,10 @@ def header_is_traced(self, header_name):
         :rtype: bool
         """
         normalized_header_name = normalize_header_name(header_name)
-        log.debug('Checking header \'%s\' tracing in whitelist %s', normalized_header_name, self._whitelist_headers)
+        log.debug("Checking header '%s' tracing in whitelist %s", normalized_header_name, self._whitelist_headers)
         return normalized_header_name in self._whitelist_headers
 
     def __repr__(self):
-        return '<{} traced_headers={} trace_query_string={}>'.format(
-            self.__class__.__name__, self._whitelist_headers, self.trace_query_string)
+        return "<{} traced_headers={} trace_query_string={}>".format(
+            self.__class__.__name__, self._whitelist_headers, self.trace_query_string
+        )

ddtrace/settings/integration.py

@@ -21,6 +21,7 @@ class IntegrationConfig(AttrDict):
         config.flask['service_name'] = 'my-service-name'
         config.flask.service_name = 'my-service-name'
     """
+
     def __init__(self, global_config, name, *args, **kwargs):
         """
         :param global_config:
@@ -32,19 +33,19 @@ def __init__(self, global_config, name, *args, **kwargs):
 
         # Set internal properties for this `IntegrationConfig`
         # DEV: By-pass the `__setattr__` overrides from `AttrDict` to set real properties
-        object.__setattr__(self, 'global_config', global_config)
-        object.__setattr__(self, 'integration_name', name)
-        object.__setattr__(self, 'hooks', Hooks())
-        object.__setattr__(self, 'http', HttpConfig())
+        object.__setattr__(self, "global_config", global_config)
+        object.__setattr__(self, "integration_name", name)
+        object.__setattr__(self, "hooks", Hooks())
+        object.__setattr__(self, "http", HttpConfig())
 
         # Set default analytics configuration, default is disabled
         # DEV: Default to `None` which means do not set this key
         # Inject environment variables for integration
-        analytics_enabled_env = get_env(name, 'analytics_enabled')
+        analytics_enabled_env = get_env(name, "analytics_enabled")
         if analytics_enabled_env is not None:
             analytics_enabled_env = asbool(analytics_enabled_env)
-        self.setdefault('analytics_enabled', analytics_enabled_env)
-        self.setdefault('analytics_sample_rate', float(get_env(name, 'analytics_sample_rate', default=1.0)))
+        self.setdefault("analytics_enabled", analytics_enabled_env)
+        self.setdefault("analytics_sample_rate", float(get_env(name, "analytics_sample_rate", default=1.0)))
 
     def __deepcopy__(self, memodict=None):
         new = IntegrationConfig(self.global_config, self.integration_name, deepcopy(dict(self), memodict))
@@ -92,7 +93,7 @@ def get_analytics_sample_rate(self, use_global_config=False):
         configuration
         """
         if self._is_analytics_enabled(use_global_config):
-            analytics_sample_rate = getattr(self, 'analytics_sample_rate', None)
+            analytics_sample_rate = getattr(self, "analytics_sample_rate", None)
             # return True if attribute is None or attribute not found
             if analytics_sample_rate is None:
                 return True
@@ -105,5 +106,5 @@ def get_analytics_sample_rate(self, use_global_config=False):
 
     def __repr__(self):
         cls = self.__class__
-        keys = ', '.join(self.keys())
-        return '{}.{}({})'.format(cls.__module__, cls.__name__, keys)
+        keys = ", ".join(self.keys())
+        return "{}.{}({})".format(cls.__module__, cls.__name__, keys)

ddtrace/tracer.py

@@ -19,6 +19,7 @@
 from .utils.deprecation import deprecated, RemovedInDDTrace10Warning
 from .vendor.dogstatsd import DogStatsd
 from . import compat
+from . import _hooks
 
 
 log = get_logger(__name__)
@@ -141,6 +142,30 @@ def __init__(self, url=None, dogstatsd_url=DEFAULT_DOGSTATSD_URL):
             writer=writer,
         )
 
+        self._hooks = _hooks.Hooks()
+
+    def on_start_span(self, func):
+        """Register a function to execute when a span start.
+
+        Can be used as a decorator.
+
+        :param func: The function to call when starting a span.
+                     The started span will be passed as argument.
+        """
+        self._hooks.register(self.__class__.start_span, func)
+        return func
+
+    def deregister_on_start_span(self, func):
+        """Unregister a function registered to execute when a span starts.
+
+        Can be used as a decorator.
+
+        :param func: The function to stop calling when starting a span.
+        """
+
+        self._hooks.deregister(self.__class__.start_span, func)
+        return func
+
     @property
     def debug_logging(self):
         return self.log.isEnabledFor(logging.DEBUG)
@@ -456,6 +481,8 @@ def start_span(self, name, child_of=None, service=None, resource=None, span_type
             # service(s) that may have been added.
             self._update_dogstatsd_constant_tags()
 
+        self._hooks.emit(self.__class__.start_span, span)
+
         return span
 
     def _update_dogstatsd_constant_tags(self):

pyproject.toml

@@ -86,7 +86,6 @@ exclude = '''
     | opentracer/
     | profiling/exporter/pprof_pb2.py$
     | propagation/
-    | settings/
     | vendor/
   )
   | tests/