test(profiling): add tests for asyncio.wait and asyncio.gather

Author: KowalskiThomas

tests/profiling/collector/pprof_utils.py

@@ -109,6 +109,9 @@ def __init__(self, locations: Optional[Any] = None, exception_type=None, *args,
         self.exception_type = exception_type
         super().__init__(*args, **kwargs)
 
+    def __str__(self):
+        return f"StackEvent(locations={self.locations}, exception_type={self.exception_type})"
+
 
 # A simple class to hold the expected attributes of a lock sample.
 class LockEvent(EventBaseClass):
@@ -297,15 +300,17 @@ def assert_lock_event(profile: pprof_pb2.Profile, sample: pprof_pb2.Sample, expe
     assert_base_event(profile.string_table, sample, expected_event)
 
 
-def assert_sample_has_locations(profile, sample, expected_locations: Optional[List[StackLocation]]):
+def assert_sample_has_locations(
+    profile: pprof_pb2.Profile, sample: pprof_pb2.Sample, expected_locations: Optional[Sequence[StackLocation]]
+) -> None:
     if not expected_locations:
         return
 
     expected_locations_idx = 0
     checked = False
 
     # For debug printing
-    sample_loc_strs = []
+    sample_loc_strs: list[str] = []
     # in a sample there can be multiple locations, we need to check
     # whether there's a consecutive subsequence of locations that match
     # the expected locations
@@ -354,38 +359,6 @@ def assert_profile_has_sample(
     samples: List[pprof_pb2.Sample],
     expected_sample: StackEvent,
 ):
-    # Print all samples with line number + function name + labels
-    print(f"\n=== Printing all {len(samples)} samples ===")
-    for i, sample in enumerate(samples):
-        print(f"\nSample {i}:")
-
-        # Print locations (stack trace)
-        print("  Stack trace:")
-        for j, location_id in enumerate(sample.location_id):
-            location = get_location_with_id(profile, location_id)
-            if location.line:
-                line = location.line[0]
-                function = get_function_with_id(profile, line.function_id)
-                function_name = profile.string_table[function.name]
-                filename = profile.string_table[function.filename]
-                print(f"    [{j}] {filename}:{line.line} in {function_name}()")
-
-        # Print labels
-        print("  Labels:")
-        for label in sample.label:
-            key_str = profile.string_table[label.key]
-            if label.str:
-                value_str = profile.string_table[label.str]
-                print(f"    {key_str}: {value_str}")
-            elif label.num:
-                print(f"    {key_str}: {label.num}")
-
-        # Print values
-        if sample.value:
-            print(f"  Values: {sample.value}")
-
-    print("=== End of samples ===\n")
-
     found = False
     for sample in samples:
         try:
@@ -397,4 +370,4 @@ def assert_profile_has_sample(
             if DEBUG_TEST:
                 print(e)
 
-    assert found, "Expected samples not found in profile"
+    assert found, f"Expected samples not found in profile, {str(expected_sample)}"

tests/profiling_v2/collector/test_asyncio_gather.py

@@ -0,0 +1,129 @@
+import pytest
+
+
+@pytest.mark.subprocess(
+    env=dict(
+        DD_PROFILING_OUTPUT_PPROF="/tmp/test_asyncio_utils_gather",
+    ),
+    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() -> None:
+    import asyncio
+    import os
+    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
+
+    assert stack_v2.is_available, stack_v2.failure_msg
+
+    sleep_time = 0.2
+    loop_run_time = 3
+
+    async def inner1() -> None:
+        start_time = time.time()
+        while time.time() < start_time + loop_run_time:
+            await asyncio.sleep(sleep_time)
+
+    async def inner2() -> None:
+        start_time = time.time()
+        while time.time() < start_time + loop_run_time:
+            await asyncio.sleep(sleep_time)
+
+    async def outer() -> None:
+        t1 = asyncio.create_task(inner1(), name="inner 1")
+        t2 = asyncio.create_task(inner2(), name="inner 2")
+        await asyncio.gather(t1, t2)
+
+    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
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        main_task = loop.create_task(outer(), name="outer")
+        loop.run_until_complete(main_task)
+
+    p.stop()
+
+    output_filename = os.environ["DD_PROFILING_OUTPUT_PPROF"] + "." + str(os.getpid())
+
+    profile = pprof_utils.parse_newest_profile(output_filename)
+
+    samples_with_span_id = pprof_utils.get_samples_with_label_key(profile, "span id")
+    assert len(samples_with_span_id) > 0
+
+    # get samples with task_name
+    samples = pprof_utils.get_samples_with_label_key(profile, "task name")
+    # The next fails if stack_v2 is not properly configured with asyncio task
+    # tracking via ddtrace.profiling._asyncio
+    assert len(samples) > 0
+
+    pprof_utils.assert_profile_has_sample(
+        profile,
+        samples,
+        expected_sample=pprof_utils.StackEvent(
+            thread_name="MainThread",
+            task_name="outer",
+            span_id=span_id,
+            local_root_span_id=local_root_span_id,
+            locations=[
+                pprof_utils.StackLocation(
+                    function_name="outer", filename="test_asyncio_gather.py", line_no=outer.__code__.co_firstlineno + 3
+                ),
+                # TODO: We should add the locations of the gathered Tasks here as they should be in the same Stack
+            ],
+        ),
+    )
+
+    pprof_utils.assert_profile_has_sample(
+        profile,
+        samples,
+        expected_sample=pprof_utils.StackEvent(
+            thread_name="MainThread",
+            task_name="outer",  # TODO: This is a bug and we need to fix it, it should be "inner 1"
+            span_id=span_id,
+            local_root_span_id=local_root_span_id,
+            locations=[
+                pprof_utils.StackLocation(
+                    function_name="inner2",
+                    filename="test_asyncio_gather.py",
+                    line_no=inner2.__code__.co_firstlineno + 3,
+                ),
+                pprof_utils.StackLocation(
+                    function_name="outer", filename="test_asyncio_gather.py", line_no=outer.__code__.co_firstlineno + 3
+                ),
+            ],
+        ),
+    )
+
+    pprof_utils.assert_profile_has_sample(
+        profile,
+        samples,
+        expected_sample=pprof_utils.StackEvent(
+            thread_name="MainThread",
+            task_name="inner 1",
+            span_id=span_id,
+            local_root_span_id=local_root_span_id,
+            locations=[
+                pprof_utils.StackLocation(
+                    function_name="inner1",
+                    filename="test_asyncio_gather.py",
+                    line_no=inner1.__code__.co_firstlineno + 3,
+                ),
+                pprof_utils.StackLocation(
+                    function_name="outer", filename="test_asyncio_gather.py", line_no=outer.__code__.co_firstlineno + 3
+                ),
+            ],
+        ),
+    )