fix(profiling): track existing asyncio loop if it exists

Author: KowalskiThomas

ddtrace/profiling/_asyncio.py

@@ -35,6 +35,19 @@ def _task_get_name(task: "asyncio.Task") -> str:
     return "Task-%d" % id(task)
 
 
+def _call_init_asyncio(asyncio: ModuleType) -> None:
+    from asyncio import tasks as asyncio_tasks
+
+    if sys.hexversion >= 0x030C0000:
+        scheduled_tasks = asyncio_tasks._scheduled_tasks.data  # type: ignore[attr-defined]
+        eager_tasks = asyncio_tasks._eager_tasks  # type: ignore[attr-defined]
+    else:
+        scheduled_tasks = asyncio_tasks._all_tasks.data  # type: ignore[attr-defined]
+        eager_tasks = None
+
+    stack_v2.init_asyncio(asyncio_tasks._current_tasks, scheduled_tasks, eager_tasks)  # type: ignore[attr-defined]
+
+
 @ModuleWatchdog.after_module_imported("asyncio")
 def _(asyncio):
     # type: (ModuleType) -> None
@@ -57,7 +70,15 @@ def _(asyncio):
     if THREAD_LINK is None:
         THREAD_LINK = _threading._ThreadLink()
 
-    init_stack_v2 = config.stack.v2_enabled and stack_v2.is_available
+    init_stack_v2: bool = config.stack.v2_enabled and stack_v2.is_available
+
+    if init_stack_v2:
+        try:
+            running_loop = asyncio.get_running_loop()
+            stack_v2.track_asyncio_loop(typing.cast(int, ddtrace_threading.current_thread().ident), running_loop)
+        except RuntimeError:
+            # No existing loop to track, continue
+            pass
 
     @partial(wrap, sys.modules["asyncio.events"].BaseDefaultEventLoopPolicy.set_event_loop)
     def _(f, args, kwargs):
@@ -91,17 +112,28 @@ def _(f, args, kwargs):
                 for child in children:
                     stack_v2.link_tasks(parent, child)
 
-        if sys.hexversion >= 0x030C0000:
-            scheduled_tasks = asyncio.tasks._scheduled_tasks.data
-            eager_tasks = asyncio.tasks._eager_tasks
-        else:
-            scheduled_tasks = asyncio.tasks._all_tasks.data
-            eager_tasks = None
-
-        stack_v2.init_asyncio(asyncio.tasks._current_tasks, scheduled_tasks, eager_tasks)
+        _call_init_asyncio(asyncio)
 
 
 def get_event_loop_for_thread(thread_id: int) -> typing.Union["asyncio.AbstractEventLoop", None]:
     global THREAD_LINK
 
     return THREAD_LINK.get_object(thread_id) if THREAD_LINK is not None else None
+
+
+def link_existing_loop_to_current_thread() -> None:
+    import asyncio
+
+    running_loop: typing.Union["asyncio.AbstractEventLoop", None] = None
+    try:
+        running_loop = asyncio.get_running_loop()
+        stack_v2.track_asyncio_loop(typing.cast(int, ddtrace_threading.current_thread().ident), running_loop)
+    except RuntimeError:
+        # No existing loop to track, continue
+        pass  # nosec: we only want to detect there is no existing loop to track
+    finally:
+        assert THREAD_LINK is not None  # nosec: assert is used for typing
+        if running_loop is not None:
+            THREAD_LINK.link_object(running_loop)
+
+    _call_init_asyncio(asyncio)

ddtrace/profiling/collector/threading.py

@@ -5,6 +5,7 @@
 
 from ddtrace.internal._unpatched import _threading as ddtrace_threading
 from ddtrace.internal.datadog.profiling import stack_v2
+from ddtrace.profiling import _asyncio
 from ddtrace.settings.profiling import config
 
 from . import _lock
@@ -68,3 +69,5 @@ def thread_bootstrap_inner(self, *args, **kwargs):
         # Instrument any living threads
         for thread_id, thread in ddtrace_threading._active.items():  # type: ignore[attr-defined]
             stack_v2.register_thread(thread_id, thread.native_id, thread.name)
+
+        _asyncio.link_existing_loop_to_current_thread()

releasenotes/notes/profiling-fix-untracked-existing-event-loop-6842ff15328dae9c.yaml

@@ -0,0 +1,4 @@
+fixes:
+  - |
+    profiling: this fix resolves an issue where importing the profiler module after an asyncio Event Loop had been
+    started would make the Profiler blind to the existing Event Loop and its Tasks.