Python协程的底层实现

快来打我* 2022-05-21 04:06 193阅读 0赞

# 生成器 #

协程的核心就是上下文切换，在Python中最简单的实现是用生成器

生成器有个方法 `send()` 可以从调用者向生成器函数发送数据，这样就可以在生成器中 `yield future` 表示要等待 future 的结果，然后把上下文切换到调用者，等 future 结果准备好后调用者再 `send(future.result())` 给生成器发送结果，并把上下文切换到生成器函数

def generator_function():
        # 正常情况应用loop.create_future()
        result = yield asyncio.Future()
        print('future结果：', result)
        return 2
    
    
    def main():
        generator = generator_function()
        try:
            future = generator.send(None)
            # 假设某个回调调用了future.set_result
            future.set_result(1)
            future = generator.send(future.result())
        except StopIteration as e:
            print('generator_function结果：', e.value)

输出：

future结果： 1
    generator_function结果： 2

但是在生成器函数中调用子生成器会很麻烦：

def generator_function2():
        # 正常情况应用loop.create_future()
        result = yield asyncio.Future()
        print('future结果：', result)
        return 2
    
    
    def generator_function():
        generator2 = generator_function2()
        try:
            future = generator2.send(None)
            result = yield future
            future = generator2.send(result)
        except StopIteration as e:
            print('generator_function2结果：', e.value)
    
        return 3
    
    
    def main():
        generator = generator_function()
        try:
            future = generator.send(None)
            # 假设某个回调调用了future.set_result
            future.set_result(1)
            future = generator.send(future.result())
        except StopIteration as e:
            print('generator_function结果：', e.value)

输出：

future结果： 1
    generator_function2结果： 2
    generator_function结果： 3

于是有了 `yield from` 的语法糖，可以把流程控制交给子生成器，即子生成器 `yield` 直接切换上下文到调用者，调用者 `send()` 直接给子生成器发送数据。这样上面的例子可以写成：

def generator_function2():
        # 正常情况应用loop.create_future()
        result = yield asyncio.Future()
        print('future结果：', result)
        return 2
    
    
    def generator_function():
        result = yield from generator_function2()
        print('generator_function2结果：', result)
        return 3
    
    
    def main():
        generator = generator_function()
        try:
            future = generator.send(None)
            # 假设某个回调调用了future.set_result
            future.set_result(1)
            future = generator.send(future.result())
        except StopIteration as e:
            print('generator_function结果：', e.value)

输出同上

但是用生成器实现协程语义上不明确，而且不能实现异步生成器（既是协程又是生成器），于是 [PEP 492][] 提出了用 `async` `await` 作为协程语法

# async, await #

`async def` 定义的函数称为协程函数，它永远返回一个协程对象，即使函数里没有用到 `await`

`await` 后面可以跟一个 awaitable 对象，它的返回值是 awaitable 对象的结果。一个实现了 `__await__()` 方法的对象或者协程对象都是 awaitable 对象。`__await__()` 方法返回一个生成器（即这个方法是一个生成器函数），它的实现和上面的生成器协程一样， `yield future` 表示要等待future的结果。当执行协程遇到 `await` 时，流程控制交给后面的 awaitable 对象，直到最底层用 `yield future` 上下文才切换到调用者

为了使 `await` 兼容生成器实现的协程，可以用 `@asyncio.coroutine` 装饰器装饰 `yield from` 实现的协程（其实它就是给生成器函数加了个 flag `CO_ITERABLE_COROUTINE`）。生成器实现的协程返回的对象（生成器）没有 `__await__()` 方法，但它也是 awaitable 对象

协程对象和生成器一样实现了 `send(), throw(), close()` 方法，但是不可以直接迭代（ `__await__()` 方法返回的生成器可以迭代），知道这个就可以实现手动执行协程了：

# 另一种写法，用 yield from 实现的协程
    # @asyncio.coroutine
    # def coroutine_function2():
    # # 正常情况应用loop.create_future()
    # result = yield from asyncio.Future()
    # print('future结果：', result)
    # return 2
    
    
    # 用 async, await 实现的协程
    async def coroutine_function2():
        # 正常情况应用loop.create_future()
        result = await asyncio.Future()
        print('future结果：', result)
        return 2
    
    
    async def coroutine_function():
        result = await coroutine_function2()
        print('coroutine_function2结果：', result)
        return 3
    
    
    def main():
        coroutine = coroutine_function()
        # 正常情况应用asyncio.ensure_future()执行协程
        try:
            future = coroutine.send(None)
            # 假设某个回调调用了future.set_result
            future.set_result(1)
            future = coroutine.send(future.result())
        except StopIteration as e:
            print('coroutine_function结果：', e.value)

# 事件循环 #

其实事件循环本身跟协程没有什么关系，它只负责添加回调（ `call_soon, call_later` ），维护 `scheduled, ready` 队列，在有事件时调用回调而已。这里不研究了，感兴趣的可以看它的实现，大部分在 `asyncio.base_events.BaseEventLoop`

# Task #

真正实现执行协程的是 `Task`。正常情况执行一个协程用 `asyncio.ensure_future()` ，参数为协程对象，它的内部又调用了 `loop.create_task()` 创建一个 `Task` ， `Task` 的实现在 `asyncio.tasks`

`Task` 继承自 `Future` ，它的结果就是协程的返回值。

`Task` 创建时就在事件循环里添加回调开始执行协程：

def __init__(self, coro, *, loop=None):
            assert coroutines.iscoroutine(coro), repr(coro)
            super().__init__(loop=loop)
            if self._source_traceback:
                del self._source_traceback[-1]
            self._coro = coro
            self._fut_waiter = None
            self._must_cancel = False
            # 添加回调
            self._loop.call_soon(self._step)
            self.__class__._all_tasks.add(self)

`_step` 负责协程的一次迭代

def _step(self, exc=None):
            assert not self.done(), \
                '_step(): already done: {!r}, {!r}'.format(self, exc)
            if self._must_cancel:
                if not isinstance(exc, futures.CancelledError):
                    exc = futures.CancelledError()
                self._must_cancel = False
            coro = self._coro
            self._fut_waiter = None
    
            self.__class__._current_tasks[self._loop] = self
            # Call either coro.throw(exc) or coro.send(None).
            try:
                if exc is None:
                    # We use the `send` method directly, because coroutines
                    # don't have `__iter__` and `__next__` methods.
                    # 迭代一次协程，await的返回值总是future.result()，所以这里不指定发送数据也可以
                    result = coro.send(None)
                else:
                    result = coro.throw(exc)
            except StopIteration as exc:
                # 协程已经执行完毕，这里设置result或者exception
                if self._must_cancel:
                    # Task is cancelled right before coro stops.
                    self._must_cancel = False
                    self.set_exception(futures.CancelledError())
                else:
                    self.set_result(exc.value)
            except futures.CancelledError:
                # 协程被取消
                super().cancel()  # I.e., Future.cancel(self).
            except Exception as exc:
                # 其他异常
                self.set_exception(exc)
            except BaseException as exc:
                self.set_exception(exc)
                raise
            else:
                # 没有异常，result应该是一个future
                blocking = getattr(result, '_asyncio_future_blocking', None)
                if blocking is not None:
                    # Yielded Future must come from Future.__iter__().
                    if result._loop is not self._loop:
                        # 错误，result是另一个事件循环的future
                        self._loop.call_soon(
                            self._step,
                            RuntimeError(
                                'Task {!r} got Future {!r} attached to a '
                                'different loop'.format(self, result)))
                    elif blocking:
                        if result is self:
                            # 错误，await自己
                            self._loop.call_soon(
                                self._step,
                                RuntimeError(
                                    'Task cannot await on itself: {!r}'.format(
                                        self)))
                        else:
                            # 正常情况，这里给result添加一个完成时的回调self._wakeup，此协程在result完成前进入睡眠
                            result._asyncio_future_blocking = False
                            result.add_done_callback(self._wakeup)
                            self._fut_waiter = result
                            if self._must_cancel:
                                if self._fut_waiter.cancel():
                                    self._must_cancel = False
                    else:
                        # 错误，在生成器实现的协程中使用了yield而不是yield from
                        self._loop.call_soon(
                            self._step,
                            RuntimeError(
                                'yield was used instead of yield from '
                                'in task {!r} with {!r}'.format(self, result)))
                elif result is None:
                    # 正常情况，在生成器实现的协程中使用裸yield交出控制权
                    # Bare yield relinquishes control for one event loop iteration.
                    self._loop.call_soon(self._step)
                elif inspect.isgenerator(result):
                    # 错误，yield一个生成器
                    # Yielding a generator is just wrong.
                    self._loop.call_soon(
                        self._step,
                        RuntimeError(
                            'yield was used instead of yield from for '
                            'generator in task {!r} with {}'.format(
                                self, result)))
                else:
                    # 错误，yield了其他东西
                    # Yielding something else is an error.
                    self._loop.call_soon(
                        self._step,
                        RuntimeError(
                            'Task got bad yield: {!r}'.format(result)))
            finally:
                self.__class__._current_tasks.pop(self._loop)
                self = None  # Needed to break cycles when an exception occurs.

等待的 future 完成后事件循环调用 `_wakeup` 唤醒协程

def _wakeup(self, future):
            try:
                future.result()
            except Exception as exc:
                # This may also be a cancellation.
                self._step(exc)
            else:
                # Don't pass the value of `future.result()` explicitly,
                # as `Future.__iter__` and `Future.__await__` don't need it.
                # If we call `_step(value, None)` instead of `_step()`,
                # Python eval loop would use `.send(value)` method call,
                # instead of `__next__()`, which is slower for futures
                # that return non-generator iterators from their `__iter__`.
                self._step()
            self = None  # Needed to break cycles when an exception occurs.

这是用事件循环和 `Task` 执行的协程：

async def coroutine_function2():
        loop = asyncio.get_event_loop()
        future = loop.create_future()
    
        # 假设某个IO操作3秒后完成
        loop.call_later(3, lambda: future.set_result(1))
    
        result = await future
        print('future结果：', result)
        return 2
    
    
    async def coroutine_function():
        result = await coroutine_function2()
        print('coroutine_function2结果：', result)
        return 3
    
    
    def main():
        loop = asyncio.get_event_loop()
        # run_until_complete内部调用了ensure_future()
        result = loop.run_until_complete(coroutine_function())
        print('coroutine_function结果：', result)
        loop.close()

[PEP 492]: https://www.python.org/dev/peps/pep-0492/