Skip to content

Commit 010e186

Browse files
KowalskiThomasKowalskiThomas
committed
test(profiling): add tests for asyncio.wait and asyncio.gather
1 parent 79661be commit 010e186

File tree

4 files changed

+364
-259
lines changed

4 files changed

+364
-259
lines changed

‎tests/profiling/collector/pprof_utils.py‎

Lines changed: 0 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -354,38 +354,6 @@ def assert_profile_has_sample(
354354
samples: List[pprof_pb2.Sample],
355355
expected_sample: StackEvent,
356356
):
357-
# Print all samples with line number + function name + labels
358-
print(f"\n=== Printing all {len(samples)} samples ===")
359-
for i, sample in enumerate(samples):
360-
print(f"\nSample {i}:")
361-
362-
# Print locations (stack trace)
363-
print(" Stack trace:")
364-
for j, location_id in enumerate(sample.location_id):
365-
location = get_location_with_id(profile, location_id)
366-
if location.line:
367-
line = location.line[0]
368-
function = get_function_with_id(profile, line.function_id)
369-
function_name = profile.string_table[function.name]
370-
filename = profile.string_table[function.filename]
371-
print(f" [{j}] {filename}:{line.line} in {function_name}()")
372-
373-
# Print labels
374-
print(" Labels:")
375-
for label in sample.label:
376-
key_str = profile.string_table[label.key]
377-
if label.str:
378-
value_str = profile.string_table[label.str]
379-
print(f" {key_str}: {value_str}")
380-
elif label.num:
381-
print(f" {key_str}: {label.num}")
382-
383-
# Print values
384-
if sample.value:
385-
print(f" Values: {sample.value}")
386-
387-
print("=== End of samples ===\n")
388-
389357
found = False
390358
for sample in samples:
391359
try:

‎tests/profiling_v2/collector/test_asyncio_gather.py‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
import pytest
2+
3+
4+
@pytest.mark.subprocess(
5+
env=dict(
6+
DD_PROFILING_OUTPUT_PPROF="/tmp/test_asyncio_utils_gather",
7+
),
8+
err=None,
9+
)
10+
# For macOS: err=None ignores expected stderr from tracer failing to connect to agent (not relevant to this test)
11+
def test_asyncio_gather() -> None:
12+
import asyncio
13+
import os
14+
import time
15+
import uuid
16+
17+
from ddtrace import ext
18+
from ddtrace.internal.datadog.profiling import stack_v2
19+
from ddtrace.profiling import profiler
20+
from ddtrace.trace import tracer
21+
from tests.profiling.collector import pprof_utils
22+
23+
assert stack_v2.is_available, stack_v2.failure_msg
24+
25+
sleep_time = 0.2
26+
loop_run_time = 3
27+
28+
async def inner1() -> None:
29+
start_time = time.time()
30+
while time.time() < start_time + loop_run_time:
31+
await asyncio.sleep(sleep_time)
32+
33+
async def inner2() -> None:
34+
start_time = time.time()
35+
while time.time() < start_time + loop_run_time:
36+
await asyncio.sleep(sleep_time)
37+
38+
async def outer() -> None:
39+
t1 = asyncio.create_task(inner1(), name="inner 1")
40+
t2 = asyncio.create_task(inner2(), name="inner 2")
41+
await asyncio.gather(t1, t2)
42+
43+
resource = str(uuid.uuid4())
44+
span_type = ext.SpanTypes.WEB
45+
46+
p = profiler.Profiler(tracer=tracer)
47+
p.start()
48+
with tracer.trace("test_asyncio", resource=resource, span_type=span_type) as span:
49+
span_id = span.span_id
50+
local_root_span_id = span._local_root.span_id
51+
52+
loop = asyncio.new_event_loop()
53+
asyncio.set_event_loop(loop)
54+
main_task = loop.create_task(outer(), name="outer")
55+
loop.run_until_complete(main_task)
56+
57+
p.stop()
58+
59+
output_filename = os.environ["DD_PROFILING_OUTPUT_PPROF"] + "." + str(os.getpid())
60+
61+
profile = pprof_utils.parse_newest_profile(output_filename)
62+
63+
samples_with_span_id = pprof_utils.get_samples_with_label_key(profile, "span id")
64+
assert len(samples_with_span_id) > 0
65+
66+
# get samples with task_name
67+
samples = pprof_utils.get_samples_with_label_key(profile, "task name")
68+
# The next fails if stack_v2 is not properly configured with asyncio task
69+
# tracking via ddtrace.profiling._asyncio
70+
assert len(samples) > 0
71+
72+
pprof_utils.assert_profile_has_sample(
73+
profile,
74+
samples,
75+
expected_sample=pprof_utils.StackEvent(
76+
thread_name="MainThread",
77+
task_name="outer",
78+
span_id=span_id,
79+
local_root_span_id=local_root_span_id,
80+
locations=[
81+
pprof_utils.StackLocation(
82+
function_name="outer", filename="test_asyncio_gather.py", line_no=outer.__code__.co_firstlineno + 3
83+
),
84+
# TODO: We should add the locations of the gathered Tasks here as they should be in the same Stack
85+
],
86+
),
87+
)
88+
89+
try:
90+
pprof_utils.assert_profile_has_sample(
91+
profile,
92+
samples,
93+
expected_sample=pprof_utils.StackEvent(
94+
thread_name="MainThread",
95+
task_name="outer", # TODO: This is a bug and we need to fix it, it should be "inner 1"
96+
span_id=span_id,
97+
local_root_span_id=local_root_span_id,
98+
locations=[
99+
pprof_utils.StackLocation(
100+
function_name="inner2",
101+
filename="test_asyncio_gather.py",
102+
line_no=inner2.__code__.co_firstlineno + 3,
103+
),
104+
pprof_utils.StackLocation(
105+
function_name="outer",
106+
filename="test_asyncio_gather.py",
107+
line_no=outer.__code__.co_firstlineno + 3,
108+
),
109+
],
110+
),
111+
)
112+
113+
pprof_utils.assert_profile_has_sample(
114+
profile,
115+
samples,
116+
expected_sample=pprof_utils.StackEvent(
117+
thread_name="MainThread",
118+
task_name="inner 1",
119+
span_id=span_id,
120+
local_root_span_id=local_root_span_id,
121+
locations=[
122+
pprof_utils.StackLocation(
123+
function_name="inner1",
124+
filename="test_asyncio_gather.py",
125+
line_no=inner1.__code__.co_firstlineno + 3,
126+
),
127+
pprof_utils.StackLocation(
128+
function_name="outer",
129+
filename="test_asyncio_gather.py",
130+
line_no=outer.__code__.co_firstlineno + 3,
131+
),
132+
],
133+
),
134+
)
135+
except AssertionError:
136+
pprof_utils.assert_profile_has_sample(
137+
profile,
138+
samples,
139+
expected_sample=pprof_utils.StackEvent(
140+
thread_name="MainThread",
141+
task_name="inner 2", # TODO: This is a bug and we need to fix it, it should be "inner 1"
142+
span_id=span_id,
143+
local_root_span_id=local_root_span_id,
144+
locations=[
145+
pprof_utils.StackLocation(
146+
function_name="inner2",
147+
filename="test_asyncio_gather.py",
148+
line_no=inner2.__code__.co_firstlineno + 3,
149+
),
150+
pprof_utils.StackLocation(
151+
function_name="outer",
152+
filename="test_asyncio_gather.py",
153+
line_no=outer.__code__.co_firstlineno + 3,
154+
),
155+
],
156+
),
157+
)
158+
159+
pprof_utils.assert_profile_has_sample(
160+
profile,
161+
samples,
162+
expected_sample=pprof_utils.StackEvent(
163+
thread_name="MainThread",
164+
task_name="outer",
165+
span_id=span_id,
166+
local_root_span_id=local_root_span_id,
167+
locations=[
168+
pprof_utils.StackLocation(
169+
function_name="inner1",
170+
filename="test_asyncio_gather.py",
171+
line_no=inner1.__code__.co_firstlineno + 3,
172+
),
173+
pprof_utils.StackLocation(
174+
function_name="outer",
175+
filename="test_asyncio_gather.py",
176+
line_no=outer.__code__.co_firstlineno + 3,
177+
),
178+
],
179+
),
180+
)

0 commit comments

Comments
 (0)