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.

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
@@ -99,10 +119,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:
         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 +157,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 +196,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