从零开始学asyncio(中)

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

• 从零开始学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
Host:cn.bing.com

'.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
Host:cn.bing.com

'.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
Host:cn.bing.com

'.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的对象.