Detect stack context inconsistency in gen.Runner.

The existing checks in StackContext.__exit__ are run after the damage has
been done and may raise exceptions in the wrong place; gen.Runner
is a more reliable place to detect this problem.

Throwing the exception into the generator provides a better stack trace
for the problem and also simplifies some of the complexity of testing for
this feature.

diff --git a/tornado/gen.py b/tornado/gen.py
index 9f1846a5e4636be75390a8cad4f4c63d214dd1a8..480781e71d16d2ccfea6ed6740079b96d3e567cc 100644 (file)
--- a/tornado/gen.py
+++ b/tornado/gen.py
@@ -89,7 +89,7 @@ import types
 
 from tornado.concurrent import Future, TracebackFuture, is_future
 from tornado.ioloop import IOLoop
-from tornado.stack_context import ExceptionStackContext, wrap
+from tornado import stack_context
 
 
 class KeyReuseError(Exception):
@@ -194,7 +194,7 @@ def _make_coroutine_wrapper(func, replace_callback):
                 typ, value, tb = sys.exc_info()
             future.set_exc_info((typ, value, tb))
             return True
-        with ExceptionStackContext(handle_exception) as deactivate:
+        with stack_context.ExceptionStackContext(handle_exception) as deactivate:
             future.add_done_callback(lambda f: deactivate())
             try:
                 result = func(*args, **kwargs)
@@ -502,6 +502,7 @@ class Runner(object):
                     except Exception:
                         self.exc_info = sys.exc_info()
                 try:
+                    orig_stack_contexts = stack_context._state.contexts
                     if self.exc_info is not None:
                         self.had_exception = True
                         exc_info = self.exc_info
@@ -509,6 +510,11 @@ class Runner(object):
                         yielded = self.gen.throw(*exc_info)
                     else:
                         yielded = self.gen.send(next)
+                    if stack_context._state.contexts is not orig_stack_contexts:
+                        self.gen.throw(
+                            stack_context.StackContextInconsistentError(
+                                'stack_context inconsistency (probably caused '
+                                'by yield within a "with StackContext" block)'))
                 except (StopIteration, Return) as e:
                     self.finished = True
                     self.future = _null_future
@@ -562,7 +568,7 @@ class Runner(object):
             else:
                 result = None
             self.set_result(key, result)
-        return wrap(inner)
+        return stack_context.wrap(inner)
 
     def handle_exception(self, typ, value, tb):
         if not self.running and not self.finished:

diff --git a/tornado/test/stack_context_test.py b/tornado/test/stack_context_test.py
index 2919330521e0357e416bed607a769daa3b2fe5fb..d65a5b21465ae4defb826b7634800d7d6e8c9a27 100644 (file)
--- a/tornado/test/stack_context_test.py
+++ b/tornado/test/stack_context_test.py
@@ -219,22 +219,13 @@ class StackContextTest(AsyncTestCase):
     def test_yield_in_with(self):
         @gen.engine
         def f():
-            try:
-                self.callback = yield gen.Callback('a')
-                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.Wait('a')
-            except StackContextInconsistentError:
-                # In python <= 3.3, this suspended generator is never garbage
-                # collected, so it remains suspended in the 'yield' forever.
-                # Starting in 3.4, it is made collectable by raising
-                # a GeneratorExit exception from the yield, which gets
-                # converted into a StackContextInconsistentError by the
-                # exit of the 'with' block.
-                pass
+            self.callback = yield gen.Callback('a')
+            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.Wait('a')
 
         with self.assertRaises(StackContextInconsistentError):
             f()
@@ -257,11 +248,8 @@ class StackContextTest(AsyncTestCase):
         # As above, but with ExceptionStackContext instead of StackContext.
         @gen.engine
         def f():
-            try:
-                with ExceptionStackContext(lambda t, v, tb: False):
-                    yield gen.Task(self.io_loop.add_callback)
-            except StackContextInconsistentError:
-                pass
+            with ExceptionStackContext(lambda t, v, tb: False):
+                yield gen.Task(self.io_loop.add_callback)
 
         with self.assertRaises(StackContextInconsistentError):
             f()