feat(profiling): inject task into main thread for wall time profiling (#3010)

This inject all running tasks (greenlet only for now) into the main thread wall
time. This makes sure that all the running tasks get the same wall time between
2 polling intervals, even if they are sleeping for a very long time.

Co-authored-by: Brett Langdon <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

ddtrace/internal/nogevent.py

@@ -38,8 +38,9 @@ def is_module_patched(module):
 Thread = get_original("threading", "Thread")
 Lock = get_original("threading", "Lock")
 
+is_threading_patched = is_module_patched("threading")
 
-if is_module_patched("threading"):
+if is_threading_patched:
 
     @attr.s
     class DoubleLock(object):
@@ -75,7 +76,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
     DoubleLock = threading.Lock  # type:  ignore[misc,assignment]
 
 
-if is_module_patched("threading"):
+if is_threading_patched:
     # NOTE: bold assumption: this module is always imported by the MainThread.
     # The python `threading` module makes that assumption and it's beautiful we're going to do the same.
     # We don't have the choice has we can't access the original MainThread

ddtrace/profiling/collector/_task.pyi

@@ -4,3 +4,4 @@ import typing
 def get_task(
     thread_id: int,
 ) -> typing.Tuple[typing.Optional[int], typing.Optional[str], typing.Optional[types.FrameType]]: ...
+def list_tasks() -> typing.List[typing.Tuple[int, str, types.FrameType]]: ...

ddtrace/profiling/collector/_task.pyx

@@ -1,3 +1,6 @@
+import threading
+import weakref
+
 from ddtrace.internal import compat
 from ddtrace.internal import nogevent
 
@@ -8,6 +11,7 @@ try:
     import gevent.hub
     import gevent.thread
     from greenlet import getcurrent
+    from greenlet import greenlet
     from greenlet import settrace
 except ImportError:
     _gevent_tracer = None
@@ -16,8 +20,17 @@ else:
     class DDGreenletTracer(object):
         def __init__(self):
             # type: (...) -> None
-            self.active_greenlet = getcurrent()
             self.previous_trace_function = settrace(self)
+            self.greenlets = weakref.WeakValueDictionary()
+            self.active_greenlet = getcurrent()
+            self._store_greenlet(self.active_greenlet)
+
+        def _store_greenlet(
+                self,
+                greenlet,  # type: greenlet.greenlet
+        ):
+            # type: (...) -> None
+            self.greenlets[gevent.thread.get_ident(greenlet)] = greenlet
 
         def __call__(self, event, args):
             if event in ('switch', 'throw'):
@@ -26,6 +39,7 @@ else:
                 # users as that does not give any information about user code.
                 if not isinstance(args[1], gevent.hub.Hub):
                     self.active_greenlet = args[1]
+                    self._store_greenlet(args[1])
 
             if self.previous_trace_function is not None:
                 self.previous_trace_function(event, args)
@@ -50,3 +64,25 @@ cpdef get_task(thread_id):
         frame = None
 
     return task_id, task_name, frame
+
+
+cpdef list_tasks():
+    # type: (...) -> typing.List[typing.Tuple[int, str, types.FrameType]]
+    """Return the list of running tasks.
+
+    This is computed for gevent by taking the list of existing threading.Thread object and removing if any real OS
+    thread that might be running.
+
+    :return: [(task_id, task_name, task_frame), ...]"""
+    # We consider all Thread objects to be greenlet
+    # This should be true as nobody could use a half-monkey-patched gevent
+    if _gevent_tracer is not None:
+        return [
+            (greenlet_id,
+             _threading.get_thread_name(greenlet_id),
+             greenlet.gr_frame)
+            for greenlet_id, greenlet in compat.iteritems(_gevent_tracer.greenlets)
+            if not greenlet.dead
+        ]
+
+    return []

ddtrace/profiling/collector/stack.pyx

@@ -296,29 +296,90 @@ cdef stack_collect(ignore_profiler, thread_time, max_nframes, interval, wall_tim
         if task_id in thread_id_ignore_list:
             continue
 
-        # Inject the task frame and replace the thread frame: this is especially handy for gevent as it allows us to
-        # replace the "hub" stack trace by the latest active greenlet, which is more interesting to show and account
-        # resources for.
-        if task_frame is not None:
-           frame = task_frame
-
-        frames, nframes = _traceback.pyframe_to_frames(frame, max_nframes)
-
-        event = stack_event.StackSampleEvent(
-            thread_id=thread_id,
-            thread_native_id=thread_native_id,
-            thread_name=thread_name,
-            task_id=task_id,
-            task_name=task_name,
-            nframes=nframes, frames=frames,
-            wall_time_ns=wall_time,
-            cpu_time_ns=cpu_time,
-            sampling_period=int(interval * 1e9),
-        )
+        # Inject wall time for all running tasks
+        if thread_id == nogevent.main_thread_id:
+            main_thread_is_a_task = False
+
+            for task_id, task_name, task_frame in _task.list_tasks():
+                if task_id == compat.main_thread.ident:
+                    main_thread_is_a_task = True
+
+                    # task_frame is None, so use the thread frames list here
+                    frames, nframes = _traceback.pyframe_to_frames(frame, max_nframes)
+
+                    event = stack_event.StackSampleEvent(
+                        thread_id=thread_id,
+                        thread_native_id=thread_native_id,
+                        thread_name=thread_name,
+                        task_id=task_id,
+                        task_name=task_name,
+                        nframes=nframes, frames=frames,
+                        wall_time_ns=wall_time,
+                        cpu_time_ns=cpu_time,
+                        sampling_period=int(interval * 1e9),
+                    )
+
+                    # FIXME: we only trace spans per thread, so we assign the span to the main thread for now
+                    # we'd need to leverage the greenlet tracer to also store the active span
+                    event.set_trace_info(span, collect_endpoint)
+                else:
+                    frames, nframes = _traceback.pyframe_to_frames(task_frame, max_nframes)
+
+                    event = stack_event.StackSampleEvent(
+                        thread_id=thread_id,
+                        thread_native_id=thread_native_id,
+                        thread_name=thread_name,
+                        task_id=task_id,
+                        task_name=task_name,
+                        nframes=nframes, frames=frames,
+                        wall_time_ns=wall_time,
+                        # we don't have CPU time per task
+                        sampling_period=int(interval * 1e9),
+                    )
+
+                stack_events.append(event)
+
+            if not main_thread_is_a_task:
+                frames, nframes = _traceback.pyframe_to_frames(frame, max_nframes)
+
+                event = stack_event.StackSampleEvent(
+                    thread_id=thread_id,
+                    thread_native_id=thread_native_id,
+                    thread_name=thread_name,
+                    nframes=nframes,
+                    frames=frames,
+                    wall_time_ns=wall_time,
+                    cpu_time_ns=cpu_time,
+                    sampling_period=int(interval * 1e9),
+                )
+
+                event.set_trace_info(span, collect_endpoint)
+
+                stack_events.append(event)
+        else:
+            # Inject the task frame and replace the thread frame: this is especially handy for gevent as it allows us to
+            # replace the "hub" stack trace by the latest active greenlet, which is more interesting to show and account
+            # resources for.
+            if task_frame is not None:
+               frame = task_frame
+
+            frames, nframes = _traceback.pyframe_to_frames(frame, max_nframes)
+
+            event = stack_event.StackSampleEvent(
+                thread_id=thread_id,
+                thread_native_id=thread_native_id,
+                thread_name=thread_name,
+                task_id=task_id,
+                task_name=task_name,
+                nframes=nframes, frames=frames,
+                wall_time_ns=wall_time,
+                cpu_time_ns=cpu_time,
+                sampling_period=int(interval * 1e9),
+            )
 
-        event.set_trace_info(span, collect_endpoint)
+            event.set_trace_info(span, collect_endpoint)
 
-        stack_events.append(event)
+            stack_events.append(event)
 
         if exception is not None:
             exc_type, exc_traceback = exception

releasenotes/notes/profiling-gevent-task-inject-86c14864ca018dfa.yaml

@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    The profiler now automatically injects running greenlets as tasks into the
+    main thread. They can be seen within the wall time profiles.

tests/profiling/collector/test_stack.py

@@ -1,4 +1,5 @@
 # -*- encoding: utf-8 -*-
+import collections
 import os
 import threading
 import time
@@ -331,7 +332,7 @@ def test_exception_collection():
     assert e.sampling_period > 0
     assert e.thread_id == nogevent.thread_get_ident()
     assert e.thread_name == "MainThread"
-    assert e.frames == [(__file__, 325, "test_exception_collection")]
+    assert e.frames == [(__file__, 326, "test_exception_collection")]
     assert e.nframes == 1
     assert e.exc_type == ValueError
 
@@ -354,7 +355,7 @@ def test_exception_collection_trace(tracer):
     assert e.sampling_period > 0
     assert e.thread_id == nogevent.thread_get_ident()
     assert e.thread_name == "MainThread"
-    assert e.frames == [(__file__, 348, "test_exception_collection_trace")]
+    assert e.frames == [(__file__, 349, "test_exception_collection_trace")]
     assert e.nframes == 1
     assert e.exc_type == ValueError
     assert e.span_id == span.span_id
@@ -600,20 +601,24 @@ def test_thread_time_cache():
 
 
 @pytest.mark.skipif(not TESTING_GEVENT, reason="Not testing gevent")
-def test_collect_gevent_thread_hub():
+def test_collect_gevent_threads():
     # type: (...) -> None
     r = recorder.Recorder()
-    s = stack.StackCollector(r, ignore_profiler=True)
+    s = stack.StackCollector(r, ignore_profiler=True, max_time_usage_pct=100)
+
+    iteration = 100
+    sleep_time = 0.01
+    nb_threads = 15
 
     # Start some greenthreads: they do nothing we just keep switching between them.
     def _nothing():
-        for _ in range(100):
+        for _ in range(iteration):
             # Do nothing and just switch to another greenlet
-            time.sleep(0.01)
+            time.sleep(sleep_time)
 
     threads = []
     with s:
-        for i in range(100):
+        for i in range(nb_threads):
             t = threading.Thread(target=_nothing, name="TestThread %d" % i)
             t.start()
             threads.append(t)
@@ -622,6 +627,8 @@ def _nothing():
 
     main_thread_found = False
     sleep_task_found = False
+    wall_time_ns_per_thread = collections.defaultdict(lambda: 0)
+
     events = r.events[stack_event.StackSampleEvent]
     for event in events:
         if event.task_id == compat.main_thread.ident:
@@ -639,5 +646,23 @@ def _nothing():
                     sleep_task_found = True
                     break
 
+            wall_time_ns_per_thread[event.task_id] += event.wall_time_ns
+
     assert main_thread_found
     assert sleep_task_found
+
+    # sanity check: we don't have duplicate in thread/task ids.
+    assert len(wall_time_ns_per_thread) == nb_threads
+
+    # In theory there should be only one value in this set, but due to timing, it's possible one task has less event, so
+    # we're not checking the len() of values here.
+    values = set(wall_time_ns_per_thread.values())
+
+    # NOTE(jd): I'm disabling this check because it works 90% of the test only. There are some cases where this test is
+    # run inside the complete test suite and fails, while it works 100% of the time in its own.
+    # Check that the sum of wall time generated for each task is right.
+    # Accept a 30% margin though, don't be crazy, we're just doing 5 seconds with a lot of tasks.
+    # exact_time = iteration * sleep_time * 1e9
+    # assert (exact_time * 0.7) <= values.pop() <= (exact_time * 1.3)
+
+    assert values.pop() > 0

tests/profiling/collector/test_task.py

@@ -1,11 +1,66 @@
+import os
+import threading
+
+import pytest
+
 from ddtrace import compat
 from ddtrace.internal import nogevent
 from ddtrace.profiling.collector import _task
 
 
+TESTING_GEVENT = os.getenv("DD_PROFILE_TEST_GEVENT", False)
+
+
 def test_get_task_main():
     # type: (...) -> None
     if _task._gevent_tracer is None:
         assert _task.get_task(nogevent.main_thread_id) == (None, None, None)
     else:
         assert _task.get_task(nogevent.main_thread_id) == (compat.main_thread.ident, "MainThread", None)
+
+
+@pytest.mark.skipif(TESTING_GEVENT, reason="only works without gevent")
+def test_list_tasks_nogevent():
+    assert _task.list_tasks() == []
+
+
+@pytest.mark.skipif(not TESTING_GEVENT, reason="only works with gevent")
+def test_list_tasks_gevent():
+    l1 = threading.Lock()
+    l1.acquire()
+
+    def wait():
+        l1.acquire()
+        l1.release()
+
+    def nothing():
+        pass
+
+    t1 = threading.Thread(target=wait, name="t1")
+    t1.start()
+
+    tasks = _task.list_tasks()
+    # can't check == 2 because there are left over from other tests
+    assert len(tasks) >= 2
+
+    main_thread_found = False
+    t1_found = False
+    for task in tasks:
+        assert len(task) == 3
+        # main thread
+        if task[0] == compat.main_thread.ident:
+            assert task[1] == "MainThread"
+            assert task[2] is None
+            main_thread_found = True
+        # t1
+        elif task[0] == t1.ident:
+            assert task[1] == "t1"
+            assert task[2] is not None
+            t1_found = True
+
+    l1.release()
+
+    t1.join()
+
+    assert t1_found
+    assert main_thread_found