Add a test to exercise asyncio stack traces in out-of-process profilers

Signed-off-by: Pablo Galindo <[email protected]>

Author: pablogsal

Include/internal/pycore_runtime.h

@@ -108,6 +108,7 @@ typedef struct _Py_DebugOffsets {
         uint64_t instr_ptr;
         uint64_t localsplus;
         uint64_t owner;
+        uint64_t stackpointer;
     } interpreter_frame;
 
     // Code object offset;
@@ -152,6 +153,13 @@ typedef struct _Py_DebugOffsets {
         uint64_t ob_size;
     } list_object;
 
+    // PySet object offset;
+    struct _set_object {
+        uint64_t size;
+        uint64_t used;
+        uint64_t table;
+    } set_object;
+
     // PyDict object offset;
     struct _dict_object {
         uint64_t size;
@@ -192,6 +200,14 @@ typedef struct _Py_DebugOffsets {
         uint64_t size;
         uint64_t collecting;
     } gc;
+
+    struct _gen_object {
+        uint64_t size;
+        uint64_t gi_name;
+        uint64_t gi_iframe;
+        uint64_t gi_task;
+        uint64_t gi_frame_state;
+    } gen_object;
 } _Py_DebugOffsets;
 
 /* Reference tracer state */

Include/internal/pycore_runtime_init.h

@@ -21,6 +21,7 @@ extern "C" {
 #include "pycore_runtime_init_generated.h"  // _Py_bytes_characters_INIT
 #include "pycore_signal.h"        // _signals_RUNTIME_INIT
 #include "pycore_tracemalloc.h"   // _tracemalloc_runtime_state_INIT
+#include "pycore_genobject.h"
 
 
 extern PyTypeObject _PyExc_MemoryError;
@@ -73,6 +74,7 @@ extern PyTypeObject _PyExc_MemoryError;
                 .instr_ptr = offsetof(_PyInterpreterFrame, instr_ptr), \
                 .localsplus = offsetof(_PyInterpreterFrame, localsplus), \
                 .owner = offsetof(_PyInterpreterFrame, owner), \
+                .stackpointer = offsetof(_PyInterpreterFrame, stackpointer), \
             }, \
             .code_object = { \
                 .size = sizeof(PyCodeObject), \
@@ -106,6 +108,11 @@ extern PyTypeObject _PyExc_MemoryError;
                 .ob_item = offsetof(PyListObject, ob_item), \
                 .ob_size = offsetof(PyListObject, ob_base.ob_size), \
             }, \
+            .set_object = { \
+                .size = sizeof(PySetObject), \
+                .used = offsetof(PySetObject, used), \
+                .table = offsetof(PySetObject, table), \
+            }, \
             .dict_object = { \
                 .size = sizeof(PyDictObject), \
                 .ma_keys = offsetof(PyDictObject, ma_keys), \
@@ -135,6 +142,13 @@ extern PyTypeObject _PyExc_MemoryError;
                 .size = sizeof(struct _gc_runtime_state), \
                 .collecting = offsetof(struct _gc_runtime_state, collecting), \
             }, \
+            .gen_object = { \
+                .size = sizeof(PyGenObject), \
+                .gi_name = offsetof(PyGenObject, gi_name), \
+                .gi_iframe = offsetof(PyGenObject, gi_iframe), \
+                .gi_task = offsetof(PyGenObject, gi_task), \
+                .gi_frame_state = offsetof(PyGenObject, gi_frame_state), \
+            }, \
         }, \
         .allocators = { \
             .standard = _pymem_allocators_standard_INIT(runtime), \

Lib/test/test_external_inspection.py

@@ -13,6 +13,7 @@
 try:
     from _testexternalinspection import PROCESS_VM_READV_SUPPORTED
     from _testexternalinspection import get_stack_trace
+    from _testexternalinspection import get_async_stack_trace
 except ImportError:
     raise unittest.SkipTest("Test only runs when _testexternalinspection is available")
 
@@ -74,6 +75,79 @@ def foo():
             ]
             self.assertEqual(stack_trace, expected_stack_trace)
 
+    @unittest.skipIf(sys.platform != "darwin" and sys.platform != "linux", "Test only runs on Linux and MacOS")
+    @unittest.skipIf(sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, "Test only runs on Linux with process_vm_readv support")
+    def test_async_remote_stack_trace(self):
+        # Spawn a process with some realistic Python code
+        script = textwrap.dedent("""\
+            import asyncio
+            import time
+            import os
+            import sys
+            import test.test_asyncio.test_stack as ts
+
+            def c5():
+                fifo = sys.argv[1]
+                with open(sys.argv[1], "w") as fifo:
+                    fifo.write("ready")
+                time.sleep(10000)
+
+            async def c4():
+                await asyncio.sleep(0)
+                c5()
+
+            async def c3():
+                await c4()
+
+            async def c2():
+                await c3()
+
+            async def c1(task):
+                await task
+
+            async def main():
+                async with asyncio.TaskGroup() as tg:
+                    task = tg.create_task(c2(), name="c2_root")
+                    tg.create_task(c1(task), name="sub_main_1")
+                    tg.create_task(c1(task), name="sub_main_2")
+
+            asyncio.run(main())
+            """)
+        stack_trace = None
+        with os_helper.temp_dir() as work_dir:
+            script_dir = os.path.join(work_dir, "script_pkg")
+            os.mkdir(script_dir)
+            fifo = f"{work_dir}/the_fifo"
+            os.mkfifo(fifo)
+            script_name = _make_test_script(script_dir, 'script', script)
+            try:
+                p = subprocess.Popen([sys.executable, script_name,  str(fifo)])
+                with open(fifo, "r") as fifo_file:
+                    response = fifo_file.read()
+                self.assertEqual(response, "ready")
+                stack_trace = get_async_stack_trace(p.pid)
+            except PermissionError:
+                self.skipTest("Insufficient permissions to read the stack trace")
+            finally:
+                os.remove(fifo)
+                p.kill()
+                p.terminate()
+                p.wait(timeout=SHORT_TIMEOUT)
+
+
+            expected_stack_trace = [
+                ["c5", "c4", "c3", "c2"],
+                "c2_root",
+                [
+                    [["main"], "Task-1", []],
+                    [["c1"], "sub_main_2", [[["main"], "Task-1", []]]],
+                    [["c1"], "sub_main_1", [[["main"], "Task-1", []]]],
+                ],
+            ]
+            self.assertEqual(stack_trace, expected_stack_trace)
+
+
+
     @unittest.skipIf(sys.platform != "darwin" and sys.platform != "linux", "Test only runs on Linux and MacOS")
     @unittest.skipIf(sys.platform == "linux" and not PROCESS_VM_READV_SUPPORTED, "Test only runs on Linux with process_vm_readv support")
     def test_self_trace(self):

Modules/_asynciomodule.c

@@ -16,6 +16,9 @@
 
 #include <stddef.h>               // offsetof()
 
+#if defined(__APPLE__)
+#  include <mach-o/loader.h>
+#endif
 
 /*[clinic input]
 module _asyncio
@@ -42,15 +45,15 @@ typedef enum {
     PyObject *prefix##_cancelled_exc;                                       \
     PyObject *prefix##_awaited_by;                                          \
     fut_state prefix##_state;                                               \
-    /* These bitfields need to be at the end of the struct
-       so that these and bitfields from TaskObj are contiguous.
+    /* Used by profilers to make traversing the stack from an external      \
+       process faster. */                                                   \
+    char prefix##_is_task;                                                  \
+    char prefix##_awaited_by_is_set;                                        \
+    /* These bitfields need to be at the end of the struct                  \
+       so that these and bitfields from TaskObj are contiguous.             \
     */                                                                      \
     unsigned prefix##_log_tb: 1;                                            \
     unsigned prefix##_blocking: 1;                                          \
-    /* Used by profilers to make traversing the stack from an external      \
-       process faster. */                                                   \
-    unsigned prefix##_is_task: 1;                                           \
-    unsigned prefix##_awaited_by_is_set: 1;
 
 typedef struct {
     FutureObj_HEAD(fut)
@@ -102,6 +105,40 @@ typedef struct {
 #endif
 
 typedef struct futureiterobject futureiterobject;
+typedef struct _Py_AsyncioModuleDebugOffsets {
+  struct _asyncio_task_object {
+    uint64_t size;
+    uint64_t task_name;
+    uint64_t task_awaited_by;
+    uint64_t task_is_task;
+    uint64_t task_awaited_by_is_set;
+    uint64_t task_coro;
+  } asyncio_task_object;
+} Py_AsyncioModuleDebugOffsets;
+
+#if defined(MS_WINDOWS)
+
+#pragma section("AsyncioDebug", read, write)
+__declspec(allocate("AsyncioDebug"))
+
+#elif defined(__APPLE__)
+
+__attribute__((section(SEG_DATA ",AsyncioDebug")))
+
+#endif
+
+Py_AsyncioModuleDebugOffsets AsyncioDebug
+#if defined(__linux__) && (defined(__GNUC__) || defined(__clang__))
+    __attribute__((section(".AsyncioDebug")))
+#endif
+    = {.asyncio_task_object = {
+           .size = sizeof(TaskObj),
+           .task_name = offsetof(TaskObj, task_name),
+           .task_awaited_by = offsetof(TaskObj, task_awaited_by),
+           .task_is_task = offsetof(TaskObj, task_is_task),
+           .task_awaited_by_is_set = offsetof(TaskObj, task_awaited_by_is_set),
+           .task_coro = offsetof(TaskObj, task_coro),
+       }};
 
 /* State of the _asyncio module */
 typedef struct {