Remove assertion for connection close instead throwing StreamClosedError

tchannel/tornado/connection.py

@@ -163,6 +163,9 @@ def _on_close(self):
         except queues.QueueEmpty:
             pass
 
+        if self._close_cb:
+            self._close_cb()
+
     def await(self):
         """Get the next call to this TChannel."""
         if self._loop_running:
@@ -208,10 +211,7 @@ def on_read_size(read_size_future):
             return read_body_future
 
         def on_error(future):
-            exception = future.exception()
-
-            if isinstance(exception, tornado.iostream.StreamClosedError):
-                self.close()
+            log.info("Failed to read data: %s", future.exception())
 
         size_width = frame.frame_rw.size_rw.width()
         read_bytes_future = self.connection.read_bytes(size_width)
@@ -290,7 +290,6 @@ def send(self, message):
         :returns:
             A Future containing the response for the message
         """
-        assert not self.closed
         assert self._loop_running, "Perform a handshake first."
         assert message.message_type in self.CALL_REQ_TYPES, (
             "Message '%s' can't use send" % repr(message)
@@ -314,8 +313,6 @@ def write(self, message):
         :param message:
             Message to write.
         """
-        assert not self.closed
-
         message.id = message.id or self.next_message_id()
 
         if message.message_type in self.CALL_REQ_TYPES:
@@ -350,10 +347,8 @@ def _write(self, message):
         return self.connection.write(body)
 
     def close(self):
-        if not self.connection.closed():
+        if not self.closed:
             self.connection.close()
-            if self._close_cb:
-                self._close_cb()
 
     @tornado.gen.coroutine
     def initiate_handshake(self, headers):
@@ -623,7 +618,6 @@ def send_request(self, request):
         :returns:
             A Future containing the response for the request
         """
-        assert not self.closed
         assert self._loop_running, "Perform a handshake first."
 
         assert request.id not in self._outstanding, (

tchannel/tornado/peer.py

@@ -27,17 +27,19 @@
 from collections import deque
 from itertools import takewhile, dropwhile
 from tornado import gen
+from tornado.iostream import StreamClosedError
 
 from ..schemes import DEFAULT as DEFAULT_SCHEME
 from ..retry import (
     DEFAULT as DEFAULT_RETRY, DEFAULT_RETRY_LIMIT
 )
-from tchannel.event import EventType
+
 from ..context import get_current_context
 from ..errors import BadRequestError
 from ..errors import NoAvailablePeerError
 from ..errors import TChannelError
 from ..errors import NetworkError
+from ..event import EventType
 from ..glossary import DEFAULT_TIMEOUT
 from ..glossary import MAX_SIZE_OF_ARG1
 from ..zipkin.annotation import Endpoint
@@ -409,6 +411,17 @@ def _send(self, connection, req):
         with timeout(response_future, req.ttl):
             try:
                 response = yield response_future
+            except StreamClosedError as error:
+                error = NetworkError(
+                    id=req.id,
+                    description=error.message,
+                    tracing=req.tracing,
+                )
+                # event: after_receive_error
+                self.tchannel.event_emitter.fire(
+                    EventType.after_receive_error, req, error,
+                )
+                raise
             except TChannelError as error:
                 # event: after_receive_error
                 self.tchannel.event_emitter.fire(
@@ -429,7 +442,6 @@ def send_with_retry(self, request, peer, retry_limit, connection):
             try:
                 response = yield self._send(connection, request)
                 raise gen.Return(response)
-            # Why are we retying on all errors????
             except TChannelError as error:
                 blacklist.add(peer.hostport)
                 (peer, connection) = yield self._prepare_for_retry(

tchannel/tornado/request.py

@@ -24,6 +24,7 @@
 
 import tornado
 import tornado.gen
+from tornado.iostream import StreamClosedError
 
 from tchannel import retry
 
@@ -177,6 +178,9 @@ def should_retry_on_error(self, error):
         if retry_flag == retry.NEVER:
             return False
 
+        if isinstance(error, StreamClosedError):
+            return True
+
         if error.code in [ErrorCode.bad_request, ErrorCode.cancelled,
                           ErrorCode.unhealthy]:
             return False

tests/integration/test_client_server.py

@@ -26,6 +26,7 @@
 from tchannel.errors import NetworkError
 from tchannel.errors import BadRequestError
 from tchannel import TChannel
+from tchannel.event import EventHook
 from tchannel.tornado.connection import StreamConnection
 from tests.util import big_arg
 
@@ -109,3 +110,24 @@ def test_endpoint_not_found(mock_server):
             endpoint='fooo',
             hostport=mock_server.hostport,
         )
+
+
+@pytest.mark.gen_test
+def test_connection_close(mock_server):
+    tchannel = TChannel(name='test')
+
+    class TestHook(EventHook):
+        def before_send_request(self, request):
+            # close the connection
+            peer = tchannel._dep_tchannel.peers.get(mock_server.hostport)
+            peer.close()
+            peer.connections[0].closed = True
+
+    tchannel.hooks.register(TestHook())
+
+    with pytest.raises(NetworkError):
+        yield tchannel.raw(
+            service='test-service',
+            hostport=mock_server.hostport,
+            endpoint='testg',
+        )

tests/tornado/test_connection.py

@@ -71,12 +71,13 @@ def test_close_callback_is_called():
     server = TChannel('server')
     server.listen()
 
-    close_cb = mock.Mock()
+    cb_future = tornado.gen.Future()
 
     conn = yield StreamConnection.outgoing(
         server.hostport, tchannel=mock.MagicMock()
     )
-    conn.set_close_callback(close_cb)
+    conn.set_close_callback(lambda: cb_future.set_result(True))
 
     conn.close()
-    close_cb.assert_called_once_with()
+
+    assert (yield cb_future)