Capture stack_context in IOLoop.add_future

diff --git a/tornado/ioloop.py b/tornado/ioloop.py
index 4e9267240d7978f9cdfd65020d26d7616d30c92d..745bdf8987d7f17267bab1dc8412584b5aa9c0fc 100644 (file)
--- a/tornado/ioloop.py
+++ b/tornado/ioloop.py
@@ -431,6 +431,7 @@ class IOLoop(object):
         """Schedules a callback on the IOLoop when the given future is finished.
         """
         assert isinstance(future, IOLoop._FUTURE_TYPES)
+        callback = stack_context.wrap(callback)
         future.add_done_callback(
             lambda future: self.add_callback(
                 functools.partial(callback, future)))

diff --git a/tornado/test/ioloop_test.py b/tornado/test/ioloop_test.py
index 0bac29eeb098279458b4010327fd4675c641b2b3..bfdc6ca58b46b5807f22fad9841268263d283a70 100644 (file)
--- a/tornado/test/ioloop_test.py
+++ b/tornado/test/ioloop_test.py
@@ -4,10 +4,12 @@
 from __future__ import absolute_import, division, with_statement
 import datetime
 import socket
+import threading
 import time
 
 from tornado.ioloop import IOLoop
 from tornado.netutil import bind_sockets
+from tornado.stack_context import ExceptionStackContext
 from tornado.testing import AsyncTestCase, LogTrapTestCase, get_unused_port
 from tornado.test.util import unittest
 
@@ -62,6 +64,35 @@ class TestIOLoopFutures(AsyncTestCase, LogTrapTestCase):
             future = self.wait()
             self.assertTrue(future.done())
             self.assertTrue(future.result() is None)
+
+    def test_add_future_stack_context(self):
+        ready = threading.Event()
+        def task():
+            # we must wait for the ioloop callback to be scheduled before
+            # the task completes to ensure that add_future adds the callback
+            # asynchronously (which is the scenario in which capturing
+            # the stack_context matters)
+            ready.wait(1)
+            assert ready.isSet(), "timed out"
+            raise Exception("worker")
+        def callback(future):
+            self.future = future
+            raise Exception("callback")
+        def handle_exception(typ, value, traceback):
+            self.exception = value
+            self.stop()
+            return True
+
+        # stack_context propagates to the ioloop callback, but the worker
+        # task just has its exceptions caught and saved in the Future.
+        with futures.ThreadPoolExecutor(1) as pool:
+            with ExceptionStackContext(handle_exception):
+                self.io_loop.add_future(pool.submit(task), callback)
+            ready.set()
+        self.wait()
+
+        self.assertEqual(self.exception.args[0], "callback")
+        self.assertEqual(self.future.exception().args[0], "worker")
 TestIOLoopFutures = unittest.skipIf(
     futures is None, "futures module not present")(TestIOLoopFutures)