Merge pull request #58 from P403n1x87/feat/memory-profiling

feat: memory profiling

README.md

@@ -58,6 +58,7 @@ options:
   -c, --cpu             sample on-CPU stacks only
   -x EXPOSURE, --exposure EXPOSURE
                         exposure time, in seconds
+  -m, --memory          Collect memory allocation events
   -n, --native          sample native stacks
   -o OUTPUT, --output OUTPUT
                         output location (can use %(pid) to insert the process ID)
@@ -101,6 +102,25 @@ You can normally send a `SIGQUIT` signal with the <kbd>CTRL</kbd>+<kbd>\\</kbd>
 key combination.
 
 
+## Memory mode
+
+Besides wall time and CPU time, Echion can be used to profile memory
+allocations. In this mode, Echion tracks the Python memory domain allocators and
+accounts for each single event. Because of the tracing nature, this mode
+introduces considerable overhead, but gives pretty accurate results that can be
+used to investigate potential memory leaks. To fully understand that data that
+is collected in this mode, one should be aware of how Echion tracks allocations
+and deallocations. When an allocation is made, Echion records the frame stack
+that was involved and maps it to the returned memory address. When a
+deallocation for a tracked memory address is made, the freed memory is accounted
+for the same stack. Therefore, objects that are allocated and deallocated during
+the tracking period account for a total of 0 allocated bytes. This means that
+all the non-negative values reported by Echion represent memory that was still
+allocated by the time the tracking ended.
+
+*Since Echion 0.3.0*.
+
+
 ## Why Echion?
 
 Sampling in-process comes with some benefits. One has easier access to more

echion/__main__.py

@@ -114,6 +114,12 @@ def main() -> None:
         help="exposure time, in seconds",
         type=int,
     )
+    parser.add_argument(
+        "-m",
+        "--memory",
+        help="Collect memory allocation events",
+        action="store_true",
+    )
     parser.add_argument(
         "-n",
         "--native",
@@ -165,6 +171,7 @@ def main() -> None:
 
     env["ECHION_INTERVAL"] = str(args.interval)
     env["ECHION_CPU"] = str(int(bool(args.cpu)))
+    env["ECHION_MEMORY"] = str(int(bool(args.memory)))
     env["ECHION_NATIVE"] = str(int(bool(args.native)))
     env["ECHION_OUTPUT"] = args.output.replace("%%(pid)", str(os.getpid()))
     env["ECHION_STEALTH"] = str(int(bool(args.stealth)))

echion/bootstrap/__init__.py

@@ -32,6 +32,7 @@ def start():
     # Set the configuration
     ec.set_interval(int(os.getenv("ECHION_INTERVAL", 1000)))
     ec.set_cpu(bool(int(os.getenv("ECHION_CPU", 0))))
+    ec.set_memory(bool(int(os.getenv("ECHION_MEMORY", 0))))
     ec.set_native(bool(int(os.getenv("ECHION_NATIVE", 0))))
     ec.set_where(bool(int(os.getenv("ECHION_WHERE", 0) or 0)))
 

echion/config.h

@@ -17,8 +17,8 @@ static unsigned int interval = 1000;
 // CPU Time mode
 static int cpu = 0;
 
-// Output stream
-static std::ofstream output;
+// Memory events
+static int memory = 0;
 
 // Native stack sampling
 static int native = 0;
@@ -55,6 +55,19 @@ set_cpu(PyObject *Py_UNUSED(m), PyObject *args)
     Py_RETURN_NONE;
 }
 
+// ----------------------------------------------------------------------------
+static PyObject *
+set_memory(PyObject *Py_UNUSED(m), PyObject *args)
+{
+    int new_memory;
+    if (!PyArg_ParseTuple(args, "p", &new_memory))
+        return NULL;
+
+    memory = new_memory;
+
+    Py_RETURN_NONE;
+}
+
 // ----------------------------------------------------------------------------
 static PyObject *
 set_native(PyObject *Py_UNUSED(m), PyObject *args)

echion/core.pyi

@@ -21,6 +21,7 @@ def init_asyncio(
 # Configuration interface
 def set_interval(interval: int) -> None: ...
 def set_cpu(cpu: bool) -> None: ...
+def set_memory(memory: bool) -> None: ...
 def set_native(native: bool) -> None: ...
 def set_where(where: bool) -> None: ...
 def set_pipe_name(name: str) -> None: ...

echion/coremodule.cc

@@ -28,6 +28,8 @@
 
 #include <echion/config.h>
 #include <echion/interp.h>
+#include <echion/memory.h>
+#include <echion/mojo.h>
 #include <echion/signals.h>
 #include <echion/stacks.h>
 #include <echion/state.h>
@@ -104,18 +106,72 @@ _start()
 {
     init_frame_cache(MAX_FRAMES * (1 + native));
 
+    try
+    {
+        mojo.open();
+    }
+    catch (MojoWriter::Error &)
+    {
+        return;
+    }
+
     install_signals();
 
 #if defined PL_DARWIN
     // Get the wall time clock resource.
     host_get_clock_service(mach_host_self(), CALENDAR_CLOCK, &cclock);
 #endif
+
+    if (where)
+    {
+        std::ofstream pipe(pipe_name, std::ios::out);
+
+        if (pipe)
+            do_where(pipe);
+
+        else
+            std::cerr << "Failed to open pipe " << pipe_name << std::endl;
+
+        running = 0;
+
+        return;
+    }
+
+    setup_where();
+
+    mojo.header();
+
+    if (memory)
+    {
+        mojo.metadata("mode", "memory");
+    }
+    else
+    {
+        mojo.metadata("mode", (cpu ? "cpu" : "wall"));
+    }
+    mojo.metadata("interval", std::to_string(interval));
+    mojo.metadata("sampler", "echion");
+
+    // DEV: Workaround for the austin-python library: we send an empty sample
+    // to set the PID. We also map the key value 0 to the empty string, to
+    // support task name frames.
+    mojo.stack(pid, 0, "MainThread");
+    mojo.string(0, "");
+    mojo.string(1, "<invalid>");
+    mojo.string(2, "<unknown>");
+    mojo.metric_time(0);
+
+    if (memory)
+        setup_memory();
 }
 
 // ----------------------------------------------------------------------------
 static inline void
 _stop()
 {
+    if (memory)
+        teardown_memory();
+
     // Clean up the thread info map. When not running async, we need to guard
     // the map lock because we are not in control of the sampling thread.
     {
@@ -125,12 +181,16 @@ _stop()
         string_table.clear();
     }
 
+    teardown_where();
+
 #if defined PL_DARWIN
     mach_port_deallocate(mach_task_self(), cclock);
 #endif
 
     restore_signals();
 
+    mojo.close();
+
     reset_frame_cache();
 }
 
@@ -142,71 +202,37 @@ _sampler()
     // hold:
     // 1. The interpreter state object lives as long as the process itself.
 
-    if (where)
-    {
-        std::ofstream pipe(pipe_name, std::ios::out);
-
-        if (pipe)
-            do_where(pipe);
-
-        else
-            std::cerr << "Failed to open pipe " << pipe_name << std::endl;
-
-        running = 0;
-
-        return;
-    }
-
-    setup_where();
-
     last_time = gettime();
 
-    output.open(std::getenv("ECHION_OUTPUT"));
-    if (!output.is_open())
-    {
-        std::cerr << "Failed to open output file " << std::getenv("ECHION_OUTPUT") << std::endl;
-        return;
-    }
-
-    mojo_header();
-
-    mojo_metadata("mode", (cpu ? "cpu" : "wall"));
-    mojo_metadata("interval", interval);
-    mojo_metadata("sampler", "echion");
-
-    // DEV: Workaround for the austin-python library: we send an empty sample
-    // to set the PID. We also map the key value 0 to the empty string, to
-    // support task name frames.
-    mojo_stack(pid, 0, "");
-    mojo_string_event(0, "");
-    mojo_string_event(1, "<invalid>");
-    mojo_string_event(2, "<unknown>");
-    mojo_metric_time(0);
-
     while (running)
     {
         microsecond_t now = gettime();
         microsecond_t end_time = now + interval;
-        microsecond_t wall_time = now - last_time;
 
-        for_each_interp(
-            [=](PyInterpreterState *interp) -> void
-            {
-                for_each_thread(
-                    interp,
-                    [=](PyThreadState *tstate, ThreadInfo &thread)
-                    { thread.sample(interp->id, tstate, wall_time); });
-            });
+        if (memory)
+        {
+            if (rss_tracker.check())
+                stack_stats.flush();
+        }
+        else
+        {
+            microsecond_t wall_time = now - last_time;
+
+            for_each_interp(
+                [=](PyInterpreterState *interp) -> void
+                {
+                    for_each_thread(
+                        interp,
+                        [=](PyThreadState *tstate, ThreadInfo &thread)
+                        { thread.sample(interp->id, tstate, wall_time); });
+                });
+        }
 
         while (gettime() < end_time && running)
             sched_yield();
 
         last_time = now;
     }
-
-    output.close();
-
-    teardown_where();
 }
 
 static void
@@ -400,6 +426,7 @@ static PyMethodDef echion_core_methods[] = {
     // Configuration interface
     {"set_interval", set_interval, METH_VARARGS, "Set the sampling interval"},
     {"set_cpu", set_cpu, METH_VARARGS, "Set whether to use CPU time instead of wall time"},
+    {"set_memory", set_memory, METH_VARARGS, "Set whether to sample memory usage"},
     {"set_native", set_native, METH_VARARGS, "Set whether to sample the native stacks"},
     {"set_where", set_where, METH_VARARGS, "Set whether to use where mode"},
     {"set_pipe_name", set_pipe_name, METH_VARARGS, "Set the pipe name"},