bpo-32314: Fix asyncio.run() to cancel runinng tasks on shutdown (#5262)

2 files changed, 35 insertions, 15 deletions

diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
index a10e706..ca9eee7 100644
--- a/Lib/asyncio/base_events.py
+++ b/Lib/asyncio/base_events.py
@@ -228,14 +228,9 @@ class BaseEventLoop(events.AbstractEventLoop):
         self._coroutine_origin_tracking_enabled = False
         self._coroutine_origin_tracking_saved_depth = None
 
-        if hasattr(sys, 'get_asyncgen_hooks'):
-            # Python >= 3.6
-            # A weak set of all asynchronous generators that are
-            # being iterated by the loop.
-            self._asyncgens = weakref.WeakSet()
-        else:
-            self._asyncgens = None
-
+        # A weak set of all asynchronous generators that are
+        # being iterated by the loop.
+        self._asyncgens = weakref.WeakSet()
         # Set to True when `loop.shutdown_asyncgens` is called.
         self._asyncgens_shutdown_called = False
 
@@ -354,7 +349,7 @@ class BaseEventLoop(events.AbstractEventLoop):
         """Shutdown all active asynchronous generators."""
         self._asyncgens_shutdown_called = True
 
-        if self._asyncgens is None or not len(self._asyncgens):
+        if not len(self._asyncgens):
             # If Python version is <3.6 or we don't have any asynchronous
             # generators alive.
             return
@@ -386,10 +381,10 @@ class BaseEventLoop(events.AbstractEventLoop):
                 'Cannot run the event loop while another loop is running')
         self._set_coroutine_origin_tracking(self._debug)
         self._thread_id = threading.get_ident()
-        if self._asyncgens is not None:
-            old_agen_hooks = sys.get_asyncgen_hooks()
-            sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
-                                   finalizer=self._asyncgen_finalizer_hook)
+
+        old_agen_hooks = sys.get_asyncgen_hooks()
+        sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
+                               finalizer=self._asyncgen_finalizer_hook)
         try:
             events._set_running_loop(self)
             while True:
@@ -401,8 +396,7 @@ class BaseEventLoop(events.AbstractEventLoop):
             self._thread_id = None
             events._set_running_loop(None)
             self._set_coroutine_origin_tracking(False)
-            if self._asyncgens is not None:
-                sys.set_asyncgen_hooks(*old_agen_hooks)
+            sys.set_asyncgen_hooks(*old_agen_hooks)
 
     def run_until_complete(self, future):
         """Run until the Future is done.
@@ -1374,6 +1368,7 @@ class BaseEventLoop(events.AbstractEventLoop):
         - 'message': Error message;
         - 'exception' (optional): Exception object;
         - 'future' (optional): Future instance;
+        - 'task' (optional): Task instance;
         - 'handle' (optional): Handle instance;
         - 'protocol' (optional): Protocol instance;
         - 'transport' (optional): Transport instance;
diff --git a/Lib/asyncio/runners.py b/Lib/asyncio/runners.py
index 94d9409..bb54b72 100644
--- a/Lib/asyncio/runners.py
+++ b/Lib/asyncio/runners.py
@@ -2,6 +2,7 @@ __all__ = 'run',
 
 from . import coroutines
 from . import events
+from . import tasks
 
 
 def run(main, *, debug=False):
@@ -42,7 +43,31 @@ def run(main, *, debug=False):
         return loop.run_until_complete(main)
     finally:
         try:
+            _cancel_all_tasks(loop)
             loop.run_until_complete(loop.shutdown_asyncgens())
         finally:
             events.set_event_loop(None)
             loop.close()
+
+
+def _cancel_all_tasks(loop):
+    to_cancel = [task for task in tasks.all_tasks(loop)
+                 if not task.done()]
+    if not to_cancel:
+        return
+
+    for task in to_cancel:
+        task.cancel()
+
+    loop.run_until_complete(
+        tasks.gather(*to_cancel, loop=loop, return_exceptions=True))
+
+    for task in to_cancel:
+        if task.cancelled():
+            continue
+        if task.exception() is not None:
+            loop.call_exception_handler({
+                'message': 'unhandled exception during asyncio.run() shutdown',
+                'exception': task.exception(),
+                'task': task,
+            })