Optimize asyncio internal call_soon scheduling

Author: KRRT7

Lib/test/test_asyncio/test_futures.py

@@ -823,6 +823,39 @@ def bag_appender(future):
     def _new_future(self):
         raise NotImplementedError
 
+    def test_set_result_uses_call_soon_override(self):
+        bag = []
+        base_loop_cls = type(self.loop)
+
+        class CountingLoop(base_loop_cls):
+            def __init__(self):
+                super().__init__()
+                self.calls = []
+
+            def call_soon(self, callback, *args, context=None):
+                self.calls.append((callback, args, context))
+                return super().call_soon(callback, *args, context=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.calls), 1)
+        callback, args, context = loop.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

@@ -189,6 +189,35 @@ async def notmuch():
         loop.run_until_complete(t)
         loop.close()
 
+    def test_task_constructor_uses_call_soon_override(self):
+        base_loop_cls = type(self.loop)
+
+        class CountingLoop(base_loop_cls):
+            def __init__(self):
+                super().__init__()
+                self.calls = []
+
+            def call_soon(self, callback, *args, context=None):
+                self.calls.append((callback, args, context))
+                return super().call_soon(callback, *args, context=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.calls), 1)
+        callback, args, context = loop.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

@@ -149,6 +149,7 @@ typedef struct {
     PyTypeObject *TaskType;
 
     PyObject *asyncio_mod;
+    PyObject *call_soon_private_name;
     PyObject *context_kwname;
 
     /* WeakSet containing scheduled 3rd party tasks which don't
@@ -370,27 +371,78 @@ call_soon(asyncio_state *state, PyObject *loop, PyObject *func, PyObject *arg,
           PyObject *ctx)
 {
     state->bench_call_soon_calls++;
-    PyObject *handle;
+    PyObject *handle = NULL;
 
-    if (ctx == NULL) {
-        PyObject *stack[] = {loop, func, arg};
-        size_t nargsf = 3 | PY_VECTORCALL_ARGUMENTS_OFFSET;
-        handle = PyObject_VectorcallMethod(&_Py_ID(call_soon), stack, nargsf, NULL);
+    PyTypeObject *loop_type = Py_TYPE(loop);
+    PyObject *module_name = PyDict_GetItemWithError(loop_type->tp_dict, &_Py_ID(__module__));
+    if (module_name == NULL && PyErr_Occurred()) {
+        return -1;
     }
-    else {
-        /* All refs in 'stack' are borrowed. */
-        PyObject *stack[4];
-        size_t nargs = 2;
-        stack[0] = loop;
-        stack[1] = func;
+
+    int use_fastpath = (
+        module_name != NULL &&
+        PyUnicode_Check(module_name) &&
+        (
+            (strcmp(loop_type->tp_name, "BaseEventLoop") == 0 &&
+             PyUnicode_CompareWithASCIIString(module_name, "asyncio.base_events") == 0) ||
+            (strcmp(loop_type->tp_name, "BaseSelectorEventLoop") == 0 &&
+             PyUnicode_CompareWithASCIIString(module_name, "asyncio.selector_events") == 0) ||
+            (strcmp(loop_type->tp_name, "_UnixSelectorEventLoop") == 0 &&
+             PyUnicode_CompareWithASCIIString(module_name, "asyncio.unix_events") == 0) ||
+            (strcmp(loop_type->tp_name, "_WindowsSelectorEventLoop") == 0 &&
+             PyUnicode_CompareWithASCIIString(module_name, "asyncio.windows_events") == 0) ||
+            (strcmp(loop_type->tp_name, "BaseProactorEventLoop") == 0 &&
+             PyUnicode_CompareWithASCIIString(module_name, "asyncio.proactor_events") == 0) ||
+            (strcmp(loop_type->tp_name, "ProactorEventLoop") == 0 &&
+             PyUnicode_CompareWithASCIIString(module_name, "asyncio.windows_events") == 0)
+        )
+    );
+
+    if (use_fastpath) {
+        PyObject *args;
+        if (arg == NULL) {
+            args = Py_GetConstantBorrowed(Py_CONSTANT_EMPTY_TUPLE);
+        }
+        else {
+            args = PyTuple_Pack(1, arg);
+            if (args == NULL) {
+                return -1;
+            }
+        }
+
+        PyObject *stack[] = {loop, func, args, ctx ? ctx : Py_None};
+        size_t nargsf = 4 | PY_VECTORCALL_ARGUMENTS_OFFSET;
+        handle = PyObject_VectorcallMethod(state->call_soon_private_name, stack,
+                                           nargsf, NULL);
         if (arg != NULL) {
-            stack[2] = arg;
-            nargs++;
+            Py_DECREF(args);
+        }
+
+        if (handle == NULL && PyErr_ExceptionMatches(PyExc_AttributeError)) {
+            PyErr_Clear();
+        }
+    }
+
+    if (handle == NULL) {
+        if (ctx == NULL) {
+            PyObject *stack[] = {loop, func, arg};
+            size_t nargsf = 3 | PY_VECTORCALL_ARGUMENTS_OFFSET;
+            handle = PyObject_VectorcallMethod(&_Py_ID(call_soon), stack, nargsf, NULL);
+        }
+        else {
+            PyObject *stack[4];
+            size_t nargs = 2;
+            stack[0] = loop;
+            stack[1] = func;
+            if (arg != NULL) {
+                stack[2] = arg;
+                nargs++;
+            }
+            stack[nargs] = (PyObject *)ctx;
+            size_t nargsf = nargs | PY_VECTORCALL_ARGUMENTS_OFFSET;
+            handle = PyObject_VectorcallMethod(&_Py_ID(call_soon), stack, nargsf,
+                                               state->context_kwname);
         }
-        stack[nargs] = (PyObject *)ctx;
-        size_t nargsf = nargs | PY_VECTORCALL_ARGUMENTS_OFFSET;
-        handle = PyObject_VectorcallMethod(&_Py_ID(call_soon), stack, nargsf,
-                                           state->context_kwname);
     }
 
     if (handle == NULL) {
@@ -4204,6 +4256,7 @@ module_traverse(PyObject *mod, visitproc visit, void *arg)
     Py_VISIT(state->TaskType);
 
     Py_VISIT(state->asyncio_mod);
+    Py_VISIT(state->call_soon_private_name);
     Py_VISIT(state->traceback_extract_stack);
     Py_VISIT(state->asyncio_future_repr_func);
     Py_VISIT(state->asyncio_get_event_loop);
@@ -4234,6 +4287,7 @@ module_clear(PyObject *mod)
     Py_CLEAR(state->TaskType);
 
     Py_CLEAR(state->asyncio_mod);
+    Py_CLEAR(state->call_soon_private_name);
     Py_CLEAR(state->traceback_extract_stack);
     Py_CLEAR(state->asyncio_future_repr_func);
     Py_CLEAR(state->asyncio_get_event_loop);
@@ -4285,6 +4339,11 @@ module_init(asyncio_state *state)
         goto fail;
     }
 
+    state->call_soon_private_name = PyUnicode_InternFromString("_call_soon");
+    if (state->call_soon_private_name == NULL) {
+        goto fail;
+    }
+
 #define WITH_MOD(NAME) \
     Py_CLEAR(module); \
     module = PyImport_ImportModule(NAME); \