Add asyncio subclass scheduling regression tests

Author: KRRT7

Lib/test/test_asyncio/test_futures.py

@@ -856,6 +856,75 @@ def call_soon(self, callback, *args, context=None):
         self.run_briefly()
         self.assertEqual(bag, [42])
 
+    def test_set_result_uses_private_call_soon_on_subclass_loop(self):
+        bag = []
+        base_loop_cls = type(self.loop)
+
+        class CountingLoop(base_loop_cls):
+            def __init__(self):
+                super().__init__()
+                self.private_calls = []
+
+            def _call_soon(self, callback, args, context):
+                self.private_calls.append((callback, args, context))
+                return super()._call_soon(callback, args, context)
+
+        old_loop = self.loop
+        loop = CountingLoop()
+        self.loop = loop
+        self.addCleanup(setattr, self, 'loop', old_loop)
+        self.addCleanup(loop.close)
+
+        f = self._new_future()
+        cb = self._make_callback(bag, 42)
+        f.add_done_callback(cb)
+        f.set_result('foo')
+
+        self.assertEqual(len(loop.private_calls), 1)
+        callback, args, context = loop.private_calls[0]
+        self.assertIs(callback, cb)
+        self.assertEqual(args, (f,))
+        self.assertIsNotNone(context)
+
+        self.run_briefly()
+        self.assertEqual(bag, [42])
+
+    def test_set_result_uses_instance_call_soon_shadow(self):
+        bag = []
+        base_loop_cls = type(self.loop)
+
+        class CountingLoop(base_loop_cls):
+            pass
+
+        old_loop = self.loop
+        loop = CountingLoop()
+        self.loop = loop
+        self.addCleanup(setattr, self, 'loop', old_loop)
+        self.addCleanup(loop.close)
+
+        base_call_soon = base_loop_cls.call_soon
+        calls = []
+
+        def call_soon(callback, *args, context=None):
+            calls.append((callback, args, context))
+            return base_call_soon(loop, callback, *args, context=context)
+
+        loop.call_soon = call_soon
+
+        f = self._new_future()
+        cb = self._make_callback(bag, 42)
+        f.add_done_callback(cb)
+        f.set_result('foo')
+
+        self.assertEqual(len(calls), 1)
+        callback, args, context = calls[0]
+        self.assertIs(callback, cb)
+        self.assertEqual(args, (f,))
+        self.assertIsNotNone(context)
+
+        self.run_briefly()
+        self.assertEqual(bag, [42])
+
     def test_callbacks_remove_first_callback(self):
         bag = []
         f = self._new_future()

Lib/test/test_asyncio/test_tasks.py

@@ -218,6 +218,66 @@ async def coro():
 
         loop.run_until_complete(task)
 
+    def test_task_constructor_uses_private_call_soon_on_subclass_loop(self):
+        base_loop_cls = type(self.loop)
+
+        class CountingLoop(base_loop_cls):
+            def __init__(self):
+                super().__init__()
+                self.private_calls = []
+
+            def _call_soon(self, callback, args, context):
+                self.private_calls.append((callback, args, context))
+                return super()._call_soon(callback, args, context)
+
+        async def coro():
+            return 42
+
+        loop = CountingLoop()
+        self.addCleanup(loop.close)
+        self.set_event_loop(loop)
+        self.addCleanup(asyncio.set_event_loop, None)
+
+        task = self.new_task(loop, coro())
+        self.assertEqual(len(loop.private_calls), 1)
+        callback, args, context = loop.private_calls[0]
+        self.assertIs(getattr(callback, '__self__', None), task)
+        self.assertEqual(args, ())
+        self.assertIs(task.get_context(), context)
+
+        loop.run_until_complete(task)
+
+    def test_task_constructor_uses_instance_call_soon_shadow(self):
+        base_loop_cls = type(self.loop)
+
+        class CountingLoop(base_loop_cls):
+            pass
+
+        async def coro():
+            return 42
+
+        loop = CountingLoop()
+        self.addCleanup(loop.close)
+        self.set_event_loop(loop)
+        self.addCleanup(asyncio.set_event_loop, None)
+
+        base_call_soon = base_loop_cls.call_soon
+        calls = []
+
+        def call_soon(callback, *args, context=None):
+            calls.append((callback, args, context))
+            return base_call_soon(loop, callback, *args, context=context)
+
+        loop.call_soon = call_soon
+        task = self.new_task(loop, coro())
+        self.assertEqual(len(calls), 1)
+        callback, args, context = calls[0]
+        self.assertIs(getattr(callback, '__self__', None), task)
+        self.assertEqual(args, ())
+        self.assertIs(task.get_context(), context)
+
+        loop.run_until_complete(task)
+
     def test_ensure_future_coroutine(self):
         async def notmuch():
             return 'ok'

Modules/_asynciomodule.c

@@ -369,6 +369,9 @@ get_event_loop(asyncio_state *state)
 static int
 is_asyncio_base_event_loop_subclass(PyTypeObject *loop_type)
 {
+    // The _asyncio module is imported from asyncio.events while
+    // asyncio.base_events is still initializing, so we can't cache the
+    // BaseEventLoop type during module init without creating an import cycle.
     for (PyTypeObject *base = loop_type; base != NULL; base = base->tp_base) {
         if (strcmp(base->tp_name, "BaseEventLoop") != 0) {
             continue;
@@ -395,6 +398,9 @@ is_asyncio_base_event_loop_subclass(PyTypeObject *loop_type)
 static int
 loop_has_instance_call_soon_shadow(PyObject *loop)
 {
+    // Respect instance-level monkeypatching such as:
+    //     loop.call_soon = custom_call_soon
+    // These must keep using the public slow path.
     PyTypeObject *loop_type = Py_TYPE(loop);
     PyObject *attr = NULL;
 
@@ -457,6 +463,8 @@ call_soon(asyncio_state *state, PyObject *loop, PyObject *func, PyObject *arg,
         is_base_event_loop_subclass &&
         module_name != NULL &&
         PyUnicode_Check(module_name) &&
+        // Only use the private fast path when the subclass keeps the public
+        // call_soon() contract unchanged at both the type and instance level.
         !has_instance_call_soon_shadow &&
         (is_stdlib_base_event_loop || call_soon_attr == NULL)
     );