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

Signed-off-by: Pablo Galindo <pablogsal@gmail.com>
Co-authored-by: Pablo Galindo <pablogsal@gmail.com>
Co-authored-by: Kumar Aditya <kumaraditya@python.org>
Co-authored-by: Łukasz Langa <lukasz@langa.pl>
Co-authored-by: Savannah Ostrowski <savannahostrowski@gmail.com>
Co-authored-by: Jacob Coffee <jacob@z7x.org>
Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>

Lib/asyncio/futures.py

@@ -2,6 +2,7 @@
 
 __all__ = (
     'Future', 'wrap_future', 'isfuture',
+    'future_add_to_awaited_by', 'future_discard_from_awaited_by',
 )
 
 import concurrent.futures
@@ -66,6 +67,9 @@ class Future:
     #   `yield Future()` (incorrect).
     _asyncio_future_blocking = False
 
+    # Used by the capture_call_stack() API.
+    __asyncio_awaited_by = None
+
     __log_traceback = False
 
     def __init__(self, *, loop=None):
@@ -115,6 +119,12 @@ class Future:
             raise ValueError('_log_traceback can only be set to False')
         self.__log_traceback = False
 
+    @property
+    def _asyncio_awaited_by(self):
+        if self.__asyncio_awaited_by is None:
+            return None
+        return frozenset(self.__asyncio_awaited_by)
+
     def get_loop(self):
         """Return the event loop the Future is bound to."""
         loop = self._loop
@@ -415,6 +425,49 @@ def wrap_future(future, *, loop=None):
     return new_future
 
 
+def future_add_to_awaited_by(fut, waiter, /):
+    """Record that `fut` is awaited on by `waiter`."""
+    # For the sake of keeping the implementation minimal and assuming
+    # that most of asyncio users use the built-in Futures and Tasks
+    # (or their subclasses), we only support native Future objects
+    # and their subclasses.
+    #
+    # Longer version: tracking requires storing the caller-callee
+    # dependency somewhere. One obvious choice is to store that
+    # information right in the future itself in a dedicated attribute.
+    # This means that we'd have to require all duck-type compatible
+    # futures to implement a specific attribute used by asyncio for
+    # the book keeping. Another solution would be to store that in
+    # a global dictionary. The downside here is that that would create
+    # strong references and any scenario where the "add" call isn't
+    # followed by a "discard" call would lead to a memory leak.
+    # Using WeakDict would resolve that issue, but would complicate
+    # the C code (_asynciomodule.c). The bottom line here is that
+    # it's not clear that all this work would be worth the effort.
+    #
+    # Note that there's an accelerated version of this function
+    # shadowing this implementation later in this file.
+    if isinstance(fut, _PyFuture) and isinstance(waiter, _PyFuture):
+        if fut._Future__asyncio_awaited_by is None:
+            fut._Future__asyncio_awaited_by = set()
+        fut._Future__asyncio_awaited_by.add(waiter)
+
+
+def future_discard_from_awaited_by(fut, waiter, /):
+    """Record that `fut` is no longer awaited on by `waiter`."""
+    # See the comment in "future_add_to_awaited_by()" body for
+    # details on implementation.
+    #
+    # Note that there's an accelerated version of this function
+    # shadowing this implementation later in this file.
+    if isinstance(fut, _PyFuture) and isinstance(waiter, _PyFuture):
+        if fut._Future__asyncio_awaited_by is not None:
+            fut._Future__asyncio_awaited_by.discard(waiter)
+
+
+_py_future_add_to_awaited_by = future_add_to_awaited_by
+_py_future_discard_from_awaited_by = future_discard_from_awaited_by
+
 try:
     import _asyncio
 except ImportError:
@@ -422,3 +475,7 @@ except ImportError:
 else:
     # _CFuture is needed for tests.
     Future = _CFuture = _asyncio.Future
+    future_add_to_awaited_by = _asyncio.future_add_to_awaited_by
+    future_discard_from_awaited_by = _asyncio.future_discard_from_awaited_by
+    _c_future_add_to_awaited_by = future_add_to_awaited_by
+    _c_future_discard_from_awaited_by = future_discard_from_awaited_by