从零开始学asyncio(中)

　　本篇文章主要是讲解asyncio模块的实现原理. 这个系列还有另外两篇文章:

一. asyncio模块简介

　　asyncio是python3.4开始内置的一个标准库, 可以用于编写异步的并发代码, 因此非常适合用在IO密集型操作.

　　现在运行如下代码:

import asyncio

import time

async def task(i):

    print('task{} start at {}'.format(i, time.ctime()))

    # asyncio.sleep的效果与time.sleep类似, 让程序睡眠n秒

    await asyncio.sleep(3)

    print('task{} end at {}'.format(i, time.ctime()))

tasks = asyncio.wait([task(i) for i in range(3)])

asyncio.run(tasks)

运行结果如下:

三个任务实际是处于同一线程的, 但它们的执行顺序不是start->end->start->end这种串行模式, 而是几乎同时开始, 同时结束, asyncio模块的作用就是, 使用异步的方式实现单线程并发的效果. 最简单的使用步骤如下:

首先, 在定义函数的时候使用关键字async, 这个函数就不是个普通函数了, 调用的时候不会执行内部代码, 而是返回一个coroutine对象, 即协程, 这一点与生成器函数类似.
然后, 在协程函数中的耗时操作前面加上await关键字, 注意await后面必须是可等待对象, 比如asyncio.sleep(n), 可等待对象在本文的第二节有详细的讲解.
最后, 调用asyncio.wait将协程列表打包, 打包结果给asyncio.run运行即可.

二. asyncio实现原理

　　要理解asyncio的原理, 需要理解如下几个概念: 协程, 事件循环, future/task. 其中协程就是用户自己定义的任务, 事件循环负责监听事件和回调, future/task则主要负责管理回调, 以及驱动协程.

1. 事件循环

　　事件循环负责同时对多个事件进行监听, 当监听到事件时, 就调用对应的回调函数, 进而驱动不同的任务. 上一节代码最后的asyncio.run, 其本质就是创建一个事件循环, 然后一直运行事件循环, 直到所有任务结束为止.

　　首先看看上篇文章最后的爬虫代码:

import select

import socket

import time

req = 'GET / HTTP/1.0\r\nHost:cn.bing.com\r\n\r\n'.encode('utf8')

address = ('cn.bing.com', 80)

db = []

class GenCrawler:

    '''

    这里使用一个类将生成器封装起来,如果要驱动生成器,就调用next_step方法

    另外,这个类还可以获取到使用的socket对象

    '''

    def __init__(self):

        self.sock = socket.socket()

        self.sock.setblocking(0)

        self._gen = self._crawler()

    def next_step(self):

        next(self._gen)

    def _crawler(self):

        self.sock.connect_ex(address)

        yield

        self.sock.send(req)

        response = b''

        while 1:

            yield

            chunk = self.sock.recv(1024)

            if chunk == b'':

                self.sock.close()

                break

            else:

                response += chunk

        db.append(response)

def event_loop(crawlers):

    # 首先,建立sock与crawler对象的映射关系,便于由socket对象找到对应的crawler对象

    # 建立映射的同时顺便调用crawler的next_step方法,让内部的生成器运行起来

    sock_to_crawler = {}

    for crawler in crawlers:

        sock_to_crawler[crawler.sock] = crawler

        crawler.next_step()

    # select.select需要传入三个列表,分别对应要监听的可读,可写和错误事件的socket对象集合

    readable = []

    writeable = [crawler.sock for crawler in crawlers]

    errors = []

    while 1:

        rs, ws, es = select.select(readable, writeable, errors)

        for sock in ws:

            # 当socket对象连接到服务器时,会创建可读缓冲区和可写缓冲区

            # 由于可写缓冲区创建时为空,因此连接成功时,就触发可写事件

            # 这时再转为监听可读事件,接收到数据时,就可以触发可读事件了

            writeable.remove(sock)

            readable.append(sock)

            sock_to_crawler[sock].next_step()

        for sock in rs:

            try:

                sock_to_crawler[sock].next_step()

            except StopIteration:

                # 如果生成器结束了,就说明对应的爬虫任务已经结束,不需要监听事件了

                readable.remove(sock)

        # 所有的事件都结束后,就退出循环

        if not readable and not writeable:

            break

if __name__ == '__main__':

    start = time.time()

    n = 10

    print('开始爬取...')

    event_loop([GenCrawler() for _ in range(n)])

    print('获取到{}条数据,用时{:.2f}秒'.format(len(db), time.time()-start))

这段代码使用IO多路复用对多个socket进行监听, 监听到事件时, 驱动对应的生成器运行, 运行到IO操作时, 再使用yield切换回事件循环, 从而实现并发的效果, 这个也就是asyncio中事件循环的工作原理.

　　由于asyncio中的事件循环使用的是selectors模块而非select, 现在在程序的代码中改用selectors模块:

import socket

import time

from selectors import DefaultSelector, EVENT_READ, EVENT_WRITE

req = 'GET / HTTP/1.0\r\nHost:cn.bing.com\r\n\r\n'.encode('utf8')

address = ('cn.bing.com', 80)

db = []

class EventLoop:

    def __init__(self):

        self.selector = DefaultSelector()

        self._stopped = False

    def register(self, fd, event, callback):

        self.selector.register(fd, event, callback)

    def unregister(self, fd):

        self.selector.unregister(fd)        

    def run_until_complete(self,gens):

        for gen in gens:

            next(gen)

        while not self._stopped:

            try:

                events = self.selector.select()

            except OSError:

                # 如果当前没有注册事件, 就会引发OSError异常

                continue

            for key, mask in events:

                # 这里的callback就是注册事件时传入的回调函数

                callback = key.data

                callback(key=key, mask=mask)

            # 生成器的gi_frame属性对应的是其框架(其实这属性我还没搞懂)

            # 在生成器结束(抛出stopiteration异常)后,这个属性值就会变成None

            # 因此,每次循环时都删减已经结束的生成器

            # 如果所有的生成器都结束了,就停止循环

            gens = [gen for gen in gens if gen.gi_frame is not None]

            if not gens:

                self.stop()

    def stop(self):

        self._stopped = True

        self.selector.close()

loop = EventLoop()

class GenCrawler:

    def __init__(self):

        self.sock = socket.socket()

        self.sock.setblocking(0)

        self._fd = self.sock.fileno()

        self.gen = self._crawler()

    def _crawler(self):

        self.sock.connect_ex(address)

        loop.register(self._fd, EVENT_WRITE, self.next_step)

        yield

        loop.unregister(self._fd)

        self.sock.send(req)

        response = b''

        while 1:

            loop.register(self._fd, EVENT_READ, self.next_step)

            yield

            loop.unregister(self._fd)

            chunk = self.sock.recv(1024)

            if chunk == b'':

                self.sock.close()

                break

            else:

                response += chunk

        db.append(response)def next_step(self,**kwargs):

        try:

            next(self.gen)

        except StopIteration:

            return

if __name__ == '__main__':

    start = time.time()

    print('开始爬取...')

    n = 10

    gens = [GenCrawler().gen for _ in range(n)]

    loop.run_until_complete(gens)

    print('获取到{}条数据,用时{:.2f}秒'.format(len(db), time.time()-start))

这里主要是改了EventLoop部分的代码, 使用register和unregister方法来注册和注销事件, 优点是更加灵活, 可以指定触发事件时调用的回调函数. 另外, DefaultSelector会自动选择系统中效率最高的多路复用机制, 比如kqueue和epoll.

2. async与协程

　　在定义函数的时候, 在def之前加上async, 这个函数就不是普通函数了, 而是一个协程函数:

async def coro():

    print('this is a coroutine')

直接调用协程函数并不能使之运行, 而是返回了一个协程对象, 如果要运行该协程, 可以调用这个协程对象的send方法:

c=coro()

c.send(None)

运行结果如下, 首先会运行协程函数内部的代码, 然后函数的代码运行结束, 抛出一个StopIteration异常:

　　因此, 协程函数与生成器函数是非常相似的. 但是, 协程不是可迭代对象, 因此无法使用next函数, 只能调用其自身的send方法来驱动.

从python3.6开始, 协程函数中可以使用yield语句, 此时调用这个函数, 就会返回一个async_generator对象, 即异步生成器.

不过这东西我还没用过, 先挖个坑, 需要的可以看PEP525.

补充说明

3. await和awaitable

　　在第一节中讲到, 协程中可以使用await语句, 后接awaitable对象, 即可等待对象. 以下几类都是可等待对象:

一个协程, 这个上一小节刚讲.
一个有__await__方法, 并且该方法返回一个迭代器的对象, 常见情况是这个对象的__await__方法就是个生成器函数.
使用@types.coroutine装饰的生成器函数, 其中types是python内置的一个库, 这个装饰器的实现原理是返回一个定义了__await__方法的对象.
Objects defined with CPython C API with a tp_as_async.am_await function, returning an iterator (similar to __await__ method). 这一条是从官网抄的, 我没理解, 应该和定义__await__类似吧.

现在定义一个可等待对象并测试:

class AwaitableObj:

    def __await__(self):

        v = yield '来自可等待对象的yield'

        print('可等待对象获得的值:', v)

        return '来自可等待对象的return'

async def coro():

    v = await AwaitableObj()

    print('协程获得的值:', v)

if __name__ == '__main__':

    c = coro()

    v = c.send(None)

    print('外部获得的值:', v)

    try:

        c.send('来自外部')

    except StopIteration:

        pass

这段程序有三个部分: 可等待对象, 协程和外部. 协程中使用await语句来等待可等待对象, 而外部调用send方法来驱动协程.

程序的运行结果如下:

await相当于外部与可等待对象之间的桥梁, 可等待对象中__await__方法返回的生成器, 其yield返回的值会传到外部, 而外部使用send方法传的值也会传给可等待对象的生成器. 最后__await__生成器迭代结束后, 协程获得其返回值.

　　这里需要说明一点: await语句本身并不能暂停和切换协程, 它只是阻塞协程直到后面接的可等待对象的__await__方法返回的可迭代对象运行完. 如果__await__里面有yield, 返回一个生成器, 协程才会因为这个yield语句暂停和切换.

4. future/task

　　future是asyncio模块中的一个可等待对象, 调用asyncio.get_event_loop获取到当前线程的事件循环loop, 然后调用loop.create_future, 就可以得到一个future对象. future的主要代码如下(有改动):

class Future:

    def __init__(self):

        self._callbacks = []

        self.result = None

    def add_callback(self, callback):

        self._callbacks.append(callback)

    def set_result(self, result):

        self.result = result

        for callback in self._callbacks:

            callback(self)

    def __await__(self):

        yield self

        return self.result

　　future可以理解为协程的一次暂停. 首先, 如果一个协程需要在某处暂停, 就可以实例化一个future对象并且await这个对象, 这样就会运行future对象的__await__方法, 当运行到yield self这句话时, 协程暂停, 直到外部再使用send方法驱动协程为止. 然后, future的另一特性是可以设置回调函数, 调用它的add_callback方法就行. 最后, future还有set_result这个接口, 一方面会运行future的回调函数, 另一方面可以设置其result属性的值, 该值在__awaiit__方法结束之后返回给协程. 一般的用法是, 协程在事件循环中注册事件, 然后让事件循环来调用future对象的set_result方法.

　　有了暂停, 自然也需要有驱动, task对象负责对协程进行封装和驱动. 调用asyncio.create_task并传入协程对象, 就可以得到一个task对象. task的主要代码如下(有改动):

class Task(Future):

    def __init__(self, coro):

        super().__init__()

        self.coro = coro

        f = Future()

        f.set_result(None)

        self.step(f)

    def step(self, future):

        try:

            next_future = self.coro.send(None)

        except StopIteration:

            self.set_result(future.result)

            return

        next_future.add_callback(self.step)

　　task和future应该是搭配使用的. 首先, task.step是负责对协程进行驱动的, 由于future.__await__方法会yield self, 因此每次驱动都会获得目前暂停点对应的future对象. 这时候将自己的step方法添加到future对象的回调中, 等到future对象调用set_result方法时, 就会回调到task.step方法, 从而驱动协程继续运行. 因此可以认为, future对象就是协程的一次暂停, 而调用其set_result方法就意味着这次暂停结束了, 但是这个过程需要task的协助.

　　task类是继承future类的, 这其实比较好理解, 比如一个简单的爬虫任务, 在连接服务器和接受数据等IO操作时需要使用future暂停, 并可以设置回调, 表示暂停结束的时候应该做什么. 而这个爬虫任务相当于一个大的IO操作, 因此也应该有可以设置回调以及可以await的特性. 当一个协程驱动结束, 即抛出StopIteration异常的时候, 就意味着这个task结束了, 因此此时就调用task.set_result方法, 把最后一个future对象的结果设置为task.result.

5. 爬虫代码重构

　　把上面讲的async/await, future/task等内容添加到之前的爬虫实例中, 最终代码如下:

import socket

import time

from selectors import DefaultSelector, EVENT_READ, EVENT_WRITE

req = 'GET / HTTP/1.0\r\nHost:cn.bing.com\r\n\r\n'.encode('utf8')

address = ('cn.bing.com', 80)

db = []

class EventLoop:

    def __init__(self):

        self.selector = DefaultSelector()

        self._stopped = False

    def register(self, fd, event, callback):

        self.selector.register(fd, event, callback)

    def unregister(self, fd):

        self.selector.unregister(fd)

    def run_until_complete(self, coros):

        def _done_callback(fut):

            nonlocal ncoros

            ncoros -= 1

            if ncoros == 0:

                self.stop()

        ncoros = len(coros)

        for coro in coros:

            task = Task(coro)

            task.add_callback(_done_callback)

        while not self._stopped:

            try:

                events = self.selector.select()

            except OSError:

                # 如果当前没有注册事件, 就会引发OSError异常

                continue

            for key, mask in events:

                # 这里的callback就是注册事件时传入的回调函数

                callback = key.data

                callback(key=key, mask=mask)

    def stop(self):

        self._stopped = True

        self.selector.close()

loop = EventLoop()

class Future:

    def __init__(self):

        self._callbacks = []

        self.result = None

    def add_callback(self, callback):

        self._callbacks.append(callback)

    def set_result(self, result):

        self.result = result

        for callback in self._callbacks:

            callback(self)

    def __await__(self):

        yield self

        return self.result

class Task(Future):

    def __init__(self, coro):

        super().__init__()

        self.coro = coro

        f = Future()

        f.set_result(None)

        self.step(f)

    def step(self, future):

        try:

            next_future = self.coro.send(None)

        except StopIteration:

            self.set_result(future.result)

            return

        next_future.add_callback(self.step)

class CoroCrawler:

    def __init__(self):

        self.sock = socket.socket()

        self.sock.setblocking(0)

        self._fd = self.sock.fileno()

        self.coro = self._crawler()

    async def _crawler(self):

        await self.connect()

        self.sock.send(req)

        response = await self.read_all()

        db.append(response)

    async def connect(self):

        self.sock.connect_ex(address)

        f = Future()

        def on_connect(key, mask):

            f.set_result(None)

        loop.register(self.sock.fileno(), EVENT_WRITE, on_connect)

        await f

        loop.unregister(self.sock.fileno())

    async def read_all(self):

        response = b''

        while 1:

            chunk = await self.read()

            if chunk == b'':

                self.sock.close()

                break

            response += chunk

        return response

    async def read(self):

        f = Future()

        def on_readable(key, mask):

            chunk = self.sock.recv(1024)

            f.set_result(chunk)

        loop.register(self._fd, EVENT_READ, on_readable)

        chunk = await f

        loop.unregister(self._fd)

        return chunk

if __name__ == '__main__':

    start = time.time()

    print('开始爬取...')

    n = 10

    coros = [CoroCrawler().coro for _ in range(n)]

    loop.run_until_complete(coros)

    print('获取到{}条数据,用时{:.2f}秒'.format(len(db), time.time()-start))

　　这段代码并不算复杂, 唯一需要留意的就是事件循环中的run_until_complete方法, 这个方法不再是主动去检查任务是否结束, 而是将协程包装成task对象, 然后给task对象添加了回调函数, 来在协程全部结束时, 停止事件循环. 这也就是用task包装协程的一个方便的地方: 可以在协程结束的时候运行指定的回调.

　　整个代码的实现流程如下, 这也就是用asyncio运行一个协程的流程.

三. 总结

asyncio模块是基于async/await实现的一个异步库, 用法的话, 简单来说就是首先定义好协程, 然后把协程打包扔给asyncio模块, 最后启动事件循环, 就实现异步了.
asyncio底层使用事件循环, 其本质是系统的IO多路复用机制, 通过这种机制来同时监听多个对象, 在触发事件时调用对应的回调函数, 从而实现异步和并发的效果.
Future表示协程单个断点的运行状态, Task继承自Future, 表示整个协程的运行状态, 二者都可以设置回调函数, 在结束的时候调用回调.
Task是对协程对象的封装和管理, 负责驱动协程, Future则直接对接loop循环, 接收回调函数的结果并返回.
要定义协程, 首先要使用async定义函数, 然后如果有耗时操作, 在耗时操作前面加上await. 不过, 对应的耗时操作必须是awaitable的对象.