Implement fast Task & Future.

In some cases this can boost the performance up to 30%.

Author: 1st1

tests/test_futures.py

@@ -4,7 +4,7 @@
 
 from . import includes as __includes
 from . import _patch
-from .loop import Loop as __BaseLoop
+from .loop import Loop as __BaseLoop, Future
 
 
 __all__ = ('new_event_loop', 'EventLoopPolicy')

tests/test_tasks.py

@@ -1,6 +1,4 @@
-import asyncio
-
-from asyncio import coroutines
+import sys
 
 
 def _format_coroutine(coro):
@@ -28,5 +26,49 @@ def _format_coroutine(coro):
     return _old_format_coroutine(coro)
 
 
-_old_format_coroutine = coroutines._format_coroutine
-coroutines._format_coroutine = _format_coroutine
+async def _wait_for_data(self, func_name):
+    """Wait until feed_data() or feed_eof() is called.
+
+    If stream was paused, automatically resume it.
+    """
+    if self._waiter is not None:
+        raise RuntimeError('%s() called while another coroutine is '
+                           'already waiting for incoming data' % func_name)
+
+    assert not self._eof, '_wait_for_data after EOF'
+
+    # Waiting for data while paused will make deadlock, so prevent it.
+    if self._paused:
+        self._paused = False
+        self._transport.resume_reading()
+
+    try:
+        create_future = self._loop.create_future
+    except AttributeError:
+        self._waiter = asyncio.Future(loop=self._loop)
+    else:
+        self._waiter = create_future()
+
+    try:
+        await self._waiter
+    finally:
+        self._waiter = None
+
+
+if sys.version_info < (3, 5, 2):
+    # In Python 3.5.2 (see PEP 478 for 3.5 release schedule)
+    # we won't need to patch anything.
+
+    import asyncio
+
+    from asyncio import coroutines
+    from asyncio import streams
+
+    # This is needed to support Cython 'async def' coroutines.
+    _old_format_coroutine = coroutines._format_coroutine
+    coroutines._format_coroutine = _format_coroutine
+
+    # Fix a possible deadlock, improve performance.
+    _old_wait_for_data = streams.StreamReader._wait_for_data
+    _wait_for_data.__module__ = _old_wait_for_data.__module__
+    streams.StreamReader._wait_for_data = _wait_for_data

uvloop/__init__.py

@@ -0,0 +1,229 @@
+DEF _FUT_PENDING = 1
+DEF _FUT_CANCELLED = 2
+DEF _FUT_FINISHED = 3
+
+
+cdef class BaseFuture:
+    cdef:
+        int _state
+
+        readonly Loop _loop
+        readonly list _callbacks
+        readonly object _exception
+        readonly object _result
+        public bint _blocking
+        readonly object _source_traceback
+        readonly bint _log_traceback
+
+    def __init__(self, Loop loop):
+        if loop is None:
+            loop = aio_get_event_loop()
+            if not isinstance(loop, Loop):
+                raise TypeError('uvloop.Future supports only uvloop.Loop')
+
+        self._state = _FUT_PENDING
+        self._loop = loop
+        self._callbacks = []
+        self._result = None
+        self._exception = None
+        self._blocking = False
+        self._log_traceback = False
+
+        if loop._debug:
+            self._source_traceback = tb_extract_stack(sys_getframe(0))
+        else:
+            self._source_traceback = None
+
+    cdef _schedule_callbacks(self):
+        cdef:
+            list callbacks
+            size_t cb_len = len(self._callbacks)
+            size_t i
+
+        if cb_len == 0:
+            return
+
+        callbacks = self._callbacks
+        self._callbacks = []
+
+        for i from 0 <= i < cb_len:
+            self._loop.call_soon(callbacks[i], self)
+
+    cdef _add_done_callback(self, fn):
+        if self._state != _FUT_PENDING:
+            self._loop.call_soon(fn, self)
+        else:
+            self._callbacks.append(fn)
+
+    cdef _done(self):
+        return self._state != _FUT_PENDING
+
+    cdef _cancel(self):
+        if self._done():
+            return False
+        self._state = _FUT_CANCELLED
+        self._schedule_callbacks()
+        return True
+
+    # _result would shadow the "_result" property
+    cdef _result_impl(self):
+        if self._state == _FUT_CANCELLED:
+            raise aio_CancelledError
+        if self._state != _FUT_FINISHED:
+            raise aio_InvalidStateError('Result is not ready.')
+        self._log_traceback = False
+        if self._exception is not None:
+            raise self._exception
+        return self._result
+
+    cdef _str_state(self):
+        if self._state == _FUT_PENDING:
+            return 'PENDING'
+        elif self._state == _FUT_CANCELLED:
+            return 'CANCELLED'
+        elif self._state == _FUT_FINISHED:
+            return 'FINISHED'
+        else:
+            raise RuntimeError('unknown Future state')
+
+    property _state:
+        def __get__(self):
+            return self._str_state()
+
+    def cancel(self):
+        """Cancel the future and schedule callbacks.
+
+        If the future is already done or cancelled, return False.  Otherwise,
+        change the future's state to cancelled, schedule the callbacks and
+        return True.
+        """
+        return self._cancel()
+
+    def cancelled(self):
+        """Return True if the future was cancelled."""
+        return self._state == _FUT_CANCELLED
+
+    def done(self):
+        """Return True if the future is done.
+
+        Done means either that a result / exception are available, or that the
+        future was cancelled.
+        """
+        return self._state != _FUT_PENDING
+
+    def result(self):
+        """Return the result this future represents.
+
+        If the future has been cancelled, raises CancelledError.  If the
+        future's result isn't yet available, raises InvalidStateError.  If
+        the future is done and has an exception set, this exception is raised.
+        """
+        return self._result_impl()
+
+    def exception(self):
+        """Return the exception that was set on this future.
+
+        The exception (or None if no exception was set) is returned only if
+        the future is done.  If the future has been cancelled, raises
+        CancelledError.  If the future isn't done yet, raises
+        InvalidStateError.
+        """
+        if self._state == _FUT_CANCELLED:
+            raise aio_CancelledError
+        if self._state != _FUT_FINISHED:
+            raise aio_InvalidStateError('Exception is not set.')
+        self._log_traceback = False
+        return self._exception
+
+    def add_done_callback(self, fn):
+        """Add a callback to be run when the future becomes done.
+
+        The callback is called with a single argument - the future object. If
+        the future is already done when this is called, the callback is
+        scheduled with call_soon.
+        """
+        self._add_done_callback(fn)
+
+    def remove_done_callback(self, fn):
+        """Remove all instances of a callback from the "call when done" list.
+
+        Returns the number of callbacks removed.
+        """
+        cdef:
+            size_t clen = len(self._callbacks)
+            size_t i
+            size_t ni = 0
+            object cb
+
+        for i from 0 <= i < clen:
+            cb = self._callbacks[i]
+            if cb != fn:
+                self._callbacks[ni] = cb
+                ni += 1
+
+        if ni != clen:
+            del self._callbacks[ni:]
+
+        return clen - ni
+
+    cpdef set_result(self, result):
+        """Mark the future done and set its result.
+
+        If the future is already done when this method is called, raises
+        InvalidStateError.
+        """
+        if self._state != _FUT_PENDING:
+            raise aio_InvalidStateError('{}: {!r}'.format(
+                self._str_state(), self))
+        self._result = result
+        self._state = _FUT_FINISHED
+        self._schedule_callbacks()
+
+    cpdef set_exception(self, exception):
+        """Mark the future done and set an exception.
+
+        If the future is already done when this method is called, raises
+        InvalidStateError.
+        """
+        if self._state != _FUT_PENDING:
+            raise aio_InvalidStateError('{}: {!r}'.format(
+                self._str_state(), self))
+        if isinstance(exception, type):
+            exception = exception()
+        if type(exception) is StopIteration:
+            raise TypeError("StopIteration interacts badly with generators "
+                            "and cannot be raised into a Future")
+        self._exception = exception
+        self._state = _FUT_FINISHED
+        self._schedule_callbacks()
+        self._log_traceback = True
+
+    # Copy of __await__
+    def __iter__(self):
+        if self._state == _FUT_PENDING:
+            self._blocking = True
+            yield self  # This tells Task to wait for completion.
+
+        if self._state == _FUT_PENDING:
+            raise AssertionError("yield from wasn't used with future")
+
+        return self._result_impl()  # May raise too.
+
+    # Copy of __iter__
+    def __await__(self):
+        if self._state == _FUT_PENDING:
+            self._blocking = True
+            yield self  # This tells Task to wait for completion.
+
+        if self._state == _FUT_PENDING:
+            raise AssertionError("yield from wasn't used with future")
+
+        return self._result_impl()  # May raise too.
+
+
+class Future(BaseFuture, aio_Future):
+    # Inherit asyncio.Future.__del__ and __repr__
+    pass
+
+
+cdef uvloop_Future = Future

uvloop/_patch.py

@@ -3,6 +3,7 @@
 import collections
 import concurrent.futures
 import functools
+import inspect
 import itertools
 import os
 import signal
@@ -11,4 +12,6 @@
 import ssl
 import sys
 import threading
+import traceback
+import time
 import warnings

uvloop/future.pyx

@@ -1,8 +1,10 @@
 import asyncio, asyncio.log, asyncio.base_events, \
-       asyncio.sslproto, asyncio.coroutines
+       asyncio.sslproto, asyncio.coroutines, \
+       asyncio.futures
 import collections
 import concurrent.futures
 import functools
+import inspect
 import itertools
 import os
 import signal as std_signal
@@ -11,11 +13,15 @@ import subprocess
 import ssl
 import sys
 import threading
+import traceback
 import time
 import warnings
 
 
+cdef aio_chain_future = asyncio.futures._chain_future
+cdef aio_get_event_loop = asyncio.get_event_loop
 cdef aio_CancelledError = asyncio.CancelledError
+cdef aio_InvalidStateError = asyncio.InvalidStateError
 cdef aio_TimeoutError = asyncio.TimeoutError
 cdef aio_Future = asyncio.Future
 cdef aio_Task = asyncio.Task
@@ -41,6 +47,7 @@ cdef cc_ThreadPoolExecutor = concurrent.futures.ThreadPoolExecutor
 cdef ft_partial = functools.partial
 
 cdef iter_chain = itertools.chain
+cdef inspect_isgenerator = inspect.isgenerator
 
 cdef int has_SO_REUSEPORT = hasattr(socket, 'SO_REUSEPORT')
 cdef int SO_REUSEPORT = getattr(socket, 'SO_REUSEPORT', 0)
@@ -80,6 +87,7 @@ cdef sys_ignore_environment = sys.flags.ignore_environment
 cdef sys_exc_info = sys.exc_info
 cdef sys_set_coroutine_wrapper = sys.set_coroutine_wrapper
 cdef sys_get_coroutine_wrapper = sys.get_coroutine_wrapper
+cdef sys_getframe = sys._getframe
 
 cdef ssl_SSLContext = ssl.SSLContext
 
@@ -93,11 +101,15 @@ cdef int signal_NSIG = std_signal.NSIG
 
 cdef time_sleep = time.sleep
 
+cdef tb_extract_stack = traceback.extract_stack
+cdef tb_format_list = traceback.format_list
+
 cdef warnings_warn = warnings.warn
 
 
 # Cython doesn't clean-up imported objects properly in Py3 mode,
 # so we delete refs to all modules manually (except sys)
 del asyncio, concurrent, collections
-del functools, itertools, socket, os, threading, std_signal, subprocess, ssl
-del time, warnings
+del functools, inspect, itertools, socket, os, threading
+del std_signal, subprocess, ssl
+del time, traceback, warnings