Skip to content

chore(profiling): improve typing #15796

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KowalskiThomas merged 1 commit into from
Jan 4, 2026
Merged

chore(profiling): improve typing #15796

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions tests/profiling/collector/pprof_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ def get_function_with_id(profile: pprof_pb2.Profile, function_id: int) -> pprof_

def assert_lock_events_of_type(
profile: pprof_pb2.Profile,
expected_events: List[LockEvent],
expected_events: Sequence[LockEvent],
event_type: LockEventType,
):
samples = get_samples_with_value_type(
Expand Down Expand Up @@ -227,8 +227,8 @@ def assert_lock_events_of_type(

def assert_lock_events(
profile: pprof_pb2.Profile,
expected_acquire_events: Union[List[LockEvent], None] = None,
expected_release_events: Union[List[LockEvent], None] = None,
expected_acquire_events: Union[Sequence[LockEvent], None] = None,
expected_release_events: Union[Sequence[LockEvent], None] = None,
):
if expected_acquire_events:
assert_lock_events_of_type(profile, expected_acquire_events, LockEventType.ACQUIRE)
Expand Down
43 changes: 24 additions & 19 deletions tests/profiling/test_main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# -*- encoding: utf-8 -*-
import multiprocessing
import os
from pathlib import Path
import sys

import pytest
Expand All @@ -10,7 +11,7 @@
from tests.utils import call_program


def test_call_script():
def test_call_script() -> None:
env = os.environ.copy()
env["DD_PROFILING_ENABLED"] = "1"
stdout, stderr, exitcode, _ = call_program(
Expand All @@ -20,8 +21,10 @@ def test_call_script():
assert exitcode == 0, (stdout, stderr)
else:
assert exitcode == 42, (stdout, stderr)
hello, pid = list(s.strip() for s in stdout.decode().strip().split("\n"))
assert hello == "hello world", stdout.decode().strip()

stdout = stdout.decode() if isinstance(stdout, bytes) else stdout
hello, _ = list(s.strip() for s in stdout.strip().split("\n"))
assert hello == "hello world", stdout.strip()


@pytest.mark.skipif(not os.getenv("DD_PROFILE_TEST_GEVENT", False), reason="Not testing gevent")
Expand All @@ -34,7 +37,7 @@ def test_call_script_gevent():
assert exitcode == 0, (stdout, stderr)


def test_call_script_pprof_output(tmp_path):
def test_call_script_pprof_output(tmp_path: Path) -> None:
"""This checks if the pprof output and atexit register work correctly.

The script does not run for one minute, so if the `stop_on_exit` flag is broken, this test will fail.
Expand All @@ -54,24 +57,27 @@ def test_call_script_pprof_output(tmp_path):
assert exitcode == 0, (stdout, stderr)
else:
assert exitcode == 42, (stdout, stderr)
_, pid = list(s.strip() for s in stdout.decode().strip().split("\n"))
profile = pprof_utils.parse_newest_profile(filename + "." + str(pid))

stdout = stdout.decode() if isinstance(stdout, bytes) else stdout
_, pid = list(s.strip() for s in stdout.strip().split("\n"))
profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}")
samples = pprof_utils.get_samples_with_value_type(profile, "cpu-time")
assert len(samples) > 0


@pytest.mark.skipif(sys.platform == "win32", reason="fork only available on Unix")
def test_fork(tmp_path):
def test_fork(tmp_path: Path) -> None:
filename = str(tmp_path / "pprof")
env = os.environ.copy()
env["DD_PROFILING_OUTPUT_PPROF"] = filename
env["DD_PROFILING_CAPTURE_PCT"] = "100"
stdout, stderr, exitcode, pid = call_program(
stdout, _, exitcode, pid = call_program(
"python", os.path.join(os.path.dirname(__file__), "simple_program_fork.py"), env=env
)
assert exitcode == 0
child_pid = stdout.decode().strip()
profile = pprof_utils.parse_newest_profile(filename + "." + str(pid))
stdout = stdout.decode() if isinstance(stdout, bytes) else stdout
child_pid = stdout.strip()
profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}")
parent_expected_acquire_events = [
pprof_utils.LockAcquireEvent(
caller_name="<module>",
Expand Down Expand Up @@ -134,11 +140,9 @@ def test_fork(tmp_path):

@pytest.mark.skipif(sys.platform == "win32", reason="fork only available on Unix")
@pytest.mark.skipif(not os.getenv("DD_PROFILE_TEST_GEVENT", False), reason="Not testing gevent")
def test_fork_gevent():
def test_fork_gevent() -> None:
env = os.environ.copy()
stdout, stderr, exitcode, pid = call_program(
"python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"), env=env
)
_, _, exitcode, _ = call_program("python", os.path.join(os.path.dirname(__file__), "gevent_fork.py"), env=env)
assert exitcode == 0


Expand All @@ -149,7 +153,7 @@ def test_fork_gevent():
"method",
set(methods) - {"forkserver", "fork"},
)
def test_multiprocessing(method, tmp_path):
def test_multiprocessing(method: str, tmp_path: Path) -> None:
filename = str(tmp_path / "pprof")
env = os.environ.copy()
env["DD_PROFILING_OUTPUT_PPROF"] = filename
Expand All @@ -163,8 +167,9 @@ def test_multiprocessing(method, tmp_path):
env=env,
)
assert exitcode == 0, (stdout, stderr)
pid, child_pid = list(s.strip() for s in stdout.decode().strip().split("\n"))
profile = pprof_utils.parse_newest_profile(filename + "." + str(pid))
stdout = stdout.decode() if isinstance(stdout, bytes) else stdout
pid, child_pid = list(s.strip() for s in stdout.strip().split("\n"))
profile = pprof_utils.parse_newest_profile(f"{filename}.{pid}")
samples = pprof_utils.get_samples_with_value_type(profile, "cpu-time")
assert len(samples) > 0
child_profile = pprof_utils.parse_newest_profile(filename + "." + str(child_pid))
Expand All @@ -177,7 +182,7 @@ def test_multiprocessing(method, tmp_path):
env=dict(DD_PROFILING_ENABLED="1"),
err=lambda _: "RuntimeError: the memalloc module is already started" not in _,
)
def test_memalloc_no_init_error_on_fork():
def test_memalloc_no_init_error_on_fork() -> None:
import os

pid = os.fork()
Expand All @@ -195,7 +200,7 @@ def test_memalloc_no_init_error_on_fork():
out="OK\n",
err=None,
)
def test_profiler_start_up_with_module_clean_up_in_protobuf_app():
def test_profiler_start_up_with_module_clean_up_in_protobuf_app() -> None:
# This can cause segfaults if we do module clean up with later versions of
# protobuf. This is a regression test.
from google.protobuf import empty_pb2 # noqa:F401
Expand Down
Loading