StackContext should not re-wrap previously-wrapped callbacks.

This allows for library code to "pop" back up to a previous version of the
stack.

diff --git a/tornado/stack_context.py b/tornado/stack_context.py
index 43e511d5b672308def64c6a839c83ac6318056d9..382da58602eaa76ebb09081d8be0a6367099d918 100644 (file)
--- a/tornado/stack_context.py
+++ b/tornado/stack_context.py
@@ -129,13 +129,14 @@ def wrap(fn, *args, **kwargs):
         callback(*args, **kwargs)
     else:
       callback(*args, **kwargs)
+  if getattr(fn, 'stack_context_wrapped', False):
+    return fn
   if args or kwargs:
     callback = functools.partial(fn, *args, **kwargs)
   else:
     callback = fn
   contexts = _state.contexts
-  if contexts:
-    return functools.partial(wrapped, callback, contexts, *args, **kwargs)
-  else:
-    return callback
+  result = functools.partial(wrapped, callback, contexts, *args, **kwargs)
+  result.stack_context_wrapped = True
+  return result
 

diff --git a/tornado/test/stack_context_test.py b/tornado/test/stack_context_test.py
index ee8be90b461a0e2c84737a2f80fb36ce51e20cf5..95c904303bace386b97799ab3dc06f0337db6463 100755 (executable)
--- a/tornado/test/stack_context_test.py
+++ b/tornado/test/stack_context_test.py
@@ -1,7 +1,10 @@
 #!/usr/bin/env python
 
-from tornado.testing import AsyncHTTPTestCase, LogTrapTestCase
+from tornado.stack_context import StackContext, wrap
+from tornado.testing import AsyncHTTPTestCase, AsyncTestCase, LogTrapTestCase
 from tornado.web import asynchronous, Application, RequestHandler
+import contextlib
+import functools
 import logging
 import unittest
 
@@ -33,7 +36,7 @@ class TestRequestHandler(RequestHandler):
     else:
       return 'unexpected failure'
 
-class StackContextTest(AsyncHTTPTestCase, LogTrapTestCase):
+class HTTPStackContextTest(AsyncHTTPTestCase, LogTrapTestCase):
   def get_app(self):
     return Application([('/', TestRequestHandler,
                          dict(io_loop=self.io_loop))])
@@ -48,5 +51,42 @@ class StackContextTest(AsyncHTTPTestCase, LogTrapTestCase):
     self.response = response
     self.stop()
 
+class StackContextTest(AsyncTestCase, LogTrapTestCase):
+  def setUp(self):
+    super(StackContextTest, self).setUp()
+    self.active_contexts = set()
+
+  @contextlib.contextmanager
+  def context(self, name):
+    assert name not in self.active_contexts
+    self.active_contexts.add(name)
+    yield
+    assert name in self.active_contexts
+    self.active_contexts.remove(name)
+
+  # Simulates the effect of an asynchronous library that uses its own
+  # StackContext internally and then returns control to the application.
+  def test_exit_library_context(self):
+    def library_function(callback):
+      # capture the caller's context before introducing our own
+      callback = wrap(callback)
+      with StackContext(functools.partial(self.context, 'library')):
+        self.io_loop.add_callback(
+          functools.partial(library_inner_callback, callback))
+    def library_inner_callback(callback):
+      assert 'application' in self.active_contexts
+      assert 'library' in self.active_contexts
+      # pass the callback out to the IOLoop to get out of the library
+      # context (could also use a NullContext here, but that would result
+      # in multiple instantiations of the application context)
+      self.io_loop.add_callback(callback)
+    def final_callback():
+      assert 'application' in self.active_contexts
+      assert 'library' not in self.active_contexts
+      self.stop()
+    with StackContext(functools.partial(self.context, 'application')):
+      library_function(final_callback)
+    self.wait()
+
 if __name__ == '__main__':
   unittest.main()