Fix an exception when the outermost stack_context is deactivated.

Add more tests for deactivation in various scenarios.

diff --git a/tornado/stack_context.py b/tornado/stack_context.py
index 7e44beda636994e8e3335f873a043e0b585e805f..9ad3147658f5e5026b6c84a93c42be744be72db5 100644 (file)
--- a/tornado/stack_context.py
+++ b/tornado/stack_context.py
@@ -240,10 +240,11 @@ def _remove_deactivated(contexts):
     while ctx is not None:
         parent = ctx.old_contexts[1]
 
-        while parent is not None and not parent.active:
-            parent = parent.old_contexts[1]
-
+        while parent is not None:
+            if parent.active:
+                break
             ctx.old_contexts = parent.old_contexts
+            parent = parent.old_contexts[1]
 
         ctx = parent
 
@@ -267,6 +268,7 @@ def wrap(fn):
     cap_contexts = [_state.contexts]
 
     def wrapped(*args, **kwargs):
+        ret = None
         try:
             # Capture old state
             current_state = _state.contexts
@@ -298,7 +300,7 @@ def wrap(fn):
             # Execute callback if no exception happened while restoring state
             if top is None:
                 try:
-                    fn(*args, **kwargs)
+                    ret = fn(*args, **kwargs)
                 except:
                     exc = sys.exc_info()
                     top = contexts[1]
@@ -330,6 +332,7 @@ def wrap(fn):
                 raise_exc_info(exc)
         finally:
             _state.contexts = current_state
+        return ret
 
     wrapped._wrapped = True
     return wrapped

diff --git a/tornado/test/stack_context_test.py b/tornado/test/stack_context_test.py
index 976ef4000de4269945ef588e17d6bd5b9ba956a3..192bf8914ba7200b277e586f78eafde736722b19 100644 (file)
--- a/tornado/test/stack_context_test.py
+++ b/tornado/test/stack_context_test.py
@@ -128,6 +128,52 @@ class StackContextTest(AsyncTestCase):
         self.io_loop.add_callback(f1)
         self.wait()
 
+    def test_deactivate_order(self):
+        # Stack context deactivation has separate logic for deactivation at
+        # the head and tail of the stack, so make sure it works in any order.
+        def check_contexts():
+            # Make sure that the full-context array and the exception-context
+            # linked lists are consistent with each other.
+            full_contexts, chain = _state.contexts
+            exception_contexts = []
+            while chain is not None:
+                exception_contexts.append(chain)
+                chain = chain.old_contexts[1]
+            self.assertEqual(list(reversed(full_contexts)), exception_contexts)
+            return list(self.active_contexts)
+
+        def make_wrapped_function():
+            """Wraps a function in three stack contexts, and returns
+            the function along with the deactivation functions.
+            """
+            # Remove the test's stack context to make sure we can cover
+            # the case where the last context is deactivated.
+            with NullContext():
+                with StackContext(functools.partial(self.context, 'c0')) as c0,\
+                     StackContext(functools.partial(self.context, 'c1')) as c1,\
+                     StackContext(functools.partial(self.context, 'c2')) as c2:
+                    return (wrap(check_contexts), [c0, c1, c2])
+
+        # First make sure the test mechanism works without any deactivations
+        func, deactivate_callbacks = make_wrapped_function()
+        self.assertEqual(func(), ['c0', 'c1', 'c2'])
+
+        # Deactivate the tail
+        func, deactivate_callbacks = make_wrapped_function()
+        deactivate_callbacks[0]()
+        self.assertEqual(func(), ['c1', 'c2'])
+
+        # Deactivate the middle
+        func, deactivate_callbacks = make_wrapped_function()
+        deactivate_callbacks[1]()
+        self.assertEqual(func(), ['c0', 'c2'])
+
+        # Deactivate the head
+        func, deactivate_callbacks = make_wrapped_function()
+        deactivate_callbacks[2]()
+        self.assertEqual(func(), ['c0', 'c1'])
+
+
     def test_isolation_nonempty(self):
         # f2 and f3 are a chain of operations started in context c1.
         # f2 is incidentally run under context c2, but that context should