test(profiling): port new tests over from echion (#15581)

## Description

This ports new echion tests to dd-trace-py.

Author: KowalskiThomas

tests/profiling/collector/test_asyncio_async_generator.py

@@ -0,0 +1,92 @@
+import pytest
+
+
+@pytest.mark.subprocess(
+    env=dict(
+        DD_PROFILING_OUTPUT_PPROF="/tmp/test_asyncio_executor",
+    ),
+    err=None,
+)
+# For macOS: err=None ignores expected stderr from tracer failing to connect to agent (not relevant to this test)
+def test_asyncio_executor_wall_time() -> None:
+    import asyncio
+    import os
+
+    from ddtrace.internal.datadog.profiling import stack_v2
+    from ddtrace.profiling import profiler
+    from tests.profiling.collector import pprof_utils
+
+    assert stack_v2.is_available, stack_v2.failure_msg
+    from typing import AsyncGenerator
+
+    async def deep_dependency() -> None:
+        # This is a regular (non-generator) coroutine called
+        # by an async generator.
+        await asyncio.sleep(0.15)
+
+    async def async_generator_dep(i: int) -> AsyncGenerator[int, None]:
+        # This is an async generator called by an async generator.
+        # We want to make sure that recursive async generators are correctly sampled.
+        for j in range(i):
+            await deep_dependency()
+            yield j
+
+    async def async_generator() -> AsyncGenerator[int, None]:
+        # This is an async generator called by a coroutine.
+        # We want to make sure we unwind async generators correctly.
+        for i in range(5):
+            async for j in async_generator_dep(i):
+                yield j
+
+    async def asynchronous_function() -> None:
+        # This is a normal (non-generator) coroutine that calls into an async generator.
+        # Stack samples should not stopped at this function, they should continue unwinding
+        # into the async generator.
+        async for _ in async_generator():
+            pass
+
+    async def main() -> None:
+        await asynchronous_function()
+
+    p = profiler.Profiler()
+    p.start()
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(main())
+
+    p.stop()
+
+    output_filename = os.environ["DD_PROFILING_OUTPUT_PPROF"] + "." + str(os.getpid())
+
+    profile = pprof_utils.parse_newest_profile(output_filename)
+
+    task_samples = pprof_utils.get_samples_with_label_key(profile, "task name")
+    assert len(task_samples) > 0
+
+    def loc(f_name: str) -> pprof_utils.StackLocation:
+        return pprof_utils.StackLocation(function_name=f_name, filename="", line_no=-1)
+
+    # Thread Pool Executor
+    pprof_utils.assert_profile_has_sample(
+        profile,
+        task_samples,
+        expected_sample=pprof_utils.StackEvent(
+            thread_name="MainThread",
+            task_name="Task-1",
+            locations=list(
+                reversed(
+                    [
+                        # loc("Task-1"),
+                        loc("main"),
+                        loc("asynchronous_function"),
+                        loc("async_generator"),
+                        loc("async_generator_dep"),
+                        loc("deep_dependency"),
+                        loc("sleep"),
+                    ]
+                )
+            ),
+        ),
+        print_samples_on_failure=True,
+    )

tests/profiling/collector/test_asyncio_gather_deep_coroutines.py

@@ -0,0 +1,71 @@
+import pytest
+
+
+@pytest.mark.subprocess(
+    env=dict(
+        DD_PROFILING_OUTPUT_PPROF="/tmp/test_asyncio_gather_deep_coroutines",
+    ),
+    err=None,
+)
+# For macOS: err=None ignores expected stderr from tracer failing to connect to agent (not relevant to this test)
+def test_asyncio_gather_deep_coroutines() -> None:
+    import asyncio
+    import os
+
+    from ddtrace.internal.datadog.profiling import stack_v2
+    from ddtrace.profiling import profiler
+    from tests.profiling.collector import pprof_utils
+
+    assert stack_v2.is_available, stack_v2.failure_msg
+
+    async def deeper() -> None:
+        await asyncio.sleep(1.0)
+
+    async def deep() -> None:
+        await deeper()
+
+    async def inner() -> None:
+        await deep()
+
+    async def main() -> None:
+        await asyncio.gather(inner(), inner())
+
+    p = profiler.Profiler()
+    p.start()
+
+    asyncio.run(main())
+
+    p.stop()
+
+    output_filename = os.environ["DD_PROFILING_OUTPUT_PPROF"] + "." + str(os.getpid())
+
+    profile = pprof_utils.parse_newest_profile(output_filename)
+
+    samples = pprof_utils.get_samples_with_label_key(profile, "task name")
+    assert len(samples) > 0
+
+    def loc(f_name: str) -> pprof_utils.StackLocation:
+        return pprof_utils.StackLocation(function_name=f_name, filename="", line_no=-1)
+
+    # Test that we see stacks for inner_1 and inner_2
+    pprof_utils.assert_profile_has_sample(
+        profile,
+        samples,
+        expected_sample=pprof_utils.StackEvent(
+            thread_name="MainThread",
+            locations=list(
+                reversed(
+                    [
+                        # loc("Task-1"),
+                        loc("main"),
+                        # loc("Task-2"),
+                        loc("inner"),
+                        loc("deep"),
+                        loc("deeper"),
+                        loc("sleep"),
+                    ]
+                ),
+            ),
+        ),
+        print_samples_on_failure=True,
+    )

tests/profiling/collector/test_asyncio_recursive_on_cpu_coros.py

@@ -0,0 +1,111 @@
+import pytest
+
+
+@pytest.mark.subprocess(
+    env=dict(
+        DD_PROFILING_OUTPUT_PPROF="/tmp/test_asyncio_recursive_on_cpu_coros",
+    ),
+    err=None,
+)
+# For macOS: err=None ignores expected stderr from tracer failing to connect to agent (not relevant to this test)
+def test_asyncio_recursive_on_cpu_coros():
+    import asyncio
+    import os
+    from sys import version_info as PYVERSION
+    import time
+    import uuid
+
+    from ddtrace import ext
+    from ddtrace.internal.datadog.profiling import stack_v2
+    from ddtrace.profiling import profiler
+    from ddtrace.trace import tracer
+    from tests.profiling.collector import pprof_utils
+    from tests.profiling.collector.pprof_utils import StackLocation
+
+    assert stack_v2.is_available, stack_v2.failure_msg
+
+    def sync_code() -> int:
+        target = time.time() + 1
+        result = 0
+        while time.time() < target:
+            result += 1
+
+        return result
+
+    def sync_code_outer() -> int:
+        return sync_code()
+
+    async def inner3() -> int:
+        return sync_code_outer()
+
+    async def inner2() -> int:
+        return await inner3()
+
+    async def inner1() -> int:
+        return await inner2()
+
+    async def outer():
+        return await inner1()
+
+    async def async_main():
+        return await outer()
+
+    def main_sync():
+        asyncio.run(async_main())
+
+    resource = str(uuid.uuid4())
+    span_type = ext.SpanTypes.WEB
+
+    p = profiler.Profiler(tracer=tracer)
+    p.start()
+    with tracer.trace("test_asyncio", resource=resource, span_type=span_type) as span:
+        span_id = span.span_id
+        local_root_span_id = span._local_root.span_id
+
+        main_sync()
+
+    p.stop()
+
+    output_filename = os.environ["DD_PROFILING_OUTPUT_PPROF"] + "." + str(os.getpid())
+    profile = pprof_utils.parse_newest_profile(output_filename)
+
+    def loc(f_name: str) -> StackLocation:
+        return pprof_utils.StackLocation(function_name=f_name, filename="", line_no=-1)
+
+    runner_prefix = "Runner." if PYVERSION >= (3, 11) else ""
+    base_event_loop_prefix = "BaseEventLoop." if PYVERSION >= (3, 11) else ""
+    handle_prefix = "Handle." if PYVERSION >= (3, 11) else ""
+
+    pprof_utils.assert_profile_has_sample(
+        profile,
+        list(profile.sample),
+        pprof_utils.StackEvent(
+            thread_name="MainThread",
+            span_id=span_id,
+            local_root_span_id=local_root_span_id,
+            locations=list(
+                reversed(
+                    [
+                        loc("<module>"),
+                        loc("main_sync"),
+                        loc("run"),
+                    ]
+                    + ([loc(f"{runner_prefix}run")] if PYVERSION >= (3, 11) else [])
+                    + [
+                        loc(f"{base_event_loop_prefix}run_until_complete"),
+                        loc(f"{base_event_loop_prefix}run_forever"),
+                        loc(f"{base_event_loop_prefix}_run_once"),
+                        loc(f"{handle_prefix}_run"),
+                        loc("async_main"),
+                        loc("outer"),
+                        loc("inner1"),
+                        loc("inner2"),
+                        loc("inner3"),
+                        loc("sync_code_outer"),
+                        loc("sync_code"),
+                    ]
+                )
+            ),
+        ),
+        print_samples_on_failure=True,
+    )