use add_future instead of add_done_callback to help with interleaving

uber · Sep 24, 2015 · 724413b · 724413b
1 parent bcecb95
commit 724413b
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 6 deletions.
diff --git a/tchannel/tornado/connection.py b/tchannel/tornado/connection.py
@@ -182,9 +182,12 @@ def on_read_size(read_size_future):
             size_bytes = read_size_future.result()
             size = frame.frame_rw.size_rw.read(BytesIO(size_bytes))
             read_body_future = self.connection.read_bytes(size - size_width)
-            read_body_future.add_done_callback(
+
+            tornado.ioloop.IOLoop.current().add_future(
+                read_body_future,
                 lambda future: on_body(future, size)
             )
+
             return read_body_future
 
         def on_error(future):
@@ -194,7 +197,12 @@ def on_error(future):
                 self.close()
 
         size_width = frame.frame_rw.size_rw.width()
-        self.connection.read_bytes(size_width).add_done_callback(on_read_size)
+        read_bytes_future = self.connection.read_bytes(size_width)
+
+        tornado.ioloop.IOLoop.current().add_future(
+            read_bytes_future,
+            on_read_size,
+        )
 
         return message_future
 
@@ -577,8 +585,9 @@ def stream_request(self, request):
 
         stream_future = self._stream(request, self.request_message_factory)
 
-        stream_future.add_done_callback(
-            lambda f: request.close_argstreams(force=True),
+        tornado.ioloop.IOLoop.current().add_future(
+            stream_future,
+            lambda f: request.close_argstreams(force=True)
         )
 
         return stream_future

diff --git a/tchannel/tornado/stream.py b/tchannel/tornado/stream.py
@@ -140,7 +140,8 @@ def read_chunk(future):
         # We're not ready yet
         if self.state != StreamState.completed and not len(self._stream):
             wait_future = self._condition.wait()
-            wait_future.add_done_callback(
+            tornado.ioloop.IOLoop.current().add_future(
+                wait_future,
                 lambda f: f.exception() or read_chunk(read_future)
             )
             return read_future

diff --git a/tchannel/tornado/util.py b/tchannel/tornado/util.py
@@ -85,7 +85,10 @@ def go():
         except StopIteration:
             all_done_future.set_result(None)
         else:
-            func(arg).add_done_callback(handle)
+            tornado.ioloop.IOLoop.current().add_future(
+                func(arg),
+                handle,
+            )
 
     go()