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

Signed-off-by: Pablo Galindo <[email protected]>
Co-authored-by: Pablo Galindo <[email protected]>
Co-authored-by: Kumar Aditya <[email protected]>
Co-authored-by: Łukasz Langa <[email protected]>
Co-authored-by: Savannah Ostrowski <[email protected]>
Co-authored-by: Jacob Coffee <[email protected]>
Co-authored-by: Irit Katriel <[email protected]>

Author: 7 people

Doc/library/asyncio-graph.rst

@@ -0,0 +1,145 @@
+.. currentmodule:: asyncio
+
+
+.. _asyncio-graph:
+
+========================
+Call Graph Introspection
+========================
+
+**Source code:** :source:`Lib/asyncio/graph.py`
+
+-------------------------------------
+
+asyncio has powerful runtime call graph 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 from within a Python program or by external profilers
+and debuggers.
+
+.. versionadded:: next
+
+
+.. function:: print_call_graph(future=None, /, *, file=None, depth=1, limit=None)
+
+   Print the async call graph for the current task or the provided
+   :class:`Task` or :class:`Future`.
+
+   This function prints entries starting from the top frame and going
+   down towards the invocation point.
+
+   The function receives an optional *future* argument.
+   If not passed, the current running task will be used.
+
+   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 the optional keyword-only *limit* argument is provided, each call stack
+   in the resulting graph is truncated to include at most ``abs(limit)``
+   entries. If *limit* is positive, the entries left are the closest to
+   the invocation point. If *limit* is negative, the topmost entries are
+   left. If *limit* is omitted or ``None``, all entries are present.
+   If *limit* is ``0``, the call stack is not printed at all, only
+   "awaited by" information is printed.
+
+   If *file* is omitted or ``None``, 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()
+
+.. function:: format_call_graph(future=None, /, *, depth=1, limit=None)
+
+   Like :func:`print_call_graph`, but returns a string.
+   If *future* is ``None`` and there's no current task,
+   the function returns an empty string.
+
+
+.. function:: capture_call_graph(future=None, /, *, depth=1, limit=None)
+
+   Capture the async call graph for the current task or the provided
+   :class:`Task` or :class:`Future`.
+
+   The function receives an optional *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 :class:`Future` or
+      a :class:`Task` (or their subclasses.)
+
+      ``call_stack`` is a tuple of ``FrameCallGraphEntry`` objects.
+
+      ``awaited_by`` is a tuple 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 :class:`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:`Future` or :class:`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 :func:`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:`Future` or :class:`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-graph.rst
 
 .. toctree::
    :caption: Low-level APIs

Doc/library/inspect.rst

@@ -150,6 +150,12 @@ attributes (see :ref:`import-mod-attrs` for module attributes):
 |                 | f_locals          | local namespace seen by   |
 |                 |                   | this frame                |
 +-----------------+-------------------+---------------------------+
+|                 | f_generator       | returns the generator or  |
+|                 |                   | coroutine object that     |
+|                 |                   | owns this frame, or       |
+|                 |                   | ``None`` if the frame is  |
+|                 |                   | of a regular function     |
++-----------------+-------------------+---------------------------+
 |                 | f_trace           | tracing function for this |
 |                 |                   | frame, or ``None``        |
 +-----------------+-------------------+---------------------------+
@@ -310,6 +316,10 @@ attributes (see :ref:`import-mod-attrs` for module attributes):
 
    Add ``__builtins__`` attribute to functions.
 
+.. versionchanged:: next
+
+   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

@@ -755,6 +755,11 @@ 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, Pablo Galindo Salgado, and Łukasz Langa
+  in :gh:`91048`.)
 
 base64
 ------

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,14 @@ 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_frame_state;
+    } gen_object;
 } _Py_DebugOffsets;
 
 
@@ -198,6 +250,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 +284,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 +319,12 @@ 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_frame_state = offsetof(PyGenObject, gi_frame_state), \
+    }, \
 }
 
 

Include/internal/pycore_tstate.h

@@ -22,6 +22,7 @@ typedef struct _PyThreadStateImpl {
     PyThreadState base;
 
     PyObject *asyncio_running_loop; // Strong reference
+    PyObject *asyncio_running_task; // Strong reference
 
     struct _qsbr_thread_state *qsbr;  // only used by free-threaded build
     struct llist_node mem_free_queue; // delayed free queue

Lib/asyncio/__init__.py

@@ -10,6 +10,7 @@
 from .events import *
 from .exceptions import *
 from .futures import *
+from .graph import *
 from .locks import *
 from .protocols import *
 from .runners import *
@@ -27,6 +28,7 @@
            events.__all__ +
            exceptions.__all__ +
            futures.__all__ +
+           graph.__all__ +
            locks.__all__ +
            protocols.__all__ +
            runners.__all__ +