diff --git a/ddtrace/internal/datadog/profiling/dd_wrapper/include/sample.hpp b/ddtrace/internal/datadog/profiling/dd_wrapper/include/sample.hpp index 0574311293f..4313279d1d5 100644 --- a/ddtrace/internal/datadog/profiling/dd_wrapper/include/sample.hpp +++ b/ddtrace/internal/datadog/profiling/dd_wrapper/include/sample.hpp @@ -63,7 +63,7 @@ class Sample SampleType type_mask; std::string errmsg; - // Timeline support works by endowing each sample with a timestamp. Collection of this data this data is cheap, but + // Timeline support works by endowing each sample with a timestamp. Collection of this data is cheap, but // due to the underlying pprof format, timeline support increases the sample cardinality. Rather than switching // around the frontend code too much, we push enablement down to whether or not timestamps get added to samples (a // 0 value suppresses the tag). However, Sample objects are short-lived, so we make the flag static. diff --git a/ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi b/ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi index 87bd4598fb4..75b0fae8de6 100644 --- a/ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi +++ b/ddtrace/internal/datadog/profiling/stack_v2/__init__.pyi @@ -1,5 +1,9 @@ -def register_thread(id: int, native_id: int, name: str) -> None: ... # noqa: A002 +import asyncio + +def register_thread(id: int, native_id: int, name: str) -> None: ... def unregister_thread(name: str) -> None: ... +def track_asyncio_loop(thread_id: int, loop: asyncio.EventLoop) -> None: ... +def link_tasks(parent: asyncio.EventLoop, child: asyncio.Task) -> None: ... is_available: bool failure_msg: str diff --git a/ddtrace/profiling/_asyncio.py b/ddtrace/profiling/_asyncio.py index 3d6d18c11bb..4bb8c03b109 100644 --- a/ddtrace/profiling/_asyncio.py +++ b/ddtrace/profiling/_asyncio.py @@ -1,8 +1,11 @@ # -*- encoding: utf-8 -*- from functools import partial import sys -from types import ModuleType # noqa:F401 -import typing # noqa:F401 +from types import ModuleType +import typing + +if typing.TYPE_CHECKING: + import asyncio from ddtrace.internal._unpatched import _threading as ddtrace_threading from ddtrace.internal.datadog.profiling import stack_v2 @@ -17,15 +20,15 @@ THREAD_LINK = None # type: typing.Optional[_threading._ThreadLink] -def current_task(loop=None): +def current_task(loop: asyncio.EventLoop | None = None) -> asyncio.Task | None: return None -def all_tasks(loop=None): +def all_tasks(loop: asyncio.EventLoop | None = None) -> typing.List[asyncio.Task] | None: return [] -def _task_get_name(task): +def _task_get_name(task: asyncio.Task) -> str: return "Task-%d" % id(task) @@ -61,6 +64,7 @@ def _(f, args, kwargs): stack_v2.track_asyncio_loop(typing.cast(int, ddtrace_threading.current_thread().ident), loop) return f(*args, **kwargs) finally: + assert THREAD_LINK is not None THREAD_LINK.clear_threads(set(sys._current_frames().keys())) if loop is not None: THREAD_LINK.link_object(loop) @@ -73,10 +77,14 @@ def _(f, args, kwargs): return f(*args, **kwargs) finally: children = get_argument_value(args, kwargs, 1, "children") + assert children is not None + # Pass an invalid positional index for 'loop' loop = get_argument_value(args, kwargs, -1, "loop") + # Link the parent gathering task to the gathered children parent = globals()["current_task"](loop) + for child in children: stack_v2.link_tasks(parent, child)