Tornado源码分析（二）异步上下文管理（StackContext）

异步异常与上下文

在Python黑魔法---上下文管理器最后关于上下文的使用，提到了tornado的处理方式。本篇就来一探究竟。回顾问题，异步函数执行的时候，抛出的异常已经和主函数的上下文不一致，为了解决这个问题，可以使用Python的上下文管理器进行wrapper。下面的代码，就存在异步异常在主函数中无法捕获的问题：

import tornado.ioloop
import tornado.stack_context

ioloop = tornado.ioloop.IOLoop.instance()

times = 0

def callback():
    print 'run callback'
    raise ValueError('except in callback')

def async_task():
    global times
    times += 1
    print 'run async task {}'.format(times)
    ioloop.add_callback(callback=callback)


def main():
    try:
        async_task()
    except Exception as e:
        print 'main exception {}'.format(e)
    print 'end'

运行上述代码将会返回：

run async task 1
end
run callback
ERROR:root:Exception in callback <tornado.stack_context._StackContextWrapper object at 0x10306f890>
Traceback (most recent call last):
  ...
    raise ValueError('except in callback')
ValueError: except in callback

async_task函数执行的时候，在注册了一个异步回调函数callback。可是在async_task的异常try逻辑中，callback抛出的异常无法正确的catch。也就是终端并没有输出main exception except in callback，而是仅仅输出了except in callback的异常。

初次解决

因为主函数无法捕获回调的异常，同时为了防止回调的异常蔓延到主函数，一个简单的思路就是在callback中进行try捕获。修改代码如下：

def callback():
    print 'run callback'
    try:
        raise ValueError('except in callback')
    except Exception as e:
        print 'main exception {}'.format(e)

运行结果如下：

run async task 1
end
run callback
main exception except in callback

看起来不错，在callback中写入了main函数的捕获逻辑。问题算是解决了。可是，这样的做法相当丑陋。如果主函数里针对callback异常还有别的业务逻辑，那么这样的写法就很死，甚至无法完成接下来的逻辑。

包裹上下文

针对主函数无法catch，初次尝试把catch移步到callback中。这样的问题是涉及主函数逻辑会写死。如果异步的try作为一个包裹，而不是语法修改，会不会更好呢？写个 callback代码如下：

def callback():
    print 'run callback'
    raise ValueError('except in callback')

def wrapper(func):
    try:
        func()
    except Exception as e:
        print 'main exception {}'.format(e)

def async_task():
    global times
    times += 1
    print 'run async task {}'.format(times)
    ioloop.add_callback(callback=functools.partial(wrapper, callback))

def main():
    wrapper(async_task)

运行之后，发现主函数可以catch callback中的异常了。这样做的思路其实很简单，因为callback会产生异常，并且这个异常需要蔓延传播到主函数，那么我们就挖一个坑，这个坑分别包裹callback和主函数，因为坑都是一样的，所有raise的异常可以定义在坑中。

wrapper原理.png

灵活性变大了，当然，这样做还是有限制，比如主函数需要另外一种坑，如果定义多个坑，那么还得修改 async_task中的wrapper，比较好的方式是在主函数可以动态的传递wrapper函数。这就涉及到全局变量。可以使用全局的字段存储多个不同的wrapper函数坑。

times = 0

GLOBAL_WRAPPERS = {}

def callback():
    print 'run callback'
    raise ValueError('except in callback')

def wrapper(func):
    try:
        func()
    except Exception as e:
        print 'wrapper exception {}'.format(e)

def other_wrapper(func):
    try:
        func()
    except Exception as e:
        print 'other_wrapper exception {}'.format(e)

def async_task():
    global times
    times += 1
    print 'run async task {}'.format(times)
    ioloop.add_callback(callback=functools.partial(GLOBAL_WRAPPERS['context'], callback))

def main():
    GLOBAL_WRAPPERS['context'] = wrapper
    wrapper(async_task)

    GLOBAL_WRAPPERS['context'] = other_wrapper
    other_wrapper(async_task)

定义了一个全局变量，用于保存不同的函数坑，其实这个坑可以理解为函数执行的上下文。变换不同的上下文，异步callback也会跟着进入对应的上下文。这种技巧，tornado的stack_context用到了极致，相当巧妙。

tornado stack_context 源码

对于stack_context的分析，主要采用tornado2.0的代码例子。tornado的源码附带的测试样例非常棒，不过我们还是写一个简单的使用stack_context的代码，然后再一步步看程序的执行。

times = 0

def callback():
    print 'Run callback'
    raise ValueError('except in callback')

def async_task():
    global times
    times += 1
    print 'run async task {}'.format(times)
    ioloop.add_callback(callback=callback)

@contextlib.contextmanager
def contextor():
    print 'Enter contextor'
    try:
        yield
    except Exception as e:
        print 'Handler except'
        print 'exception {}'.format(e)
    finally:
        print 'Release'

def main():
    with tornado.stack_context.StackContext(contextor):
        async_task()
    print 'End'

运行结果如下：

Enter contextor
run async task 1
Release
End
Enter contextor
Run callback
Handler except
exception except in callback
Release

从输出来看：

首先进入contextor上下文管理器上下文
执行 async 函数
退出contextor上下文管理器上下文
再次进入contextor上下文管理器上下文
执行异步的callback
callback产生异常，执行 contextor上下文管理器的异常处理代码
再次退出contextor上下文管理器上下文

所有上述的步骤，正如前面的分析，无论是主函数还是异步回调函数，都经过了stack_context的包裹（挖的坑），实现了上下文切换执行代码。具体而言，在我们的代码的with语句进行了一次包裹，ioloop.add_callback则进行了对回调的包裹。

创建Stack_context 上下文管理器

在main函数中，首先创建了Stack_context上下文管理器，然后通过with语句进入contextor上下文

def main():
     stack_context = tornado.stack_context.StackContext(contextor)
    with stack_context:
        async_task()
    print 'End'

在 stack_context.py 文件中，实例StackContext的时候，将上下文管理contextor注入其中，然后调用 with语句的时候，执行StackContext的 __enter__方法：

class StackContext(object):
    
    def __init__(self, context_factory):
            # 将上下文管理函数传到StackContext
        self.context_factory = context_factory
        
    def __enter__(self):
            # 存储旧的状态上下文
        self.old_contexts = _state.contexts
        # _state.contexts 是一个元组的结果，为StackContext和上下文管理函数 (class, arg) 这样的结构，下面就是更新 _state.contexts
        _state.contexts = (self.old_contexts + 
                           ((StackContext, self.context_factory),))
        
        try:
                 # self.context_factory 是传递进来的上下文管理函数（contextor），通过调用self.context_factory创建上下文管理器。
            self.context = self.context_factory()
            # 调用上下文管理器的__enter__ 方法，进入contextor上下文环境
            self.context.__enter__()
        except Exception:
            _state.contexts = self.old_contexts
            raise

上述代码注释解释了大部分逻辑，需要额外注意是这个 _state.context。它是一个python线程的全局变量（theading.local），其职能类似GLOBAL_WRAPPER用于保存不同的上下文。他的特点就是每个线程都能把自己的私有数据写入，同时对于别的线程又是隔离不可见。一旦执行了self.context.__enter__()代码，函数控制上下文将会转移到上下文管理器(contextor)的__enter__方法中：

def contextor():
     # StackContext 中执行 self.context = self.context_factory()将会转移到此
    print 'Enter contextor'
    try:
        yield
    except Exception as e:
        print 'Handler except'
        print 'exception {}'.format(e)
    finally:
        print 'Release'

此时可以看到控制台输出 'Enter contextor'的输出，同时被yield，函数控制权回到StackContext中的enter。

注册回调函数

接下来，进入到with语句后，__enter__方法返回后，执行async_task函数，而async_task调用了ioloop.add_callback(callback=callback)。下面来看里面的代码：

    def add_callback(self, callback):
        if not self._callbacks:
            self._wake()
        # 将callback传递给stack_context，返回一个_StackContextWrapper对象，该其中保存了callback和aysnc_task的上下文对象元组(StackContext, contextor)
        self._callbacks.append(stack_context.wrap(callback))

add_callback 会针对管道进行一下处理，具体放到ioloop再讨论，这里只需要了解callback又被stack_context包裹了，并且注册到ioloop实例的_callbacks列表里。

下面在看看这个wrap干了什么事情：

def wrap(fn):
    if fn is None or fn.__class__ is _StackContextWrapper:
        return fn
    def wrapped(callback, contexts, *args, **kwargs):
        ...
    return _StackContextWrapper(wrapped, fn, _state.contexts)

首先判断包裹的函数（callback）是否为None，并且他是否已经被_StackContextWrapper包裹了，如果满足上面的条件，就直接返回。否则则进行_StackContextWrapper包裹。_StackContextWrapper其实就是一个偏函数functools.partial。这里需要注意的是 wreapped函数（稍后会用到），fn（被包裹的callback），状态上下文 _state.contexts。 _state.contexts就是之前 Stack_context.enter方法中创建的那个（class，args）元组。这样的做法，就是为了后面包裹回调函数的上下文环境保存起来。此时的_state.contexts是一个 StackContext和contextor的元组对，将会在wrapper函数中进行再一次包裹：即StackContext(contextor)。

管理回调函数上下文

stack_context.wrap函数执行返回后，将会退出包裹contextor的上下文，即调用StackContext的 __exit__方法：

    def __exit__(self, type, value, traceback):
        try:
            return self.context.__exit__(type, value, traceback)
        finally:
                # 将全contextor的上下文出栈
            _state.contexts = self.old_contexts

该__exit__中会执行self.context的__exit__方法，即contextor中的finnaly，此时会打印出 Release。

@contextlib.contextmanager
def contextor():
    print 'Enter contextor'
    try:
        yield
    except Exception as e:
        print 'Handler except'
        print 'exception {}'.format(e)
    finally:
        print 'Release'

StackContext的finally还会把刚执行完毕的全局上下文出栈，即恢复到StackContext.wrapper(contextor)之前的上下文。

执行callback

出现异常的逻辑在callback，到目前为止，还没有执行callback函数。从上面的经验可以看出，想要执行callback，首先需要上下文管理器包裹一下callback，然后进入callback上下文，执行callback，触发异常，进入callback的exit上下文。当然，无论是之前的对contextor的wrapper还是接下来对callback的wrapper，都是用的同一个上下文管理器 contextor。

继续代码的执行，将会运行到 ioloop.start方法


callbacks = self._callbacks
self._callbacks = []
for callback in callbacks:         
    self._run_callback(callback)

然后是在_run_callback中执行 callback()函数。

    def _run_callback(self, callback):
        try:
            callback()   # 此时成callback是一个被StackContext.wrap包裹的_StackContextWrappe对象。即可以通过contextor创建上下文环境，该上下文环境与async_task的一致
        except (KeyboardInterrupt, SystemExit):
            raise
        except:
            self.handle_callback_exception(callback)

注意此时的callback，并不是定义的callback，而是经过StackContext包裹的callback，具体在StackContext.wrap(callback)调用的时候，返回了偏函数的_StackContextWrapper 对象。因此调用_StackContextWrappe(),进入下面的StackContext.wrap函数的逻辑

def wrap(fn):
    
    '''
    if fn is None or fn.__class__ is _StackContextWrapper:
        return fn
   
    def wrapped(callback, contexts, *args, **kwargs):
            # 判断当前上下文（cls， args（contextor））是否在全局中保存。对于没有嵌套的StackContext.wrap，此时的条件不成立。如果是嵌套包裹，此时就直接调用callback。
        if contexts is _state.contexts or not contexts:
            callback(*args, **kwargs)
            return
        # 将 StackContext和contextor进行包裹
        if not _state.contexts:
            new_contexts = [cls(arg) for (cls, arg) in contexts]
        elif (len(_state.contexts) > len(contexts) or
            any(a[1] is not b[1]
                for a, b in itertools.izip(_state.contexts, contexts))):
            # contexts have been removed or changed, so start over
            new_contexts = ([NullContext()] +
                            [cls(arg) for (cls,arg) in contexts])
        else:
            new_contexts = [cls(arg)
                            for (cls, arg) in contexts[len(_state.contexts):]]
                            
        if len(new_contexts) > 1:
            with _nested(*new_contexts):
                callback(*args, **kwargs)
        elif new_contexts:
                # 再一次使用 StackContext包裹一个上下文处理器 contextor
            with new_contexts[0]:
                # 将callback在被StackContext包裹contextor执行callback
                callback(*args, **kwargs)
        else:
            callback(*args, **kwargs)
    return _StackContextWrapper(wrapped, fn, _state.contexts)

上述代码很多，其实目前只需要关注new_contexts = [cls(arg) for (cls, arg) in contexts]和with new_contexts[0]:callback(*args, **kwargs)两个逻辑。

cls(arg)的做法，与main函数中的stack_context = tornado.stack_context.StackContext(contextor)。一模一样。创建一个创建Stack_context 上下文管理器。至于with new_contexts则与StackContext.wrapper(connextor)的效果一致。进入contextor上下文环境，然后执行callback，此时进入上下文管理器的时候，也会打印 Enter contextor。然后就是真正的执行callback回调函数。因为发生异常，就触发了contextor的__exit__方法，然后执行了print 'exception {}'.format(e)代码，最后退出contextor上下文环境。完成callback的调用。

回顾

如果一步步debug，还是很容易弄清楚StackContext的原理，写成文字，反而说不清。现在我们再分析代码输出结果

1. Enter contextor
2. run async task 1
3. Release
4. End
5. Enter contextor
6. Run callback
7. Handler except
8. exception except in callback
9. Release

1 StackContext(contextor)实例化创建上下文管理器，然后通过with语句调用，进入了contextor的 __enter__方法所打印输出
2 进入with上下文环境，调用 async_task输出，同时ioloop注册回调函数。通过stack_context.wrap(callback)注册并保存与async_task上下文一样的管理器，并使用_StackContext偏函数返回
3 退出with代码块，执行contextor.exit方法输出
4 主函数main继续执行打印
5 ioloop继续执行，调用callback回调，此时的callback是_StackContextWrapper对象，_StackContextWrapper调用 wrapper函数内逻辑，通过cls(args)创建一个新的上下文管理器，并通过with new_contexts[0]进入上下文管理器。
6 进入 callback函数执行
7 产生异常，触发新创建的上下文管理器的exit中的异常处理
8 输出异常
9 执行上下文管理器的finnaly分支，退出上下文管理器。

其中 2 步骤是处理上下文管理器的基础，5则是aync_task和callback上下文管理器包裹同步的关键。

大概流程图如下：

流程.png

总而言之，async_task和callback的执行上下文本来不一样。为了解决问题，定义一个上下文管理器contextor。无论再调用async_task还是callback之前，先用StackContext管理contextor。初始执行async_task和callback函数逻辑的时候，都在contextor上下文环境中，并且异常抛出也一样。简化为一下步骤为：

1 使用StackContext(contextor) 创建一个上下文管理器，并将上下文管理函数推入_state.contexts 栈中
2 执行 async_task函数，注册callback回调：将_state.contexts栈中的上下文管理函数出栈，创建一个_StackContextWrapper 对象，该对象存储了出栈的async_task上下文函数。此时ioloop注册的callback为_StackContextWrapper对象。
3 ioloop调用callback，_StackContextWrapper中，将存储的上下文函数创建一个与syanc_task 一样的上下文管理器。在这个上下文环境中执行callback函数
4 3步骤中也涉及了创建上下文管理器的_state.contexts入栈出栈操作，多嵌套的with则会操作对应的上下文函数。执行完callback（或产生异常），执行上下文管理器的exit方法。

4个步骤的关键在于通过_state.contexts栈的处理，将主函数上下文管理函数绑定给了callback。因此无论callback还是async_task的上下文，通过contextor管理器都变得一样了。

contextor就像一个桥梁，连接着async_task和callback。而StackContext就像一个工程师，如何把函数和异步回调之间架设桥梁。

总结

本篇使用了大量的文字描述stack_contextor 的原理，其实还比不过打断点执行一遍。当然，对于多个嵌套的with，stack_context模块同样使用，其关键就在于_state.context是一个上下文管理器的栈，通过他的入栈和出栈可以轻松应对嵌套环境下的上下文环境。

下面是一段多嵌套的代码和输出结果，具体原理就不再分析了：

ioloop = tornado.ioloop.IOLoop.instance()

times = 0

def callback():
    print 'Run callback'
    raise ValueError('except in callback')

def async_task():
    global times
    times += 1
    print 'run async task {}'.format(times)
    ioloop.add_callback(callback=callback)

@contextlib.contextmanager
def A():
    print("Enter A context")
    try:
        yield
    except Exception as e:
        print("A catch the exception: %s" % e)
    finally:
        print("Exit A context")


@contextlib.contextmanager
def B():
    print("Enter B context")
    try:
        yield
    except Exception as e:
        print("B catch the exception: %s" % e)
    finally:
        print("Exit B context")


def main():
    with tornado.stack_context.StackContext(A):
        with tornado.stack_context.StackContext(B):
            async_task()

main()
ioloop.start()

输入结果很明了：

Enter A context
Enter B context
run async task 1
Exit B context
Exit A context
Enter A context
Enter B context
Run callback
B catch the exception: except in callback
Exit B context
Exit A context

先进入A的上下文，再进入B中，然后运行函数注册异步回调，退出B，再退出A。ioloop执行异步函数，再进入A，再进入B，运行回调，B发生异常，catch 捕获，退出B，再退出A 。

最后编辑于：2017.12.04 02:41:46

人面猴
序言：七十年代末，一起剥皮案震惊了整个滨河市，随后出现的几起案子，更是在滨河造成了极大的恐慌，老刑警刘岩，带你破解...
沈念sama阅读 204,590评论 6赞 478
死咒
序言：滨河连续发生了三起死亡事件，死亡现场离奇诡异，居然都是意外死亡，警方通过查阅死者的电脑和手机，发现死者居然都...
沈念sama阅读 86,808评论 2赞 381
救了他两次的神仙让他今天三更去死
文/潘晓璐我一进店门，熙熙楼的掌柜王于贵愁眉苦脸地迎上来，“玉大人，你说我怎么就摊上这事。” “怎么了？”我有些...
开封第一讲书人阅读 151,151评论 0赞 337
道士缉凶录：失踪的卖姜人
文/不坏的土叔我叫张陵，是天一观的道长。经常有香客问我，道长，这世上最难降的妖魔是什么？我笑而不...
开封第一讲书人阅读 54,779评论 1赞 277
港岛之恋（遗憾婚礼）
正文为了忘掉前任，我火速办了婚礼，结果婚礼上，老公的妹妹穿的比我还像新娘。我一直安慰自己，他们只是感情好，可当我...
茶点故事阅读 63,773评论 5赞 367
恶毒庶女顶嫁案：这布局不是一般人想出来的
文/花漫我一把揭开白布。她就那样静静地躺着，像睡着了一般。火红的嫁衣衬着肌肤如雪。梳的纹丝不乱的头发上，一...
开封第一讲书人阅读 48,656评论 1赞 281
城市分裂传说
那天，我揣着相机与录音，去河边找鬼。笑死，一个胖子当着我的面吹牛，可吹牛的内容都是我干的。我是一名探鬼主播，决...
沈念sama阅读 38,022评论 3赞 398
双鸳鸯连环套：你想象不到人心有多黑
文/苍兰香墨我猛地睁开眼，长吁一口气：“原来是场噩梦啊……” “哼！你这毒妇竟也来了？” 一声冷哼从身侧响起，我...
开封第一讲书人阅读 36,678评论 0赞 258
万荣杀人案实录
序言：老挝万荣一对情侣失踪，失踪者是张志新（化名）和其女友刘颖，没想到半个月后，有当地人在树林里发现了一具尸体，经...
沈念sama阅读 41,038评论 1赞 299
护林员之死
正文独居荒郊野岭守林人离奇死亡，尸身上长有42处带血的脓包…… 初始之章·张勋以下内容为张勋视角年9月15日...
茶点故事阅读 35,659评论 2赞 321
白月光启示录
正文我和宋清朗相恋三年，在试婚纱的时候发现自己被绿了。大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
茶点故事阅读 37,756评论 1赞 330
活死人
序言：一个原本活蹦乱跳的男人离奇死亡，死状恐怖，灵堂内的尸体忽然破棺而出，到底是诈尸还是另有隐情，我是刑警宁泽，带...
沈念sama阅读 33,411评论 4赞 321
日本核电站爆炸内幕
正文年R本政府宣布，位于F岛的核电站，受9级特大地震影响，放射性物质发生泄漏。R本人自食恶果不足惜，却给世界环境...
茶点故事阅读 39,005评论 3赞 307
男人毒药：我在死后第九天来索命
文/蒙蒙一、第九天我趴在偏房一处隐蔽的房顶上张望。院中可真热闹，春花似锦、人声如沸。这庄子的主人今日做“春日...
开封第一讲书人阅读 29,973评论 0赞 19
一桩弑父案，背后竟有这般阴谋
文/苍兰香墨我抬头看了看天上的太阳。三九已至，却和暖如春，着一层夹袄步出监牢的瞬间，已是汗流浃背。一阵脚步声响...
开封第一讲书人阅读 31,203评论 1赞 260
情欲美人皮
我被黑心中介骗来泰国打工，没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留，地道东北人。一个月前我还...
沈念sama阅读 45,053评论 2赞 350
代替公主和亲
正文我出身青楼，却偏偏与公主长得像，于是被迫代替她去往敌国和亲。传闻我的和亲对象是个残疾皇子，可洞房花烛夜当晚...
茶点故事阅读 42,495评论 2赞 343