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 <[email protected]>

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 ``<NUMBER> ([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
+<https://bloomberg.github.io/memray/getting_started.html>`_ 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
+

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__ = [

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) -> Generator[None, None, 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

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