Browse files

Add an assertion for inconsistent StackContexts when used with genera…

…tors.
  • Loading branch information...
1 parent d98afd4 commit 1282b6efde5bc6ca57cd981dccd748927fa76555 @bdarnell bdarnell committed Feb 15, 2013
Showing with 73 additions and 8 deletions.
  1. +30 −6 tornado/stack_context.py
  2. +43 −2 tornado/test/stack_context_test.py
View
36 tornado/stack_context.py
@@ -77,6 +77,10 @@ def die_on_error():
from tornado.util import raise_exc_info
+class StackContextInconsistentError(Exception):
+ pass
+
+
class _State(threading.local):
def __init__(self):
self.contexts = ()
@@ -113,8 +117,10 @@ def __init__(self, context_factory, _active_cell=None):
def __enter__(self):
self.old_contexts = _state.contexts
# _state.contexts is a tuple of (class, arg, active_cell) tuples
- _state.contexts = (self.old_contexts +
- ((StackContext, self.context_factory, self.active_cell),))
+ self.new_contexts = (self.old_contexts +
+ ((StackContext, self.context_factory,
+ self.active_cell),))
+ _state.contexts = self.new_contexts
try:
self.context = self.context_factory()
self.context.__enter__()
@@ -127,7 +133,19 @@ def __exit__(self, type, value, traceback):
try:
return self.context.__exit__(type, value, traceback)
finally:
+ final_contexts = _state.contexts
_state.contexts = self.old_contexts
+ # Generator coroutines and with-statements with non-local
+ # effects interact badly. Check here for signs of
+ # the stack getting out of sync.
+ # Note that this check comes after restoring _state.context
+ # so that if it fails things are left in a (relatively)
+ # consistent state.
+ if final_contexts is not self.new_contexts:
+ raise StackContextInconsistentError(
+ 'stack_context inconsistency (may be caused by yield '
+ 'within a "with StackContext" block)')
+ self.old_contexts = self.new_contexts = None
class ExceptionStackContext(object):
@@ -149,18 +167,24 @@ def __init__(self, exception_handler, _active_cell=None):
def __enter__(self):
self.old_contexts = _state.contexts
- _state.contexts = (self.old_contexts +
- ((ExceptionStackContext, self.exception_handler,
- self.active_cell),))
+ self.new_contexts = (self.old_contexts +
+ ((ExceptionStackContext, self.exception_handler,
+ self.active_cell),))
+ _state.contexts = self.new_contexts
return lambda: operator.setitem(self.active_cell, 0, False)
def __exit__(self, type, value, traceback):
try:
if type is not None:
return self.exception_handler(type, value, traceback)
finally:
+ final_contexts = _state.contexts
_state.contexts = self.old_contexts
- self.old_contexts = None
+ if final_contexts is not self.new_contexts:
+ raise StackContextInconsistentError(
+ 'stack_context inconsistency (may be caused by yield '
+ 'within a "with StackContext" block)')
+ self.old_contexts = self.new_contexts = None
class NullContext(object):
View
45 tornado/test/stack_context_test.py
@@ -1,9 +1,10 @@
#!/usr/bin/env python
from __future__ import absolute_import, division, print_function, with_statement
+from tornado import gen
from tornado.log import app_log
-from tornado.stack_context import StackContext, wrap, NullContext
-from tornado.testing import AsyncHTTPTestCase, AsyncTestCase, ExpectLog
+from tornado.stack_context import StackContext, wrap, NullContext, StackContextInconsistentError, ExceptionStackContext
+from tornado.testing import AsyncHTTPTestCase, AsyncTestCase, ExpectLog, gen_test
from tornado.test.util import unittest
from tornado.web import asynchronous, Application, RequestHandler
import contextlib
@@ -168,6 +169,46 @@ def f3():
self.io_loop.add_callback(f1)
self.wait()
+ def test_yield_in_with(self):
+ @gen.engine
+ def f():
+ with StackContext(functools.partial(self.context, 'c1')):
+ # This yield is a problem: the generator will be suspended
+ # and the StackContext's __exit__ is not called yet, so
+ # the context will be left on _state.contexts for anything
+ # that runs before the yield resolves.
+ yield gen.Task(self.io_loop.add_callback)
+
+ with self.assertRaises(StackContextInconsistentError):
+ f()
+ self.wait()
+
+ @gen_test
+ def test_yield_outside_with(self):
+ # This pattern avoids the problem in the previous test.
+ cb = yield gen.Callback('k1')
+ with StackContext(functools.partial(self.context, 'c1')):
+ self.io_loop.add_callback(cb)
+ yield gen.Wait('k1')
+
+ def test_yield_in_with_exception_stack_context(self):
+ # As above, but with ExceptionStackContext instead of StackContext.
+ @gen.engine
+ def f():
+ with ExceptionStackContext(lambda t, v, tb: False):
+ yield gen.Task(self.io_loop.add_callback)
+
+ with self.assertRaises(StackContextInconsistentError):
+ f()
+ self.wait()
+
+ @gen_test
+ def test_yield_outside_with_exception_stack_context(self):
+ cb = yield gen.Callback('k1')
+ with ExceptionStackContext(lambda t, v, tb: False):
+ self.io_loop.add_callback(cb)
+ yield gen.Wait('k1')
+
if __name__ == '__main__':
unittest.main()

0 comments on commit 1282b6e

Please sign in to comment.