tornado 源码之 coroutine 分析

tornado 源码之 coroutine 分析

tornado 的协程原理分析
版本:4.3.0

为支持异步,tornado 实现了一个协程库。

tornado 实现的协程框架有下面几个特点:

  1. 支持 python 2.7,没有使用 yield from
    特性,纯粹使用 yield 实现
  2. 使用抛出异常的方式从协程返回值
  3. 采用 Future 类代理协程(保存协程的执行结果,当携程执行结束时,调用注册的回调函数)
  4. 使用 IOLoop 事件循环,当事件发生时在循环中调用注册的回调,驱动协程向前执行

由此可见,这是 python 协程的一个经典的实现。

本文将实现一个类似 tornado 实现的基础协程框架,并阐述相应的原理。

外部库

使用 time 来实现定时器回调的时间计算。
bisect 的 insort 方法维护一个时间有限的定时器队列。
functools 的 partial 方法绑定函数部分参数。
使用 backports_abc 导入 Generator 来判断函数是否是生成器。

import time
import bisect
import functools
from backports_abc import Generator as GeneratorType

Future

是一个穿梭于协程和调度器之间的信使。
提供了回调函数注册(当异步事件完成后,调用注册的回调)、中间结果保存、结束结果返回等功能

add_done_callback 注册回调函数,当 Future 被解决时,改回调函数被调用。
set_result 设置最终的状态,并且调用已注册的回调函数

协程中的每一个 yield 对应一个协程,相应的对应一个 Future 对象,譬如:

@coroutine
def routine_main():
    yield routine_simple()

    yield sleep(1)

这里的 routine_simple() 和 sleep(1) 分别对应一个协程,同时有一个 Future 对应。

class Future(object):
    def __init__(self):
        self._done = False
        self._callbacks = []
        self._result = None

    def _set_done(self):
        self._done = True
        for cb in self._callbacks:
            cb(self)
        self._callbacks = None

    def done(self):
        return self._done

    def add_done_callback(self, fn):
        if self._done:
            fn(self)
        else:
            self._callbacks.append(fn)

    def set_result(self, result):
        self._result = result
        self._set_done()

    def result(self):
        return self._result

IOLoop

这里的 IOLoop 去掉了 tornado 源代码中 IO 相关部分,只保留了基本需要的功能,如果命名为 CoroutineLoop 更贴切。

这里的 IOLoop 提供基本的回调功能。它是一个线程循环,在循环中完成两件事:

  1. 检测有没有注册的回调并执行
  2. 检测有没有到期的定时器回调并执行

程序中注册的回调事件,最终都会在此处执行。
可以认为,协程程序本身、协程的驱动程序 都会在此处执行。
协程本身使用 wrapper 包装,并最后注册到 IOLoop 的事件回调,所以它的从预激到结束的代码全部在 IOLoop 回调中执行。
而协程预激后,会把 Runner.run() 函数注册到 IOLoop 的事件回调,以驱动协程向前运行。

理解这一点对于理解协程的运行原理至关重要。

这就是单线程异步的基本原理。因为都在一个线程循环中执行,我们可以不用处理多线程需要面对的各种繁琐的事情。

IOLoop.start

事件循环,回调事件和定时器事件在循环中调用。

IOLoop.run_sync

执行一个协程。

将 run 注册进全局回调,在 run 中调用 func()启动协程。
注册协程结束回调 stop, 退出 run_sync 的 start 循环,事件循环随之结束。

class IOLoop(object):,
    def __init__(self):
        self._callbacks = []
        self._timers = []
        self._running = False

    @classmethod
    def instance(cls):
        if not hasattr(cls, "_instance"):
            cls._instance = cls()
        return cls._instance

    def add_future(self, future, callback):
        future.add_done_callback(
            lambda future: self.add_callback(functools.partial(callback, future)))

    def add_timeout(self, when, callback):
        bisect.insort(self._timers, (when, callback))

    def call_later(self, delay, callback):
        return self.add_timeout(time.time() + delay, callback)

    def add_callback(self, call_back):
        self._callbacks.append(call_back)

    def start(self):
        self._running = True
        while self._running:

            # 回调任务
            callbacks = self._callbacks
            self._callbacks = []
            for call_back in callbacks:
                call_back()

            # 定时器任务
            while self._timers and self._timers[0][0] < time.time():
                task = self._timers[0][1]
                del self._timers[0]
                task()

    def stop(self):
        self._running = False

    def run_sync(self, func):
        future_cell = [None]

        def run():
            try:
                future_cell[0] = func()
            except Exception:
                pass

            self.add_future(future_cell[0], lambda future: self.stop())

        self.add_callback(run)

        self.start()
        return future_cell[0].result()

coroutine

协程装饰器。
协程由 coroutine 装饰,分为两类:

  1. 含 yield 的生成器函数
  2. 无 yield 语句的普通函数

装饰协程,并通过注册回调驱动协程运行。
程序中通过 yield coroutine_func() 方式调用协程。
此时,wrapper 函数被调用:

  1. 获取协程生成器
  2. 如果是生成器,则

    1. 调用 next() 预激协程
    2. 实例化 Runner(),驱动协程
  3. 如果是普通函数,则

    1. 调用 set_result() 结束协程

协程返回 Future 对象,供外层的协程处理。外部通过操作该 Future 控制协程的运行。
每个 yield 对应一个协程,每个协程拥有一个 Future 对象。

外部协程获取到内部协程的 Future 对象,如果内部协程尚未结束,将 Runner.run() 方法注册到 内部协程的 Future 的结束回调。
这样,在内部协程结束时,会调用注册的 run() 方法,从而驱动外部协程向前执行。

各个协程通过 Future 形成一个链式回调关系。

Runner 类在下面单独小节描述。

def coroutine(func):
    return _make_coroutine_wrapper(func)

# 每个协程都有一个 future, 代表当前协程的运行状态
def _make_coroutine_wrapper(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        future = Future()

        try:
            result = func(*args, **kwargs)
        except (Return, StopIteration) as e:
            result = _value_from_stopiteration(e)
        except Exception:
            return future
        else:
            if isinstance(result, GeneratorType):
                try:
                    yielded = next(result)
                except (StopIteration, Return) as e:
                    future.set_result(_value_from_stopiteration(e))
                except Exception:
                    pass
                else:
                    Runner(result, future, yielded)
                try:
                    return future
                finally:
                    future = None
        future.set_result(result)
        return future
    return wrapper

协程返回值

因为没有使用 yield from,协程无法直接返回值,所以使用抛出异常的方式返回。

python 2 无法在生成器中使用 return 语句。但是生成器中抛出的异常可以在外部 send() 语句中捕获。
所以,使用抛出异常的方式,将返回值存储在异常的 value 属性中,抛出。外部使用诸如:

try:
    yielded = gen.send(value)
except Return as e:

这样的方式获取协程的返回值。

class Return(Exception):
    def __init__(self, value=None):
        super(Return, self).__init__()
        self.value = value
        self.args = (value,)

Runner

Runner 是协程的驱动器类。

self.result_future 保存当前协程的状态。
self.future 保存 yield 子协程传递回来的协程状态。
从子协程的 future 获取协程运行结果 send 给当前协程,以驱动协程向前执行。

注意,会判断子协程返回的 future
如果 future 已经 set_result,代表子协程运行结束,回到 while Ture 循环,继续往下执行下一个 send;
如果 future 未 set_result,代表子协程运行未结束,将 self.run 注册到子协程结束的回调,这样,子协程结束时会调用 self.run,重新驱动协程执行。

如果本协程 send() 执行过程中,捕获到 StopIteration 或者 Return 异常,说明本协程执行结束,设置 result_future 的协程返回值,此时,注册的回调函数被执行。这里的回调函数为本协程的父协程所注册的 run()。
相当于唤醒已经处于 yiled 状态的父协程,通过 IOLoop 回调 run 函数,再执行 send()。

class Runner(object):
    def __init__(self, gen, result_future, first_yielded):
        self.gen = gen
        self.result_future = result_future
        self.io_loop = IOLoop.instance()
        self.running = False
        self.future = None

        if self.handle_yield(first_yielded):
            self.run()

    def run(self):
        try:
            self.running = True
            while True:

                try:
                    # 每一个 yield 处看做一个协程,对应一个 Future
                    # 将该协程的结果 send 出去
                    # 这样外层形如  ret = yiled coroutine_func() 能够获取到协程的返回数据
                    value = self.future.result()
                    yielded = self.gen.send(value)
                except (StopIteration, Return) as e:
                    # 协程执行完成,不再注册回调
                    self.result_future.set_result(_value_from_stopiteration(e))
                    self.result_future = None
                    return
                except Exception:
                    return
                # 协程未执行结束,继续使用 self.run() 进行驱动
                if not self.handle_yield(yielded):
                    return
        finally:
            self.running = False

    def handle_yield(self, yielded):
        self.future = yielded
        if not self.future.done():
            # 给 future 增加执行结束回调函数,这样,外部使用 future.set_result 时会调用该回调
            # 而该回调是把 self.run() 注册到 IOLoop 的事件循环
            # 所以,future.set_result 会把 self.run() 注册到 IOLoop 的事件循环,从而在下一个事件循环中调用
            self.io_loop.add_future(
                self.future, lambda f: self.run())
            return False
        return True

sleep

sleep 是一个延时协程,充分展示了协程的标准实现。

  • 创建一个 Future,并返回给外部协程;
  • 外部协程发现是一个未完的状态,将 run()注册到 Future 的完成回调,同时外部协程被挂起;
  • 在设置的延时后,IOLoop 会回调 set_result 结束协程;
  • IOLoop 调用 run() 函数;
  • IOLoop 调用 send(),唤醒挂起的外部协程。

流程如下图:

tornado 源码之 coroutine 分析

def sleep(duration):
    f = Future()
    IOLoop.instance().call_later(duration, lambda: f.set_result(None))
    return f

运行

@coroutine
def routine_ur(url, wait):
    yield sleep(wait)
    print('routine_ur {} took {}s to get!'.format(url, wait))


@coroutine
def routine_url_with_return(url, wait):
    yield sleep(wait)
    print('routine_url_with_return {} took {}s to get!'.format(url, wait))
    raise Return((url, wait))

# 非生成器协程,不会为之生成单独的 Runner()
# coroutine 运行结束后,直接返回一个已经执行结束的 future
@coroutine
def routine_simple():
    print("it is simple routine")

@coroutine
def routine_simple_return():
    print("it is simple routine with return")
    raise Return("value from routine_simple_return")

@coroutine
def routine_main():
    yield routine_simple()

    yield routine_ur("url0", 1)

    ret = yield routine_simple_return()
    print(ret)

    ret = yield routine_url_with_return("url1", 1)
    print(ret)

    ret = yield routine_url_with_return("url2", 2)
    print(ret)


if __name__ == '__main__':
    IOLoop.instance().run_sync(routine_main)

运行输出为:

it is simple routine
routine_ur url0 took 1s to get!
it is simple routine with return
value from routine_simple_return
routine_url_with_return url1 took 1s to get!
('url1', 1)
routine_url_with_return url2 took 2s to get!
('url2', 2)

可以观察到协程 sleep 已经生效。

源码

simple_coroutine.py

copyright

author:bigfish
copyright: 许可协议 知识共享署名-非商业性使用 4.0 国际许可协议

相关推荐