Allow memray stats to output a JSON report

.github/workflows/build_wheels.yml

@@ -40,7 +40,7 @@ jobs:
           CIBW_ARCHS_LINUX: ${{ matrix.cibw_arch }}
           CIBW_PRERELEASE_PYTHONS: True
           CIBW_TEST_EXTRAS: test
-          CIBW_TEST_COMMAND: pytest {package}/tests
+          CIBW_TEST_COMMAND: python -m pytest {package}/tests
       - uses: actions/upload-artifact@v3
         with:
           path: ./wheelhouse/*.whl

docs/examples/mandelbrot/mandelbrot-threaded.py

@@ -0,0 +1,41 @@
+from threading import Thread
+
+import numpy as np
+
+
+def mandelbrot(height, width, x=-0.5, y=0, zoom=1, max_iterations=100):
+    # To make navigation easier we calculate these values
+    x_width = 1.5
+    y_height = 1.5 * height / width
+    x_from = x - x_width / zoom
+    x_to = x + x_width / zoom
+    y_from = y - y_height / zoom
+    y_to = y + y_height / zoom
+    # Here the actual algorithm starts
+    x = np.linspace(x_from, x_to, width).reshape((1, width))
+    y = np.linspace(y_from, y_to, height).reshape((height, 1))
+    c = x + 1j * y
+    # Initialize z to all zero
+    z = np.zeros(c.shape, dtype=np.complex128)
+    # To keep track in which iteration the point diverged
+    div_time = np.zeros(z.shape, dtype=int)
+    # To keep track on which points did not converge so far
+    m = np.full(c.shape, True, dtype=bool)
+    for i in range(max_iterations):
+        z[m] = z[m] ** 2 + c[m]
+        diverged = np.greater(
+            np.abs(z), 2, out=np.full(c.shape, False), where=m
+        )  # Find diverging
+        div_time[diverged] = i  # set the value of the diverged iteration number
+        m[np.abs(z) > 2] = False  # to remember which have diverged
+    return div_time
+
+
+if __name__ == "__main__":
+    t1 = Thread(target=mandelbrot, args=(800, 1000))
+    t1.start()
+    t2 = Thread(target=mandelbrot, args=(800, 1000))
+    t2.start()
+    t3 = Thread(target=mandelbrot, args=(800, 1000))
+    t3.start()
+    t3.join()

docs/flamegraph.rst

@@ -186,8 +186,8 @@ Non-relevant frame hiding
 The flame graph exposes a button to show or hide frames which might be
 distracting when interpreting the results, either because they were
 injected by Memray or because they are low-level implementation
-details of CPython. By default, frames tagged as non-relevant are
-hidden. You can reveal them with the "Show Non-Relevant Frames"
+details of CPython. By default, frames tagged as irrelevant are
+hidden. You can reveal them by unchecking the *Hide Irrelevant Frames*
 checkbox:
 
 .. image:: _static/images/non_relevant_checkbox.png
@@ -196,6 +196,21 @@ checkbox:
 Note that allocations in these frames will still be accounted for
 in parent frames, even if they're hidden.
 
+Inverted View
+-----------------
+
+Although the flame graphs explained above show the calling functions below
+and memory allocating functions above, flame graphs can be inverted
+so that the calling functions are at the top, while memory allocating
+functions are at the bottom. In this view, look for wide ceilings
+instead of wide plateaus to find functions with the largest allocation
+of memory.
+
+To invert the flame graph, press the *Invert* button:
+
+.. image:: _static/images/invert_button.png
+    :align: center
+
 .. _memory-leaks-view:
 
 Memory Leaks View
@@ -238,7 +253,7 @@ select a specific thread to display a flame graph for that one thread:
 To go back to the merged view, the "Reset" entry can be used in the
 dropdown menu.
 
-Note that the root node (displayed as **memray**) is always present
+Note that the root node (displayed as **<root>**) is always present
 and is displayed as thread 0.
 
 Conclusion

docs/live.rst

@@ -61,6 +61,16 @@ program by pressing the left and right arrow keys.
 
 .. image:: _static/images/live_different_thread.png
 
+Using with native tracking
+--------------------------
+
+It is possible to use :ref:`native tracking` along with the live mode. This can be achieved by passing ``--native``
+to the ``run`` command.
+
+.. code:: shell-session
+
+  $ memray run --live --native application.py
+
 Remote mode
 -----------
 
@@ -104,21 +114,10 @@ It is possible to make ``run --live-remote`` start the server on a user-specifie
   ``run`` command *before* your script/module. Otherwise, they will be treated as arguments for the script and will not
   be used by Memray.
 
-  For example, the following invocation will pass ``--live-remote`` and ``--live-port 12345`` to ``application.py``,
-  instead of having them be used by ``memray run``:
+  For example, the following invocation will pass ``--live-port 12345`` to ``application.py``,
+  instead of having it be used by ``memray run``:
 
   .. code:: shell-session
 
     $ memray run --live-remote application.py --live-port 12345
     Run 'memray live 60125' in another shell to see live results
-
-Using with native tracking
---------------------------
-
-It is possible to use :ref:`native tracking` along with the live mode. This can be achieved by passing ``--native``
-to the ``run`` command.
-
-.. code:: shell-session
-
-  $ memray run --live --native application.py
-  Run 'memray live 60125' in another shell to see live results

docs/run.rst

@@ -150,11 +150,17 @@ use ``run --live-remote``:
 
   memray3.9 run --live-remote application.py
 
-In this mode it will choose an unused port and bind to it, waiting for you to run:
+In this mode, Memray will choose an unused port, bind to it, and display a message saying:
+
+.. code:: text
+
+  Run 'memray live <port>' in another shell to see live results
+
+It will wait for you to run:
 
 .. code:: shell
 
-   memray3.9 live $port
+   memray3.9 live <port>
 
 in another terminal window to attach to it. Regardless of whether you choose to use one terminal or two, the resulting
 TUI is exactly the same. See :doc:`live` for details on how to interpret and control the TUI.

docs/stats.rst

@@ -20,6 +20,8 @@ The output includes the following:
 
 * Stack trace and **count** of the top 'n' largest allocating locations by number of allocations (*default: 5*, configurable with the ``-n`` command line param)
 
+* (for JSON output only) Metadata about the tracked process
+
 Basic Usage
 -----------
 
@@ -34,6 +36,18 @@ previously generated using :doc:`the run subcommand <run>`.
 
 The output will be printed directly to the standard output of the terminal.
 
+JSON Output
+-----------
+
+If you supply the ``--json`` flag, the ``stats`` subcommand will write its
+output to a JSON file, rather than to the terminal. Like other commands that
+output to files, the default output file name is based on the name of your
+capture file, but it can be overridden with the ``-o`` / ``--output`` option.
+By default Memray will refuse to overwrite an existing file, but you can force
+it to by supplying the ``-f`` / ``--force`` option.
+
+Note that new fields may be added to the JSON output over time, though we'll
+try to avoid removing existing fields.
 
 CLI Reference
 -------------

news/377.feature.rst

@@ -0,0 +1 @@
+Allow ``memray stats`` to output a JSON report via ``--json`` flag.

news/379.feature.rst

@@ -0,0 +1 @@
+We now publish x86-64 musllinux_1_1 wheels, compatible with Alpine Linux.

pyproject.toml

@@ -58,7 +58,7 @@ exclude="tests/integration/(native_extension|multithreaded_extension)/"
 
 [tool.cibuildwheel]
 build = ["cp38-*", "cp39-*", "cp310-*", "cp311-*"]
-skip = "*musllinux*"
+skip = "*musllinux*i686*"
 manylinux-x86_64-image = "manylinux2014"
 manylinux-i686-image = "manylinux2014"
 
@@ -102,3 +102,11 @@ omit = [
 [tool.coverage.report]
 skip_covered = true
 show_missing = true
+
+
+# Override the default linux before-all for musl linux
+[[tool.cibuildwheel.overrides]]
+select = "*-musllinux*"
+before-all = [
+  "apk add --update libunwind-dev lz4-dev gdb lldb",
+]

src/memray/_memray/hooks.cpp

@@ -190,7 +190,11 @@ realloc(void* ptr, size_t size) noexcept
 {
     assert(hooks::realloc);
 
-    void* ret = hooks::realloc(ptr, size);
+    void* ret;
+    {
+        tracking_api::RecursionGuard guard;
+        ret = hooks::realloc(ptr, size);
+    }
     if (ret) {
         if (ptr != nullptr) {
             tracking_api::Tracker::trackDeallocation(ptr, 0, hooks::Allocator::FREE);
@@ -280,8 +284,11 @@ void*
 dlopen(const char* filename, int flag) noexcept
 {
     assert(hooks::dlopen);
-
-    void* ret = hooks::dlopen(filename, flag);
+    void* ret;
+    {
+        tracking_api::RecursionGuard guard;
+        ret = hooks::dlopen(filename, flag);
+    }
     if (ret) {
         tracking_api::Tracker::invalidate_module_cache();
         if (filename && nullptr != strstr(filename, "/_greenlet.")) {
@@ -307,7 +314,11 @@ aligned_alloc(size_t alignment, size_t size) noexcept
 {
     assert(hooks::aligned_alloc);
 
-    void* ret = hooks::aligned_alloc(alignment, size);
+    void* ret;
+    {
+        tracking_api::RecursionGuard guard;
+        ret = hooks::aligned_alloc(alignment, size);
+    }
     if (ret) {
         tracking_api::Tracker::trackAllocation(ret, size, hooks::Allocator::ALIGNED_ALLOC);
     }

src/memray/_stats.pyi

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
 
-from ._memray import AllocatorType
 from ._memray import PythonStackElement
 from ._metadata import Metadata
 
@@ -11,6 +10,6 @@ class Stats:
     total_memory_allocated: int
     peak_memory_allocated: int
     allocation_count_by_size: dict[int, int]
-    allocation_count_by_allocator: dict[AllocatorType, int]
+    allocation_count_by_allocator: dict[str, int]
     top_locations_by_size: list[tuple[PythonStackElement, int]]
     top_locations_by_count: list[tuple[PythonStackElement, int]]

src/memray/commands/stats.py

@@ -1,6 +1,7 @@
 import argparse
 import os
 from pathlib import Path
+from typing import Optional
 
 from memray._errors import MemrayCommandError
 from memray._memray import compute_statistics
@@ -33,6 +34,26 @@ def valid_positive_int(value: str) -> int:
             default=5,
         )
 
+        parser.add_argument(
+            "--json",
+            help="Exports stats to a JSON file",
+            action="store_true",
+            default=False,
+        )
+        parser.add_argument(
+            "-o",
+            "--output",
+            help="Output file name for JSON output",
+            default=None,
+        )
+        parser.add_argument(
+            "-f",
+            "--force",
+            help="If the JSON output file already exists, overwrite it",
+            action="store_true",
+            default=False,
+        )
+
     def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
         result_path = Path(args.results)
         if not result_path.exists() or not result_path.is_file():
@@ -49,5 +70,24 @@ def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None
                 exit_code=1,
             )
 
+        json_output_file: Optional[Path] = None
+        if args.json:
+            if args.output:
+                json_output_file = Path(args.output)
+            else:
+                filename = str(result_path.name) + ".json"
+                if filename.startswith("memray-"):
+                    filename = filename[len("memray-") :]
+                filename = "memray-stats-" + filename
+                json_output_file = result_path.with_name(filename)
+
+            if not args.force and json_output_file.exists():
+                raise MemrayCommandError(
+                    f"File already exists, will not overwrite: {json_output_file}",
+                    exit_code=1,
+                )
+
         reporter = StatsReporter(stats, args.num_largest)
-        reporter.render()
+        reporter.render(json_output_file=json_output_file)
+        if json_output_file is not None:
+            print(f"Wrote {json_output_file}")