Python asyncio库核心源码解析

Python asyncio库是由Python之父Guido亲自主持开发的Python异步I/O库，Python3.6之后已正式成为标准库中的一员，其提供了async/await语法支持原生协程，使得在Python中进行异步编程变得非常简单。

阅读这篇文章之前最好先阅读文章深入理解Python异步编程，了解一下Python异步编程的发展过程。早期，Python是基于生成器和事件循环来实现的异步编程，当前的async/await原生协程的内部实现机制和早期的基于生成器的实现方案其实也是基本一致的，不过asyncio库还实现了很多其他功能，比如异常处理、任务的状态管理和调度等等。

Python asyncio库的核心组件包括EventLoop、Coroutine、Task、Future、Handle等，我们接下来对这些核心组件进行逐一解析，我们可以在Python标准库的Lib/asyncio目录下找到对应的源码。

1. 入门实例

---

我们先使用asyncio库编写一个简单的入门实例，然后以该实例作为代码入口梳理asyncio库的核心代码逻辑。

# -*- coding: utf-8 -*-
import asyncio


async def cal(n):
    print("cal: ", n)
    await asyncio.sleep(3)
    print("done: ", n)
    return n


if __name__ == "__main__":
    res = asyncio.run(cal(5))
    print("result: ", res)

这里asyncio.run(coro)就是程序的入口，传入一个协程对象，被async修饰的函数即为协程，程序的入口还可以写成如下的这种方式，直接使用事件循环对象。

if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    res = loop.run_until_complete(cal(5))
    print("result: ", res)

这两种方式其实是一样的，asyncio.run()方法就是对上面的代码做了一个包装。

2. asyncio.run()

---

首先点击进入asyncio.run()方法，详细代码如下，我们在阅读源码的时候可以先忽略多余的代码，只关注关键的步骤，这样才能快速将代码的核心逻辑梳理清楚。

这里首先创建了一个新的事件循环EventLoop，然后调用该事件循环的run_until_complete方法执行该协程。

注意：该入口程序只能调用一次，因为每次调用都会尝试生成一个新的EventLoop，代码中每次只允许有一个EventLoop运行。

接着我们进入run_until_complete方法，我们一起来看上面的代码，主要包含4块操作：

1. 首先将传入的参数使用tasks.ensure_future()方法转换成一个awaitable的future对象，如果传入的是协程，会调用loop.create_task()方法生成一个Task对象，Task也是Future的子类。这里为什么要转换成Task对象呢？我们后面在讲Task类的时候再详细讲解；
2. 第二步将_run_until_complete_cb方法注册为future的回调方法，当future执行完成时会回调该方法，将当前EventLoop设置为结束状态，这样EventLoop在当前循环执行完成之后就会退出；
3. 接着就是调用run_forever()进入当前EventLoop的主循环中，这块也会在后面进行详细讲解；
4. 该协程执行完毕，当前EventLoop的主循环已跳出，在最后通过future.result()返回协程的执行结果。

3. Future类

---

我们先来看下asyncio库中核心类之一，Future类其源码位于asyncio.futures模块中，和之前介绍的基于生成器的协程的Future基本一致，这里就不再贴源码了。

Future类包含的属性主要有三个：_loop属性表示所属的EventLoop，_callbacks列表属性保存添加到该future对象中的回调方法，_result属性保存该任务的执行结果。

Future类包含的方法主要有：add_done_callback方法用于注册回调，set_result方法用于在任务执行完成之后设置执行结果，并且会调用__schedule_callbacks()方法将注册的回调方法都依次执行一遍。这里执行回调方法并不是直接调用，而是通过call_soon方法在EventLoop中注册一个handle，这样在EventLoop的最近一次循环中就会执行该回调方法。

最后就是Future类还实现了__await__()方法，使得future对象awaitable，可用await修饰符修饰，在该方法中将future对象自身通过yield返回。

那future对象的set_result()方法是谁在什么时候调用的呢？参考上面的入门实例，我们点击进入sleep方法的源码，如下所示，在sleep方法中首先创建了一个future对象，然后通过EventLoop的call_later注册了一个回调方法，最后在任务结束之后返回future的执行结果。EventLoop在经过delay时间之后会调用future的_set_result_unless_cancelled方法，该方法中会调用future的set_result()方法设置任务执行结果。

这里的sleep方法是asyncio库内部实现的方法，如果是我们自己来实现异步操作，那么可以调用EventLoop的add_writer或者add_reader方法，向EventLoop中的selector注册读或写事件，当事件发生时会执行我们注册的回调方法。比如aiohttp库中就是调用的这两个方法注册读写事件，实现了异步的HTTP请求。

4. Task类

---

每个协程内部可能包含若干个Future对象，程序执行到每个await future时如果future尚未完成，那么程序就会发生中断，等到任务执行完成设置结果到future之后，程序才会从await future的地方继续往后执行，直到协程的所有内部逻辑执行完成。那么整个执行流程是通过什么串联起来的呢？这就是Task对象所起的作用了。

Task类的源码位于asyncio.tasks模块中。每个Task类都会包裹一个协程coro，用于驱动该协程中的Future对象依次执行。Task类中最重要的__step方法就是用于驱动协程一步一步执行的，由于源码中包含大量非核心代码，我们这里只提取出关键代码，如下：

class Task(Future):
    ......
    def __step(self, exc=None):
        try:
            result = coro.send(None)
        except StopIteration as exc:
            return
        result.add_done_callback(self.__wakeup, context=self._context)

在__step方法中通过调用core.send()触发协程开始执行或者从上一次await的地方继续执行，直到遇到下一个await future，将future对象返回回来，保存在result中，然后将__wakeup方法注册为该future对象的回调方法，当future完成之后就会继续回调这里__wakeup方法，__wakeup方法内部其实就是再次调用了一次__step方法。如此循环直到协程执行完成，抛出StopIteration异常，该Task执行完成。

那协程的第一次是怎么触发起来的呢？我们在Task类的构造方法__init__中会发现代码调用了EventLoop的call_soon方法，代码如下。将该任务的step方法添加到了EventLoop的执行ready队列中去，在EventLoop的下一次循环就会调用该方法触发协程的执行。

到此，协程的执行流程就大概梳理清楚了，是通过Task、Future、EventLoop三者配合驱动协程一步一步执行。

5. EventLoop

---

最后我们再来详细地看下核心组件EventLoop的内部实现。

我们再次回到EventLoop的run_forever方法，去掉多余的代码，核心逻辑只剩下4行，如下所示，该方法会一直循环执行_run_once方法，直到EventLoop的_stopping属性被设置为False就跳出循环。

def run_forever():
    while True:
        self._run_once()
        if self._stopping:
            break

那么_run_once方法是干什么的呢？我们一起进入该方法的源码。由于该方法较长，这里就不贴源码了，多读几遍就会发现该方法主要实现了如下几个步骤：

1. _scheduled是EventLoop用于存放待调度的任务队列，第一步会先将该队列进行建堆操作，这样可以按照任务优先级进行任务调度；
2. 找出最优先执行的任务的等待时间timeout，然后调用event_list = self._selector.select(timeout)等待事件的发生或者超时；
3. 调用self._process_events(event_list)方法处理返回的事件列表，该方法中会依次将每个事件对应的回调方法添加到EventLoop的_ready任务队列中；
4. 有的回调并不是通过向selector中注册事件的方式，而是根据时间直接添加到任务队列中的，所以还需要将_scheduled队列中的满足执行时间要求的任务添加到_ready队列中，_ready队列存放的即为当前循环需要执行的任务；
5. 针对_ready队列中的每个handle任务均执行handle._run()方法，触发回调函数的调用。Handle对象除了封装回调函数及其方法参数，还维护了该任务的状态等，调用handle对象的_run方法即可执行回调。

还有最后一个问题，EventLoop任务队列中的任务是怎么添加进去的呢？EventLoop中维护了两个队列，_scheduled和_ready，前者存放未来某个时间会执行的任务，通过EventLoop的call_later、call_at等方法添加任务，后者存放了最近一次循环会执行的任务，可通过call_soon等方法直接往其中添加任务。EventLoop还支持通过add_reader和add_writer方法向selector中注册读写事件和回调，事件发生时就会触发回调方法的调用。

至此，EventLoop的核心业务逻辑也梳理清楚了，

6. 总结

---

通过分析asyncio库的核心源码，我们梳理清楚了asyncio库的EventLoop、Task、Future等几个核心对象之间的关系，以及这几个对象是如何配合驱动协程一步一步执行的，整个思路和基于生成器的协程实现基本类似。