diff --git a/docs/api.rst b/docs/api.rst new file mode 100644 index 0000000..c6c9790 --- /dev/null +++ b/docs/api.rst @@ -0,0 +1,18 @@ +.. module:: pytest_memray + +pytest-memray API +================= + +Types +----- + +.. autoclass:: LeaksFilterFunction() + :members: __call__ + :show-inheritance: + +.. autoclass:: Stack() + :members: + +.. autoclass:: StackFrame() + :members: + diff --git a/docs/conf.py b/docs/conf.py index b0ff132..4183367 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -9,8 +9,10 @@ from sphinxcontrib.programoutput import Command extensions = [ + "sphinx.ext.autodoc", "sphinx.ext.extlinks", "sphinx.ext.githubpages", + "sphinx.ext.intersphinx", "sphinxarg.ext", "sphinx_inline_tabs", "sphinxcontrib.programoutput", @@ -36,6 +38,15 @@ "https://github.com/bloomberg/pytest-memray/issues/.*": "https://github.com/bloomberg/pytest-memray/pull/.*" } +# Try to resolve Sphinx references as Python objects by default. This means we +# don't need :func: or :class: etc, which keep docstrings more human readable. +default_role = "py:obj" + +# Automatically link to Python standard library types. +intersphinx_mapping = { + "python": ("https://docs.python.org/3", None), +} + def _get_output(self): code, out = prev(self) diff --git a/docs/index.rst b/docs/index.rst index 3f088d8..759a43b 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -13,4 +13,5 @@ reports like: usage configuration + api news diff --git a/docs/usage.rst b/docs/usage.rst index a409335..4d2e05a 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -31,33 +31,115 @@ reported after tests run ends: Markers ~~~~~~~ -This plugin provides markers that can be used to enforce additional checks and -validations on tests when this plugin is enabled. +This plugin provides `markers `__ +that can be used to enforce additional checks and validations on tests. -.. important:: These markers do nothing when the plugin is not enabled. +.. py:function:: pytest.mark.limit_memory(memory_limit: str) -``limit_memory`` ----------------- + 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. + 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. -The format for the string is `` ([KMGTP]B|B)``. The marker will raise -``ValueError`` if the string format cannot be parsed correctly. + The format for the string is `` ([KMGTP]B|B)``. The marker will raise + ``ValueError`` if the string format cannot be parsed correctly. -.. warning:: + .. warning:: - As the Python interpreter has its own - `object allocator `__ is possible - that memory is not immediately released to the system when objects are deleted, so - tests using this marker may need to give some room to account for this. + As the Python interpreter has its own + `object allocator `__ it's possible + that memory is not immediately released to the system when objects are deleted, + so tests using this marker may need to give some room to account for this. -Example of usage: + Example of usage: -.. code-block:: python + .. code-block:: python - @pytest.mark.limit_memory("24 MB") - def test_foobar(): - pass # do some stuff that allocates memory + @pytest.mark.limit_memory("24 MB") + def test_foobar(): + pass # do some stuff that allocates memory + + +.. py:function:: pytest.mark.limit_leaks(location_limit: str, filter_fn: LeaksFilterFunction | None = None) + + Fail the execution of the test if any call stack in the test leaks more memory than + allowed. + + .. important:: + To detect leaks, Memray needs to intercept calls to the Python allocators and + report native call frames. This is adds significant overhead, and will slow your + test down. + + When this marker is applied to a test, the plugin will analyze the memory + allocations that are made while the test body runs and not freed by the time the + test body function returns. It groups them by the call stack leading to the + allocation, and sums the amount leaked by each **distinct call stack**. If the total + amount leaked from any particular call stack is greater than the configured limit, + the test will fail. + + .. important:: + It's recommended to run your API or code in a loop when utilizing this plugin. + This practice helps in distinguishing genuine leaks from the "noise" generated + by internal caches and other incidental allocations. + + 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-only argument ``filter_fn``. This argument + represents a filtering function that will be called once for each distinct call + stack that leaked more memory than allowed. If it returns *True*, leaks from that + location will be included in the final report. If it returns *False*, leaks + associated with the stack it was called with will be ignored. If all leaks are + ignored, the test will not fail. This can be used to discard any known false + positives. + + .. tip:: + + You can pass the ``--memray-bin-path`` argument to ``pytest`` to specify + a directory where Memray will store the binary files with the results. You + can then use the ``memray`` CLI to further investigate the allocations and the + leaks using any Memray reporters you'd like. Check `the memray docs + `_ for more + information. + + Example of usage: + + .. code-block:: python + + @pytest.mark.limit_leaks("1 MB") + def test_foobar(): + # Run the function we're testing in a loop to ensure + # we can differentiate leaks from memory held by + # caches inside the Python interpreter. + for _ in range(100): + do_some_stuff() + + .. warning:: + It is **very** challenging to write tests that do not "leak" memory in some way, + due to circumstances beyond your control. + + There are many caches inside the Python interpreter itself. Just a few examples: + + - The `re` module caches compiled regexes. + - The `logging` module caches whether a given log level is active for + a particular logger the first time you try to log something at that level. + - A limited number of objects of certain heavily used types are cached for reuse + so that `object.__new__` does not always need to allocate memory. + - The mapping from bytecode index to line number for each Python function is + cached when it is first needed. + + There are many more such caches. Also, within pytest, any message that you log or + print is captured, so that it can be included in the output if the test fails. + + Memray sees these all as "leaks", because something was allocated while the test + ran and it was not freed by the time the test body finished. We don't know that + it's due to an implementation detail of the interpreter or pytest that the memory + wasn't freed. Morever, because these caches are implementation details, the + amount of memory allocated, the call stack of the allocation, and even the + allocator that was used can all change from one version to another. + + Because of this, you will almost certainly need to allow some small amount of + leaked memory per call stack, or use the ``filter_fn`` argument to filter out + false-positive leak reports based on the call stack they're associated with. diff --git a/src/pytest_memray/__init__.py b/src/pytest_memray/__init__.py index 052834c..1b96d5b 100644 --- a/src/pytest_memray/__init__.py +++ b/src/pytest_memray/__init__.py @@ -1,7 +1,13 @@ from __future__ import annotations from ._version import __version__ as __version__ +from .marks import LeaksFilterFunction +from .marks import Stack +from .marks import StackFrame __all__ = [ "__version__", + "LeaksFilterFunction", + "Stack", + "StackFrame", ] diff --git a/src/pytest_memray/marks.py b/src/pytest_memray/marks.py index 2faa30d..ece673f 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 Iterable +from typing import Optional +from typing import Protocol from typing import Tuple from typing import cast from memray import AllocationRecord +from memray import FileReader from pytest import Config from .utils import parse_memory_string @@ -14,12 +19,90 @@ PytestSection = Tuple[str, str] +@dataclass +class StackFrame: + """One frame of a call stack. + + Each frame has attributes to tell you what code was executing. + """ + + function: str + """The function being executed, or ``"???"`` if unknown.""" + + filename: str + """The source file being executed, or ``"???"`` if unknown.""" + + lineno: int + """The line number of the executing line, or ``0`` if unknown.""" + + +@dataclass +class Stack: + """The call stack that led to some memory allocation. + + You can inspect the frames which make up the call stack. + """ + + frames: Tuple[StackFrame, ...] + """The frames that make up the call stack, most recent first.""" + + +class LeaksFilterFunction(Protocol): + """A callable that can decide whether to ignore some memory leaks. + + This can be used to suppress leak reports from locations that are known to + leak. For instance, you might know that objects of a certain type are + cached by the code you're invoking, and so you might want to ignore all + reports of leaked memory allocated below that type's constructor. + + You can provide any callable with the following signature as the + ``filter_fn`` keyword argument for the `.limit_leaks` marker: + """ + + def __call__(self, stack: Stack) -> bool: + """Return whether allocations from this stack should be reported. + + Return ``True`` if you want the leak to be reported, or ``False`` if + you want it to be suppressed. + """ + ... + + @dataclass class _MemoryInfo: - """Type that holds all memray-related info for a failed test.""" + """Type that holds memory-related info for a failed test.""" max_memory: float + allocations: list[AllocationRecord] + num_stacks: int + native_stacks: bool total_allocated_memory: int + + @property + def section(self) -> PytestSection: + """Return a tuple in the format expected by section reporters.""" + body = _generate_section_text( + self.allocations, self.native_stacks, self.num_stacks + ) + return ( + "memray-max-memory", + "List of allocations:\n" + body, + ) + + @property + def long_repr(self) -> str: + """Generate a longrepr user-facing error message.""" + return ( + f"Test was limited to {sizeof_fmt(self.max_memory)} " + f"but allocated {sizeof_fmt(self.total_allocated_memory)}" + ) + + +@dataclass +class _LeakedInfo: + """Type that holds leaked memory-related info for a failed test.""" + + max_memory: float allocations: list[AllocationRecord] num_stacks: int native_stacks: bool @@ -27,55 +110,119 @@ class _MemoryInfo: @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 is using {total_memory_str} out of limit of {max_memory_str}", - "List of allocations: ", - ] - for record in self.allocations: - size = record.size - stack_trace = ( - record.hybrid_stack_trace() - if self.native_stacks - else record.stack_trace() - ) - if not stack_trace: - continue - padding = " " * 4 - text_lines.append(f"{padding}- {sizeof_fmt(size)} allocated here:") - stacks_left = self.num_stacks - for function, file, line in stack_trace: - if stacks_left <= 0: - break - text_lines.append(f"{padding*2}{function}:{file}:{line}") - stacks_left -= 1 - - return "memray-max-memory", "\n".join(text_lines) + body = _generate_section_text( + self.allocations, self.native_stacks, self.num_stacks + ) + return ( + "memray-leaked-memory", + "List of leaked allocations:\n" + body, + ) @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 limited to {max_memory_str} but allocated {total_memory_str}" + return ( + f"Test was allowed to leak {sizeof_fmt(self.max_memory)} " + "per location but at least one location leaked more" + ) + + +def _generate_section_text( + allocations: list[AllocationRecord], native_stacks: bool, num_stacks: int +) -> str: + text_lines = [] + for record in allocations: + size = record.size + stack_trace = ( + record.hybrid_stack_trace() if native_stacks else record.stack_trace() + ) + if not stack_trace: + continue + padding = " " * 4 + text_lines.append(f"{padding}- {sizeof_fmt(size)} allocated here:") + stacks_left = num_stacks + for function, file, line in stack_trace: + if stacks_left <= 0: + text_lines.append(f"{padding*2}...") + break + text_lines.append(f"{padding*2}{function}:{file}:{line}") + stacks_left -= 1 + + return "\n".join(text_lines) + + +def _passes_filter( + stack: Iterable[Tuple[str, str, int]], filter_fn: Optional[LeaksFilterFunction] +) -> bool: + if filter_fn is None: + return True + + frames = tuple(StackFrame(*frame) for frame in stack) + return filter_fn(Stack(frames)) def limit_memory( - limit: str, *, _allocations: list[AllocationRecord], _config: Config + limit: str, *, _result_file: Path, _config: Config ) -> _MemoryInfo | None: """Limit memory used by the test.""" + reader = FileReader(_result_file) + allocations: list[AllocationRecord] = 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 num_stacks: int = cast(int, value_or_ini(_config, "stacks")) native_stacks: bool = cast(bool, value_or_ini(_config, "native")) return _MemoryInfo( - max_memory, total_allocated_memory, _allocations, num_stacks, native_stacks + max_memory=max_memory, + allocations=allocations, + num_stacks=num_stacks, + native_stacks=native_stacks, + total_allocated_memory=total_allocated_memory, + ) + + +def limit_leaks( + location_limit: str, + *, + filter_fn: Optional[LeaksFilterFunction] = None, + _result_file: Path, + _config: Config, +) -> _LeakedInfo | None: + reader = FileReader(_result_file) + allocations: list[AllocationRecord] = list( + reader.get_leaked_allocation_records(merge_threads=True) + ) + + memory_limit = parse_memory_string(location_limit) + + leaked_allocations = list( + allocation + for allocation in allocations + if ( + allocation.size >= memory_limit + and _passes_filter(allocation.hybrid_stack_trace(), filter_fn) + ) + ) + + if not leaked_allocations: + return None + + num_stacks: int = max(cast(int, value_or_ini(_config, "stacks")), 5) + return _LeakedInfo( + max_memory=memory_limit, + allocations=leaked_allocations, + num_stacks=num_stacks, + native_stacks=True, ) __all__ = [ "limit_memory", + "limit_leaks", + "LeaksFilterFunction", + "Stack", + "StackFrame", ] diff --git a/src/pytest_memray/plugin.py b/src/pytest_memray/plugin.py index 1933388..c9804ca 100644 --- a/src/pytest_memray/plugin.py +++ b/src/pytest_memray/plugin.py @@ -17,6 +17,7 @@ from typing import List from typing import Tuple from typing import cast +from typing import Protocol from _pytest.terminal import TerminalReporter from memray import AllocationRecord @@ -34,12 +35,32 @@ from pytest import hookimpl from .marks import limit_memory +from .marks import limit_leaks from .utils import WriteEnabledDirectoryAction from .utils import positive_int from .utils import sizeof_fmt from .utils import value_or_ini -MARKERS = {"limit_memory": limit_memory} + +class SectionMetadata(Protocol): + long_repr: str + section: Tuple[str, str] + + +class PluginFn(Protocol): + def __call__( + *args: Any, + _result_file: Path, + _config: Config, + **kwargs: Any, + ) -> SectionMetadata | None: + ... + + +MARKERS = { + "limit_memory": limit_memory, + "limit_leaks": limit_leaks, +} N_TOP_ALLOCS = 5 N_HISTOGRAM_BINS = 5 @@ -134,6 +155,9 @@ def pytest_pyfunc_call(self, pyfuncitem: Function) -> object | None: yield return + if len(markers) > 1: + raise ValueError("Only one Memray marker can be applied to each test") + def _build_bin_path() -> Path: if self._tmp_dir is None and not os.getenv("MEMRAY_RESULT_PATH"): of_id = pyfuncitem.nodeid.replace("::", "-") @@ -151,6 +175,9 @@ def _build_bin_path() -> Path: value_or_ini(self.config, "trace_python_allocators") ) + if markers and "limit_leaks" in markers: + native = trace_python_allocators = True + @functools.wraps(func) def wrapper(*args: Any, **kwargs: Any) -> object | None: test_result: object | Any = None @@ -198,19 +225,17 @@ def pytest_runtest_makereport( return None for marker in item.iter_markers(): - marker_fn = MARKERS.get(marker.name) - if not marker_fn: + maybe_marker_fn = MARKERS.get(marker.name) + if not maybe_marker_fn: continue + marker_fn: PluginFn = cast(PluginFn, maybe_marker_fn) 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, + _result_file=result.result_file, _config=self.config, ) if res: diff --git a/src/pytest_memray/py.typed b/src/pytest_memray/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py index dd19afb..0ae0ded 100644 --- a/tests/test_pytest_memray.py +++ b/tests/test_pytest_memray.py @@ -594,3 +594,151 @@ def test_memory_alloc_fails(): ) result = pytester.runpytest("-Werror", "--memray") assert result.ret == ExitCode.OK + + +@pytest.mark.parametrize( + "size, outcome", + [ + (1, ExitCode.OK), + (1024 * 1 / 10, ExitCode.OK), + (1024 * 1, ExitCode.TESTS_FAILED), + (1024 * 10, ExitCode.TESTS_FAILED), + ], +) +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.limit_leaks("5KB") + def test_memory_alloc_fails(): + for _ in range(10): + allocator.valloc({size}) + # No free call here + """ + ) + + result = pytester.runpytest("--memray") + + assert result.ret == outcome + + +@pytest.mark.parametrize( + "size, outcome", + [ + (1, ExitCode.OK), + (1024 * 1 / 10, ExitCode.OK), + (1024 * 1, ExitCode.TESTS_FAILED), + (1024 * 10, ExitCode.TESTS_FAILED), + ], +) +def test_leak_marker_in_a_thread( + pytester: Pytester, size: int, outcome: ExitCode +) -> None: + pytester.makepyfile( + f""" + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + import threading + def allocating_func(): + for _ in range(10): + allocator.valloc({size}) + # No free call here + @pytest.mark.limit_leaks("5KB") + def test_memory_alloc_fails(): + t = threading.Thread(target=allocating_func) + t.start() + t.join() + """ + ) + + result = pytester.runpytest("--memray") + assert result.ret == outcome + + +def test_leak_marker_filtering_function(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + LEAK_SIZE = 1024 + allocator = MemoryAllocator() + + def this_should_not_be_there(): + allocator.valloc(LEAK_SIZE) + # No free call here + + def filtering_function(stack): + for frame in stack.frames: + if frame.function == "this_should_not_be_there": + return False + return True + + @pytest.mark.limit_leaks("5KB", filter_fn=filtering_function) + def test_memory_alloc_fails(): + for _ in range(10): + this_should_not_be_there() + """ + ) + + result = pytester.runpytest("--memray") + + assert result.ret == ExitCode.OK + + +def test_leak_marker_does_work_if_memray_not_passed(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + @pytest.mark.limit_leaks("0B") + def test_memory_alloc_fails(): + allocator.valloc(512) + # No free call here + """ + ) + + result = pytester.runpytest() + + assert result.ret == ExitCode.TESTS_FAILED + + +def test_multiple_markers_are_not_supported(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + @pytest.mark.limit_leaks("0MB") + @pytest.mark.limit_memory("0MB") + def test_bar(): + pass + """ + ) + + result = pytester.runpytest("--memray") + assert result.ret == ExitCode.TESTS_FAILED + + output = result.stdout.str() + assert "Only one Memray marker can be applied to each test" in output + + +def test_multiple_markers_are_not_supported_with_global_marker( + pytester: Pytester, +) -> None: + pytester.makepyfile( + """ + import pytest + pytestmark = pytest.mark.limit_memory("1 MB") + @pytest.mark.limit_leaks("0MB") + def test_bar(): + pass + """ + ) + + result = pytester.runpytest("--memray") + assert result.ret == ExitCode.TESTS_FAILED + + output = result.stdout.str() + assert "Only one Memray marker can be applied to each test" in output