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

@@ -417,6 +417,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_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

@@ -25,6 +25,10 @@ extern "C" {
 #include "pycore_typeobject.h"      // struct _types_runtime_state
 #include "pycore_unicodeobject.h"   // struct _Py_unicode_runtime_state
 
+#if defined(__APPLE__)
+#  include <mach-o/loader.h>
+#endif
+
 struct _getargs_runtime_state {
     struct _PyArg_Parser *static_parsers;
 };
@@ -59,6 +63,37 @@ typedef struct _Py_AuditHookEntry {
     void *userData;
 } _Py_AuditHookEntry;
 
+// 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
+
 typedef struct _Py_DebugOffsets {
     char cookie[8];
     uint64_t version;
@@ -108,6 +143,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 +188,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;
@@ -192,6 +236,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

@@ -22,6 +22,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;
@@ -74,6 +75,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), \
@@ -107,6 +109,12 @@ 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), \
+                .mask = offsetof(PySetObject, mask), \
+            }, \
             .dict_object = { \
                 .size = sizeof(PyDictObject), \
                 .ma_keys = offsetof(PyDictObject, ma_keys), \
@@ -136,6 +144,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/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__ +