tornado 原始碼之 coroutine 分析

thebigfish發表於2019-02-28
原始碼

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 國際許可協議

相關文章