bpo-32357: Optimize asyncio.iscoroutine() for non-native coroutines (#4915)

Author: 1st1

Lib/asyncio/coroutines.py

@@ -1,14 +1,13 @@
 __all__ = 'coroutine', 'iscoroutinefunction', 'iscoroutine'
 
+import collections.abc
 import functools
 import inspect
 import os
 import sys
 import traceback
 import types
 
-from collections.abc import Awaitable, Coroutine
-
 from . import base_futures
 from . import constants
 from . import format_helpers
@@ -162,7 +161,7 @@ def coro(*args, **kw):
                 except AttributeError:
                     pass
                 else:
-                    if isinstance(res, Awaitable):
+                    if isinstance(res, collections.abc.Awaitable):
                         res = yield from await_meth()
             return res
 
@@ -199,12 +198,24 @@ def iscoroutinefunction(func):
 # Prioritize native coroutine check to speed-up
 # asyncio.iscoroutine.
 _COROUTINE_TYPES = (types.CoroutineType, types.GeneratorType,
-                    Coroutine, CoroWrapper)
+                    collections.abc.Coroutine, CoroWrapper)
+_iscoroutine_typecache = set()
 
 
 def iscoroutine(obj):
     """Return True if obj is a coroutine object."""
-    return isinstance(obj, _COROUTINE_TYPES)
+    if type(obj) in _iscoroutine_typecache:
+        return True
+
+    if isinstance(obj, _COROUTINE_TYPES):
+        # Just in case we don't want to cache more than 100
+        # positive types.  That shouldn't ever happen, unless
+        # someone stressing the system on purpose.
+        if len(_iscoroutine_typecache) < 100:
+            _iscoroutine_typecache.add(type(obj))
+        return True
+    else:
+        return False
 
 
 def _format_coroutine(coro):

Lib/test/test_asyncio/test_tasks.py

@@ -62,6 +62,20 @@ def __call__(self, *args):
         pass
 
 
+class CoroLikeObject:
+    def send(self, v):
+        raise StopIteration(42)
+
+    def throw(self, *exc):
+        pass
+
+    def close(self):
+        pass
+
+    def __await__(self):
+        return self
+
+
 class BaseTaskTests:
 
     Task = None
@@ -2085,6 +2099,12 @@ def test_create_task_with_noncoroutine(self):
                                     "a coroutine was expected, got 123"):
             self.new_task(self.loop, 123)
 
+        # test it for the second time to ensure that caching
+        # in asyncio.iscoroutine() doesn't break things.
+        with self.assertRaisesRegex(TypeError,
+                                    "a coroutine was expected, got 123"):
+            self.new_task(self.loop, 123)
+
     def test_create_task_with_oldstyle_coroutine(self):
 
         @asyncio.coroutine
@@ -2095,6 +2115,12 @@ def coro():
         self.assertIsInstance(task, self.Task)
         self.loop.run_until_complete(task)
 
+        # test it for the second time to ensure that caching
+        # in asyncio.iscoroutine() doesn't break things.
+        task = self.new_task(self.loop, coro())
+        self.assertIsInstance(task, self.Task)
+        self.loop.run_until_complete(task)
+
     def test_create_task_with_async_function(self):
 
         async def coro():
@@ -2104,6 +2130,23 @@ async def coro():
         self.assertIsInstance(task, self.Task)
         self.loop.run_until_complete(task)
 
+        # test it for the second time to ensure that caching
+        # in asyncio.iscoroutine() doesn't break things.
+        task = self.new_task(self.loop, coro())
+        self.assertIsInstance(task, self.Task)
+        self.loop.run_until_complete(task)
+
+    def test_create_task_with_asynclike_function(self):
+        task = self.new_task(self.loop, CoroLikeObject())
+        self.assertIsInstance(task, self.Task)
+        self.assertEqual(self.loop.run_until_complete(task), 42)
+
+        # test it for the second time to ensure that caching
+        # in asyncio.iscoroutine() doesn't break things.
+        task = self.new_task(self.loop, CoroLikeObject())
+        self.assertIsInstance(task, self.Task)
+        self.assertEqual(self.loop.run_until_complete(task), 42)
+
     def test_bare_create_task(self):
 
         async def inner():

Misc/NEWS.d/next/Library/2017-12-18-00-36-41.bpo-32357.t1F3sn.rst

@@ -0,0 +1,5 @@
+Optimize asyncio.iscoroutine() and loop.create_task() for non-native
+coroutines (e.g. async/await compiled with Cython).
+
+'loop.create_task(python_coroutine)' used to be 20% faster than
+'loop.create_task(cython_coroutine)'.  Now, the latter is as fast.

Modules/_asynciomodule.c

@@ -49,6 +49,9 @@ static PyObject *current_tasks;
    all running event loops.  {EventLoop: Task} */
 static PyObject *all_tasks;
 
+/* An isinstance type cache for the 'is_coroutine()' function. */
+static PyObject *iscoroutine_typecache;
+
 
 typedef enum {
     STATE_PENDING,
@@ -118,6 +121,71 @@ static PyObject* future_new_iter(PyObject *);
 static inline int future_call_schedule_callbacks(FutureObj *);
 
 
+static int
+_is_coroutine(PyObject *coro)
+{
+    /* 'coro' is not a native coroutine, call asyncio.iscoroutine()
+       to check if it's another coroutine flavour.
+
+       Do this check after 'future_init()'; in case we need to raise
+       an error, __del__ needs a properly initialized object.
+    */
+    PyObject *res = PyObject_CallFunctionObjArgs(
+        asyncio_iscoroutine_func, coro, NULL);
+    if (res == NULL) {
+        return -1;
+    }
+
+    int is_res_true = PyObject_IsTrue(res);
+    Py_DECREF(res);
+    if (is_res_true <= 0) {
+        return is_res_true;
+    }
+
+    if (PySet_Size(iscoroutine_typecache) < 100) {
+        /* Just in case we don't want to cache more than 100
+           positive types.  That shouldn't ever happen, unless
+           someone stressing the system on purpose.
+        */
+        if (PySet_Add(iscoroutine_typecache, (PyObject*) Py_TYPE(coro))) {
+            return -1;
+        }
+    }
+
+    return 1;
+}
+
+
+static inline int
+is_coroutine(PyObject *coro)
+{
+    if (PyCoro_CheckExact(coro)) {
+        return 1;
+    }
+
+    /* Check if `type(coro)` is in the cache.
+       Caching makes is_coroutine() function almost as fast as
+       PyCoro_CheckExact() for non-native coroutine-like objects
+       (like coroutines compiled with Cython).
+
+       asyncio.iscoroutine() has its own type caching mechanism.
+       This cache allows us to avoid the cost of even calling
+       a pure-Python function in 99.9% cases.
+    */
+    int has_it = PySet_Contains(
+        iscoroutine_typecache, (PyObject*) Py_TYPE(coro));
+    if (has_it == 0) {
+        /* type(coro) is not in iscoroutine_typecache */
+        return _is_coroutine(coro);
+    }
+
+    /* either an error has occured or
+       type(coro) is in iscoroutine_typecache
+    */
+    return has_it;
+}
+
+
 static int
 get_running_loop(PyObject **loop)
 {
@@ -1778,37 +1846,20 @@ static int
 _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop)
 /*[clinic end generated code: output=9f24774c2287fc2f input=8d132974b049593e]*/
 {
-    PyObject *res;
-
     if (future_init((FutureObj*)self, loop)) {
         return -1;
     }
 
-    if (!PyCoro_CheckExact(coro)) {
-        /* 'coro' is not a native coroutine, call asyncio.iscoroutine()
-           to check if it's another coroutine flavour.
-
-           Do this check after 'future_init()'; in case we need to raise
-           an error, __del__ needs a properly initialized object.
-        */
-        res = PyObject_CallFunctionObjArgs(
-            asyncio_iscoroutine_func, coro, NULL);
-        if (res == NULL) {
-            return -1;
-        }
-
-        int tmp = PyObject_Not(res);
-        Py_DECREF(res);
-        if (tmp < 0) {
-            return -1;
-        }
-        if (tmp) {
-            self->task_log_destroy_pending = 0;
-            PyErr_Format(PyExc_TypeError,
-                         "a coroutine was expected, got %R",
-                         coro, NULL);
-            return -1;
-        }
+    int is_coro = is_coroutine(coro);
+    if (is_coro == -1) {
+        return -1;
+    }
+    if (is_coro == 0) {
+        self->task_log_destroy_pending = 0;
+        PyErr_Format(PyExc_TypeError,
+                     "a coroutine was expected, got %R",
+                     coro, NULL);
+        return -1;
     }
 
     self->task_fut_waiter = NULL;
@@ -3007,8 +3058,9 @@ module_free(void *m)
     Py_CLEAR(asyncio_InvalidStateError);
     Py_CLEAR(asyncio_CancelledError);
 
-    Py_CLEAR(current_tasks);
     Py_CLEAR(all_tasks);
+    Py_CLEAR(current_tasks);
+    Py_CLEAR(iscoroutine_typecache);
 
     module_free_freelists();
 }
@@ -3028,6 +3080,11 @@ module_init(void)
         goto fail;
     }
 
+    iscoroutine_typecache = PySet_New(NULL);
+    if (iscoroutine_typecache == NULL) {
+        goto fail;
+    }
+
 #define WITH_MOD(NAME) \
     Py_CLEAR(module); \
     module = PyImport_ImportModule(NAME); \