From d5882270764879e1ab44e55ec8b5281f542b535c Mon Sep 17 00:00:00 2001 From: Pablo Galindo Date: Tue, 15 Nov 2022 17:19:54 +0000 Subject: [PATCH] Add a new marker to check for memory leaks Users have indicated that it will be very useful if the plugin exposes a way to detect memory leaks in tests. This is possible, but is a bit tricky as the interpreter can allocate memory for internal caches, as well as user functions. To make this more reliable, the new marker will take two parameters: * The watermark of memory to ignore. If the memory leaked by the test is higher than this value, the test will fail and it will pass otherwise. * The number of warmup runs. This allows to run the test multiple times (assuming it passes) before actually checking for leaks. This allows to warmup user and interpreter caches. Signed-off-by: Pablo Galindo --- docs/usage.rst | 59 ++++++++++++++++++++++++++ src/pytest_memray/marks.py | 84 ++++++++++++++++++++++++++++++++++--- src/pytest_memray/plugin.py | 59 ++++++++++++++++++++++---- tests/test_pytest_memray.py | 65 ++++++++++++++++++++++++++++ 4 files changed, 253 insertions(+), 14 deletions(-) diff --git a/docs/usage.rst b/docs/usage.rst index e27c3e5..e54de2e 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -40,6 +40,10 @@ validations on tests when this plugin is enabled. ``limit_memory`` ---------------- +.. py:function:: limit_memory(memory_limit: str) + + Fail the execution of the test if the test allocates more memory than allowed + When this marker is applied to a test, it will cause the test to fail if the execution of the test allocates more memory than allowed. It takes a single argument with a string indicating the maximum memory that the test can allocate. @@ -61,3 +65,58 @@ Example of usage: @pytest.mark.limit_memory("24 MB") def test_foobar(): pass # do some stuff that allocates memory + +``check_leaks`` +--------------- + +.. py:function:: check_leaks(max_leaked_memory: str , *, warmups: int=0) + + Fail the execution of the test if the test leaks more memory than allowed. + +.. warning:: + + Running this marker can cause the test to take extra time to run, as it will + activate some additional check in ``memray`` to ensure that all leaks are + detected. Additionally, if the ``warmups`` argument is provided, the test + will be run multiple times before the final execution, which will also + increase the total time it takes to complete the test run. + +When this marker is applied to a test, it will cause the test to fail if the execution +of the test leaks more memory than allowed. It takes a single positional argument with a +string indicating the maximum memory that the test is allowed to leak. + +The format for the string is `` ([KMGTP]B|B)``. The marker will raise +``ValueError`` if the string format cannot be parsed correctly. + +The marker also takes an optional keyword-onlly argument ``warmups``. This argument +indicates the number of times the test will be executed before the memory leak check +is performed. This is useful to avoid false positives when the test is allocating +memory in caches that will survive the test execution or to account for memory +allocated (and not released) by the interpreter under normal execution. + +.. important:: + + The Python interpreter allocates memory in some internal caches that may not be + released when the test ends. This means that tests using this marker may need to + give some room to account for this or to use the ``warmups`` argument. The amount + of memory used by the interpreter and the location assigned to it can also change + between minor versions of Python. + +Using this marker with "0B" as the maximum allowed leaked memory can be a bit +challenging and will surely require a nonzero value for the ``warmups`` +argument. + +If you need to further investigate the leaks, you can use the ``--memray-bin-path`` +to specify a directory where ``memray`` will store the binary files with the results. +This will allow you to use the ``memray`` CLI to inspect the results using the different +reporters available. Check `the memray docs +`_ for more information. + +Example of usage: + +.. code-block:: python + + @pytest.mark.check_leaks("1 MB", warmups=3) + def test_foobar(): + pass # do some stuff that cannot leak memory + diff --git a/src/pytest_memray/marks.py b/src/pytest_memray/marks.py index 6efbaf0..ebad301 100644 --- a/src/pytest_memray/marks.py +++ b/src/pytest_memray/marks.py @@ -1,10 +1,15 @@ from __future__ import annotations from dataclasses import dataclass +from pathlib import Path +from typing import Any +from typing import Callable from typing import Optional from typing import Tuple from memray import AllocationRecord +from memray import FileReader +from memray import Tracker from .utils import parse_memory_string from .utils import sizeof_fmt @@ -46,15 +51,84 @@ def long_repr(self) -> str: return f"Test was limited to {max_memory_str} but allocated {total_memory_str}" -def limit_memory( - limit: str, *, _allocations: list[AllocationRecord] -) -> Optional[_MemoryInfo]: +@dataclass +class _LeakedInfo: + """Type that holds all memray-related info for a failed test.""" + + max_memory: float + total_allocated_memory: int + allocations: list[AllocationRecord] + + @property + def section(self) -> PytestSection: + """Return a tuple in the format expected by section reporters.""" + total_memory_str = sizeof_fmt(self.total_allocated_memory) + max_memory_str = sizeof_fmt(self.max_memory) + text_lines = [ + f"Test leaked {total_memory_str} out of limit of {max_memory_str}", + "List of leaked allocations: ", + ] + for record in self.allocations: + size = record.size + stack_trace = record.stack_trace() + if not stack_trace: + continue + (function, file, line), *_ = stack_trace + if "pytest_memray" in file: + # Do not report leaks that happen originated from the plugin + continue + text_lines.append(f"\t- {function}:{file}:{line} -> {sizeof_fmt(size)}") + return "memray-leaked-memory", "\n".join(text_lines) + + @property + def long_repr(self) -> str: + """Generate a longrepr user-facing error message.""" + total_memory_str = sizeof_fmt(self.total_allocated_memory) + max_memory_str = sizeof_fmt(self.max_memory) + return ( + f"Test was allowed to leak {max_memory_str} but leaked {total_memory_str}" + ) + + +def limit_memory(limit: str, *, results_file: Path) -> Optional[_MemoryInfo]: """Limit memory used by the test.""" + reader = FileReader(results_file) + allocations = list( + (reader.get_high_watermark_allocation_records(merge_threads=True)) + ) max_memory = parse_memory_string(limit) - total_allocated_memory = sum(record.size for record in _allocations) + total_allocated_memory = sum(record.size for record in allocations) if total_allocated_memory < max_memory: return None - return _MemoryInfo(max_memory, total_allocated_memory, _allocations) + return _MemoryInfo(max_memory, total_allocated_memory, allocations) + + +def check_leaks_runner( + func: Callable[[], Any], + results_file: Path, + *marker_args: Any, + warmups: int = 1, + **marker_kwargs: Any, +) -> Any: + result = None + for _ in range(warmups): + with Tracker("/dev/null"): + result = func() + with Tracker(results_file, trace_python_allocators=True): + result = func() + return result + + +def check_leaks( + limit: str, *, results_file: Path, **kwargs: Any +) -> Optional[_LeakedInfo]: + reader = FileReader(results_file) + allocations = list((reader.get_leaked_allocation_records(merge_threads=True))) + max_memory = parse_memory_string(limit) + total_leaked_memory = sum(record.size for record in allocations) + if total_leaked_memory < max_memory: + return None + return _LeakedInfo(max_memory, total_leaked_memory, allocations) __all__ = [ diff --git a/src/pytest_memray/plugin.py b/src/pytest_memray/plugin.py index ee9fde9..2b75b41 100644 --- a/src/pytest_memray/plugin.py +++ b/src/pytest_memray/plugin.py @@ -11,9 +11,13 @@ from pathlib import Path from tempfile import TemporaryDirectory from typing import Any +from typing import Callable +from typing import Dict from typing import Generator from typing import Iterable from typing import List +from typing import Mapping +from typing import Protocol from typing import Tuple from typing import cast @@ -32,11 +36,27 @@ from pytest import TestReport from pytest import hookimpl +from .marks import check_leaks +from .marks import check_leaks_runner from .marks import limit_memory from .utils import WriteEnabledDirectoryAction from .utils import sizeof_fmt -MARKERS = {"limit_memory": limit_memory} +TestFunction = Callable[[], Any] + + +class Runner(Protocol): + def __call__( + self, func: TestFunction, results_file: Path, *args: Any, **kwargs: Any + ) -> Any: + ... + + +MARKERS: Dict[str, Callable[..., Any]] = { + "limit_memory": limit_memory, + "check_leaks": check_leaks, +} +MARKERS_RUNNERS: Dict[str, Runner] = {"check_leaks": check_leaks_runner} N_TOP_ALLOCS = 5 N_HISTOGRAM_BINS = 5 @@ -84,7 +104,9 @@ class Manager: def __init__(self, config: Config) -> None: self.results: dict[str, Result] = {} self.config = config - path: Path | None = config.getvalue("memray_bin_path") + the_path = config.getvalue("memray_bin_path") + assert the_path is None or isinstance(the_path, Path) + path: Path | None = the_path if path is None: self._tmp_dir: None | TemporaryDirectory[str] = TemporaryDirectory() self.result_path: Path = Path(self._tmp_dir.name) @@ -99,10 +121,28 @@ def pytest_unconfigure(self, config: Config) -> Generator[None, None, None]: if self._tmp_dir is not None: self._tmp_dir.cleanup() + def _default_runner( + self, func: Callable[[], Any], results_file: Path, *args: Any, **kwargs: Any + ) -> Any: + result = None + with Tracker(results_file): + result = func() + return result + @hookimpl(hookwrapper=True) # type: ignore[misc] # Untyped decorator - def pytest_pyfunc_call(self, pyfuncitem: Function) -> object | None: + def pytest_pyfunc_call(self, pyfuncitem: Function) -> Iterable[object | None]: func = pyfuncitem.obj + runner: Runner = self._default_runner + marker_args: Tuple[Any, ...] = tuple() + marker_kwargs: Mapping[str, Any] = dict() + for marker in pyfuncitem.iter_markers(): + if marker.name in MARKERS_RUNNERS: + runner = MARKERS_RUNNERS[marker.name] + marker_args = marker.args + marker_kwargs = marker.kwargs + break + def _build_bin_path() -> Path: if self._tmp_dir is None: of_id = pyfuncitem.nodeid.replace("::", "-") @@ -119,8 +159,10 @@ def _build_bin_path() -> Path: def wrapper(*args: Any, **kwargs: Any) -> object | None: try: result_file = _build_bin_path() - with Tracker(result_file): - result: object | None = func(*args, **kwargs) + runner_function = functools.partial(func, *args, **kwargs) + result: object | None = runner( + runner_function, result_file, *marker_args, **marker_kwargs + ) try: metadata = FileReader(result_file).metadata except OSError: @@ -156,10 +198,9 @@ def pytest_runtest_makereport( result = self.results.get(item.nodeid) if not result: continue - reader = FileReader(result.result_file) - func = reader.get_high_watermark_allocation_records - allocations = list((func(merge_threads=True))) - res = marker_fn(*marker.args, **marker.kwargs, _allocations=allocations) + res = marker_fn( + *marker.args, **marker.kwargs, results_file=result.result_file + ) if res: report.outcome = "failed" report.longrepr = res.long_repr diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py index 01dfd16..fe8ce9e 100644 --- a/tests/test_pytest_memray.py +++ b/tests/test_pytest_memray.py @@ -389,3 +389,68 @@ def test_hello_world(): # called it multiple times per retry. assert mock.call_count == 2 assert result.ret == ExitCode.TESTS_FAILED + + +@pytest.mark.parametrize( + "size, outcome", + [ + (1024 * 1, ExitCode.TESTS_FAILED), + (1024 * 2, ExitCode.TESTS_FAILED), + (1024 * 1 - 1, ExitCode.OK), + (1024 * 0.5, ExitCode.OK), + ], +) +def test_leak_marker(pytester: Pytester, size: int, outcome: ExitCode) -> None: + pytester.makepyfile( + f""" + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + + @pytest.mark.check_leaks("1KB", warmups=1) + def test_memory_alloc_fails(): + allocator.valloc({size}) + # No free call here + """ + ) + + result = pytester.runpytest("--memray") + + assert result.ret == outcome + + +def test_leak_marker_warmups(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + + @pytest.mark.check_leaks("1MB", warmups=10) + def test_the_thing(): + pass + """ + ) + + with patch("pytest_memray.marks.Tracker") as mock: + result = pytester.runpytest("--memray") + + assert result.ret == ExitCode.OK + assert mock.call_count == 10 + 1 + + +def test_leak_marker_does_not_work_if_memray_inactive(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + + @pytest.mark.check_leaks("0B", warmups=1) + def test_memory_alloc_fails(): + allocator.valloc(512) + # No free call here + """ + ) + + result = pytester.runpytest("") + + assert result.ret == ExitCode.OK