GH-91048: Add utils for capturing async call stack for asyncio programs and enable profiling

Doc/library/asyncio-stack.rst

@@ -0,0 +1,141 @@
+.. currentmodule:: asyncio
+
+
+.. _asyncio-stack:
+
+===================
+Stack Introspection
+===================
+
+**Source code:** :source:`Lib/asyncio/stack.py`
+
+-------------------------------------
+
+asyncio has powerful runtime call stack introspection utilities
+to trace the entire call graph of a running coroutine or task, or
+a suspended *future*.  These utilities and the underlying machinery
+can be used by users in their Python code or by external profilers
+and debuggers.
+
+.. versionadded:: 3.14
+
+
+.. function:: print_call_graph(*, future=None, file=None, depth=1)
+
+   Print the async call graph for the current task or the provided
+   :class:`Task` or :class:`Future`.
+
+   The function receives an optional keyword-only *future* argument.
+   If not passed, the current running task will be used. If there's no
+   current task, the function returns ``None``.
+
+   If the function is called on *the current task*, the optional
+   keyword-only *depth* argument can be used to skip the specified
+   number of frames from top of the stack.
+
+   If *file* is not specified the function will print to :data:`sys.stdout`.
+
+   **Example:**
+
+   The following Python code:
+
+   .. code-block:: python
+
+      import asyncio
+
+      async def test():
+         asyncio.print_call_graph()
+
+      async def main():
+         async with asyncio.TaskGroup() as g:
+            g.create_task(test())
+
+      asyncio.run(main())
+
+   will print::
+
+      * Task(name='Task-2', id=0x1039f0fe0)
+      + Call stack:
+      |   File 't2.py', line 4, in async test()
+      + Awaited by:
+         * Task(name='Task-1', id=0x103a5e060)
+            + Call stack:
+            |   File 'taskgroups.py', line 107, in async TaskGroup.__aexit__()
+            |   File 't2.py', line 7, in async main()
+
+   For rendering the call stack to a string the following pattern
+   should be used:
+
+   .. code-block:: python
+
+      import io
+
+      ...
+
+      buf = io.StringIO()
+      asyncio.print_call_graph(file=buf)
+      output = buf.getvalue()
+
+
+.. function:: capture_call_graph(*, future=None)
+
+   Capture the async call graph for the current task or the provided
+   :class:`Task` or :class:`Future`.
+
+   The function receives an optional keyword-only *future* argument.
+   If not passed, the current running task will be used. If there's no
+   current task, the function returns ``None``.
+
+   If the function is called on *the current task*, the optional
+   keyword-only ``depth`` argument can be used to skip the specified
+   number of frames from top of the stack.
+
+   Returns a ``FutureCallGraph`` data class object:
+
+   * ``FutureCallGraph(future, call_stack, awaited_by)``
+
+      Where 'future' is a reference to a *Future* or a *Task*
+      (or their subclasses.)
+
+      ``call_stack`` is a list of ``FrameCallGraphEntry`` objects.
+
+      ``awaited_by`` is a list of ``FutureCallGraph`` objects.
+
+   * ``FrameCallGraphEntry(frame)``
+
+      Where ``frame`` is a frame object of a regular Python function
+      in the call stack.
+
+
+Low level utility functions
+===========================
+
+To introspect an async call graph asyncio requires cooperation from
+control flow structures, such as :func:`shield` or :class:`TaskGroup`.
+Any time an intermediate ``Future`` object with low-level APIs like
+:meth:`Future.add_done_callback() <asyncio.Future.add_done_callback>` is
+involved, the following two functions should be used to inform *asyncio*
+about how exactly such intermediate future objects are connected with
+the tasks they wrap or control.
+
+
+.. function:: future_add_to_awaited_by(future, waiter, /)
+
+   Record that *future* is awaited on by *waiter*.
+
+   Both *future* and *waiter* must be instances of
+   :class:`asyncio.Future <Future>` or :class:`asyncio.Task <Task>` or
+   their subclasses, otherwise the call would have no effect.
+
+   A call to ``future_add_to_awaited_by()`` must be followed by an
+   eventual call to the ``future_discard_from_awaited_by()`` function
+   with the same arguments.
+
+
+.. function:: future_discard_from_awaited_by(future, waiter, /)
+
+   Record that *future* is no longer awaited on by *waiter*.
+
+   Both *future* and *waiter* must be instances of
+   :class:`asyncio.Future <Future>` or :class:`asyncio.Task <Task>` or
+   their subclasses, otherwise the call would have no effect.

Doc/library/asyncio.rst

@@ -99,6 +99,7 @@ You can experiment with an ``asyncio`` concurrent context in the :term:`REPL`:
    asyncio-subprocess.rst
    asyncio-queue.rst
    asyncio-exceptions.rst
+   asyncio-stack.rst
 
 .. toctree::
    :caption: Low-level APIs

Doc/library/inspect.rst

@@ -162,6 +162,12 @@ attributes (see :ref:`import-mod-attrs` for module attributes):
 |                 |                   | per-opcode events are     |
 |                 |                   | requested                 |
 +-----------------+-------------------+---------------------------+
+|                 | f_generator       | returns the generator or  |
+|                 |                   | coroutine object that     |
+|                 |                   | owns this frame, or       |
+|                 |                   | ``None`` if the frame is  |
+|                 |                   | of a regular function     |
++-----------------+-------------------+---------------------------+
 |                 | clear()           | used to clear all         |
 |                 |                   | references to local       |
 |                 |                   | variables                 |
@@ -310,6 +316,10 @@ attributes (see :ref:`import-mod-attrs` for module attributes):
 
    Add ``__builtins__`` attribute to functions.
 
+.. versionchanged:: 3.14
+
+   Add ``f_generator`` attribute to frames.
+
 .. function:: getmembers(object[, predicate])
 
    Return all the members of an object in a list of ``(name, value)``

Doc/whatsnew/3.14.rst

@@ -434,6 +434,12 @@ asyncio
   reduces memory usage.
   (Contributed by Kumar Aditya in :gh:`107803`.)
 
+* :mod:`asyncio` has new utility functions for introspecting and printing
+  the program's call graph: :func:`asyncio.capture_call_graph` and
+  :func:`asyncio.print_call_graph`.
+  (Contributed by Yury Selivanov and Pablo Galindo Salgado in :gh:`91048`.)
+
+
 Deprecated
 ==========
 

Include/cpython/genobject.h

@@ -32,6 +32,8 @@ PyAPI_DATA(PyTypeObject) PyCoro_Type;
 PyAPI_FUNC(PyObject *) PyCoro_New(PyFrameObject *,
     PyObject *name, PyObject *qualname);
 
+PyAPI_FUNC(void) _PyCoro_SetTask(PyObject *coro, PyObject *task);
+
 
 /* --- Asynchronous Generators -------------------------------------------- */
 

Include/internal/pycore_debug_offsets.h

@@ -11,6 +11,41 @@ extern "C" {
 
 #define _Py_Debug_Cookie "xdebugpy"
 
+#if defined(__APPLE__)
+#  include <mach-o/loader.h>
+#endif
+
+// Macros to burn global values in custom sections so out-of-process
+// profilers can locate them easily
+
+#define GENERATE_DEBUG_SECTION(name, declaration) \
+    _GENERATE_DEBUG_SECTION_WINDOWS(name)         \
+    _GENERATE_DEBUG_SECTION_APPLE(name)           \
+    declaration                                   \
+    _GENERATE_DEBUG_SECTION_LINUX(name)
+
+#if defined(MS_WINDOWS)
+#define _GENERATE_DEBUG_SECTION_WINDOWS(name)                       \
+    _Pragma(Py_STRINGIFY(section(Py_STRINGIFY(name), read, write))) \
+    __declspec(allocate(Py_STRINGIFY(name)))
+#else
+#define _GENERATE_DEBUG_SECTION_WINDOWS(name)
+#endif
+
+#if defined(__APPLE__)
+#define _GENERATE_DEBUG_SECTION_APPLE(name) \
+    __attribute__((section(SEG_DATA "," Py_STRINGIFY(name))))
+#else
+#define _GENERATE_DEBUG_SECTION_APPLE(name)
+#endif
+
+#if defined(__linux__) && (defined(__GNUC__) || defined(__clang__))
+#define _GENERATE_DEBUG_SECTION_LINUX(name) \
+    __attribute__((section("." Py_STRINGIFY(name))))
+#else
+#define _GENERATE_DEBUG_SECTION_LINUX(name)
+#endif
+
 #ifdef Py_GIL_DISABLED
 # define _Py_Debug_gilruntimestate_enabled offsetof(struct _gil_runtime_state, enabled)
 # define _Py_Debug_Free_Threaded 1
@@ -69,6 +104,7 @@ typedef struct _Py_DebugOffsets {
         uint64_t instr_ptr;
         uint64_t localsplus;
         uint64_t owner;
+        uint64_t stackpointer;
     } interpreter_frame;
 
     // Code object offset;
@@ -113,6 +149,14 @@ 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;
+        uint64_t mask;
+    } set_object;
+
     // PyDict object offset;
     struct _dict_object {
         uint64_t size;
@@ -153,6 +197,15 @@ typedef struct _Py_DebugOffsets {
         uint64_t size;
         uint64_t collecting;
     } gc;
+
+    // Generator object offset;
+    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;
 
 
@@ -198,6 +251,7 @@ typedef struct _Py_DebugOffsets {
         .instr_ptr = offsetof(_PyInterpreterFrame, instr_ptr), \
         .localsplus = offsetof(_PyInterpreterFrame, localsplus), \
         .owner = offsetof(_PyInterpreterFrame, owner), \
+        .stackpointer = offsetof(_PyInterpreterFrame, stackpointer), \
     }, \
     .code_object = { \
         .size = sizeof(PyCodeObject), \
@@ -231,6 +285,12 @@ typedef struct _Py_DebugOffsets {
         .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), \
+        .mask = offsetof(PySetObject, mask), \
+    }, \
     .dict_object = { \
         .size = sizeof(PyDictObject), \
         .ma_keys = offsetof(PyDictObject, ma_keys), \
@@ -260,6 +320,13 @@ typedef struct _Py_DebugOffsets {
         .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), \
+    }, \
 }
 
 

Include/internal/pycore_genobject.h

@@ -22,6 +22,13 @@ extern "C" {
     PyObject *prefix##_qualname;                                            \
     _PyErr_StackItem prefix##_exc_state;                                    \
     PyObject *prefix##_origin_or_finalizer;                                 \
+    /* A *borrowed* reference to a task that drives the coroutine.          \
+       The field is meant to be used by profilers and debuggers only.       \
+       The main invariant is that a task can't get GC'ed while              \
+       the coroutine it drives is alive and vice versa.                     \
+       Profilers can use this field to reconstruct the full async           \
+       call stack of program. */                                            \
+    PyObject *prefix##_task;                                                \
     char prefix##_hooks_inited;                                             \
     char prefix##_closed;                                                   \
     char prefix##_running_async;                                            \

Include/internal/pycore_runtime.h

@@ -27,7 +27,6 @@ extern "C" {
 #include "pycore_typeobject.h"      // struct _types_runtime_state
 #include "pycore_unicodeobject.h"   // struct _Py_unicode_runtime_state
 
-
 /* Full Python runtime state */
 
 /* _PyRuntimeState holds the global state for the CPython runtime.

Include/internal/pycore_runtime_init.h

@@ -23,6 +23,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;

Lib/asyncio/__init__.py

@@ -14,6 +14,7 @@
 from .protocols import *
 from .runners import *
 from .queues import *
+from .stack import *
 from .streams import *
 from .subprocess import *
 from .tasks import *
@@ -31,6 +32,7 @@
            protocols.__all__ +
            runners.__all__ +
            queues.__all__ +
+           stack.__all__ +
            streams.__all__ +
            subprocess.__all__ +
            tasks.__all__ +