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

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

@@ -749,6 +749,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__ +